libsnapshot: Implement OpenReader for CompressedSnapshotWriter.

Bug: 168554689
Test: vts_libsnapshot_test
Test: full OTA with update_device.py
Test: incremental OTA with update_device.py
Change-Id: I3878abfd767d2e47cf8486bc2c06233da2f1ef08
diff --git a/fs_mgr/libsnapshot/Android.bp b/fs_mgr/libsnapshot/Android.bp
index 95606d7..f154138 100644
--- a/fs_mgr/libsnapshot/Android.bp
+++ b/fs_mgr/libsnapshot/Android.bp
@@ -239,6 +239,7 @@
     srcs: [
         "partition_cow_creator_test.cpp",
         "snapshot_metadata_updater_test.cpp",
+        "snapshot_reader_test.cpp",
         "snapshot_test.cpp",
     ],
     shared_libs: [
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/cow_reader.h b/fs_mgr/libsnapshot/include/libsnapshot/cow_reader.h
index 814e104..b863ff2 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/cow_reader.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/cow_reader.h
@@ -39,11 +39,17 @@
     // maximum number of bytes that can be written to the returned buffer.
     //
     // The returned buffer is owned by IByteSink, but must remain valid until
-    // the ready operation has completed (or the entire buffer has been
+    // the read operation has completed (or the entire buffer has been
     // covered by calls to ReturnData).
     //
     // After calling GetBuffer(), all previous buffers returned are no longer
     // valid.
+    //
+    // GetBuffer() is intended to be sequential. A returned size of N indicates
+    // that the output stream will advance by N bytes, and the ReturnData call
+    // indicates that those bytes have been fulfilled. Therefore, it is
+    // possible to have ReturnBuffer do nothing, if the implementation doesn't
+    // care about incremental writes.
     virtual void* GetBuffer(size_t requested, size_t* actual) = 0;
 
     // Called when a section returned by |GetBuffer| has been filled with data.
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot_writer.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot_writer.h
index f76f545..45c7ec4 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot_writer.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot_writer.h
@@ -42,9 +42,10 @@
   protected:
     android::base::borrowed_fd GetSourceFd();
 
+    std::optional<std::string> source_device_;
+
   private:
     android::base::unique_fd source_fd_;
-    std::optional<std::string> source_device_;
 };
 
 // Send writes to a COW or a raw device directly, based on a threshold.
@@ -52,7 +53,7 @@
   public:
     CompressedSnapshotWriter(const CowOptions& options);
 
-    // Sets the COW device, if needed.
+    // Sets the COW device; this is required.
     bool SetCowDevice(android::base::unique_fd&& cow_device);
 
     bool Finalize() override;
diff --git a/fs_mgr/libsnapshot/snapshot_reader.cpp b/fs_mgr/libsnapshot/snapshot_reader.cpp
index 0d47468..a4a652a 100644
--- a/fs_mgr/libsnapshot/snapshot_reader.cpp
+++ b/fs_mgr/libsnapshot/snapshot_reader.cpp
@@ -14,13 +14,17 @@
 // limitations under the License.
 //
 
-#include <ext4_utils/ext4_utils.h>
-
 #include "snapshot_reader.h"
 
+#include <android-base/file.h>
+#include <android-base/logging.h>
+#include <ext4_utils/ext4_utils.h>
+
 namespace android {
 namespace snapshot {
 
+using android::base::borrowed_fd;
+
 // Not supported.
 bool ReadOnlyFileDescriptor::Open(const char*, int, mode_t) {
     errno = EINVAL;
@@ -73,5 +77,252 @@
     return true;
 }
 
+bool CompressedSnapshotReader::SetCow(std::unique_ptr<CowReader>&& cow) {
+    cow_ = std::move(cow);
+
+    CowHeader header;
+    if (!cow_->GetHeader(&header)) {
+        return false;
+    }
+    block_size_ = header.block_size;
+
+    // Populate the operation map.
+    op_iter_ = cow_->GetOpIter();
+    while (!op_iter_->Done()) {
+        const CowOperation* op = &op_iter_->Get();
+        if (op->new_block >= ops_.size()) {
+            ops_.resize(op->new_block + 1, nullptr);
+        }
+        ops_[op->new_block] = op;
+        op_iter_->Next();
+    }
+
+    return true;
+}
+
+void CompressedSnapshotReader::SetSourceDevice(const std::string& source_device) {
+    source_device_ = {source_device};
+}
+
+void CompressedSnapshotReader::SetBlockDeviceSize(uint64_t block_device_size) {
+    block_device_size_ = block_device_size;
+}
+
+borrowed_fd CompressedSnapshotReader::GetSourceFd() {
+    if (source_fd_ < 0) {
+        if (!source_device_) {
+            LOG(ERROR) << "CompressedSnapshotReader needs source device, but none was set";
+            errno = EINVAL;
+            return {-1};
+        }
+        source_fd_.reset(open(source_device_->c_str(), O_RDONLY | O_CLOEXEC));
+        if (source_fd_ < 0) {
+            PLOG(ERROR) << "open " << *source_device_;
+            return {-1};
+        }
+    }
+    return source_fd_;
+}
+
+class MemoryByteSink : public IByteSink {
+  public:
+    MemoryByteSink(void* buf, size_t count) {
+        buf_ = reinterpret_cast<uint8_t*>(buf);
+        pos_ = buf_;
+        end_ = buf_ + count;
+    }
+
+    void* GetBuffer(size_t requested, size_t* actual) override {
+        *actual = std::min(remaining(), requested);
+        if (!*actual) {
+            return nullptr;
+        }
+
+        uint8_t* start = pos_;
+        pos_ += *actual;
+        return start;
+    }
+
+    bool ReturnData(void*, size_t) override { return true; }
+
+    uint8_t* buf() const { return buf_; }
+    uint8_t* pos() const { return pos_; }
+    size_t remaining() const { return end_ - pos_; }
+
+  private:
+    uint8_t* buf_;
+    uint8_t* pos_;
+    uint8_t* end_;
+};
+
+ssize_t CompressedSnapshotReader::Read(void* buf, size_t count) {
+    // Find the start and end chunks, inclusive.
+    uint64_t start_chunk = offset_ / block_size_;
+    uint64_t end_chunk = (offset_ + count - 1) / block_size_;
+
+    // Chop off the first N bytes if the position is not block-aligned.
+    size_t start_offset = offset_ % block_size_;
+
+    MemoryByteSink sink(buf, count);
+
+    size_t initial_bytes = std::min(block_size_ - start_offset, sink.remaining());
+    ssize_t rv = ReadBlock(start_chunk, &sink, start_offset, initial_bytes);
+    if (rv < 0) {
+        return -1;
+    }
+    offset_ += rv;
+
+    for (uint64_t chunk = start_chunk + 1; chunk < end_chunk; chunk++) {
+        ssize_t rv = ReadBlock(chunk, &sink, 0);
+        if (rv < 0) {
+            return -1;
+        }
+        offset_ += rv;
+    }
+
+    if (sink.remaining()) {
+        ssize_t rv = ReadBlock(end_chunk, &sink, 0, {sink.remaining()});
+        if (rv < 0) {
+            return -1;
+        }
+        offset_ += rv;
+    }
+
+    errno = 0;
+
+    DCHECK(sink.pos() - sink.buf() == count);
+    return count;
+}
+
+// Discard the first N bytes of a sink request, or any excess bytes.
+class PartialSink : public MemoryByteSink {
+  public:
+    PartialSink(void* buffer, size_t size, size_t ignore_start)
+        : MemoryByteSink(buffer, size), ignore_start_(ignore_start) {}
+
+    void* GetBuffer(size_t requested, size_t* actual) override {
+        // Throw away the first N bytes if needed.
+        if (ignore_start_) {
+            *actual = std::min({requested, ignore_start_, sizeof(discard_)});
+            ignore_start_ -= *actual;
+            return discard_;
+        }
+        // Throw away any excess bytes if needed.
+        if (remaining() == 0) {
+            *actual = std::min(requested, sizeof(discard_));
+            return discard_;
+        }
+        return MemoryByteSink::GetBuffer(requested, actual);
+    }
+
+  private:
+    size_t ignore_start_;
+    char discard_[4096];
+};
+
+ssize_t CompressedSnapshotReader::ReadBlock(uint64_t chunk, IByteSink* sink, size_t start_offset,
+                                            const std::optional<uint64_t>& max_bytes) {
+    size_t bytes_to_read = block_size_;
+    if (max_bytes) {
+        bytes_to_read = *max_bytes;
+    }
+
+    // The offset is relative to the chunk; we should be reading no more than
+    // one chunk.
+    CHECK(start_offset + bytes_to_read <= block_size_);
+
+    const CowOperation* op = nullptr;
+    if (chunk < ops_.size()) {
+        op = ops_[chunk];
+    }
+
+    size_t actual;
+    void* buffer = sink->GetBuffer(bytes_to_read, &actual);
+    if (!buffer || actual < bytes_to_read) {
+        // This should never happen unless we calculated the read size wrong
+        // somewhere. MemoryByteSink always fulfills the entire requested
+        // region unless there's not enough buffer remaining.
+        LOG(ERROR) << "Asked for buffer of size " << bytes_to_read << ", got " << actual;
+        errno = EINVAL;
+        return -1;
+    }
+
+    if (!op || op->type == kCowCopyOp) {
+        borrowed_fd fd = GetSourceFd();
+        if (fd < 0) {
+            // GetSourceFd sets errno.
+            return -1;
+        }
+
+        if (op) {
+            chunk = op->source;
+        }
+
+        off64_t offset = (chunk * block_size_) + start_offset;
+        if (!android::base::ReadFullyAtOffset(fd, buffer, bytes_to_read, offset)) {
+            PLOG(ERROR) << "read " << *source_device_;
+            // ReadFullyAtOffset sets errno.
+            return -1;
+        }
+    } else if (op->type == kCowZeroOp) {
+        memset(buffer, 0, bytes_to_read);
+    } else if (op->type == kCowReplaceOp) {
+        PartialSink partial_sink(buffer, bytes_to_read, start_offset);
+        if (!cow_->ReadData(*op, &partial_sink)) {
+            LOG(ERROR) << "CompressedSnapshotReader failed to read replace op";
+            errno = EIO;
+            return -1;
+        }
+    } else {
+        LOG(ERROR) << "CompressedSnapshotReader unknown op type: " << op->type;
+        errno = EINVAL;
+        return -1;
+    }
+
+    // MemoryByteSink doesn't do anything in ReturnBuffer, so don't bother calling it.
+    return bytes_to_read;
+}
+
+off64_t CompressedSnapshotReader::Seek(off64_t offset, int whence) {
+    switch (whence) {
+        case SEEK_SET:
+            offset_ = offset;
+            break;
+        case SEEK_END:
+            offset_ = static_cast<off64_t>(block_device_size_) + offset;
+            break;
+        case SEEK_CUR:
+            offset_ += offset;
+            break;
+        default:
+            LOG(ERROR) << "Unrecognized seek whence: " << whence;
+            errno = EINVAL;
+            return -1;
+    }
+    return offset_;
+}
+
+uint64_t CompressedSnapshotReader::BlockDevSize() {
+    return block_device_size_;
+}
+
+bool CompressedSnapshotReader::Close() {
+    cow_ = nullptr;
+    source_fd_ = {};
+    return true;
+}
+
+bool CompressedSnapshotReader::IsSettingErrno() {
+    return true;
+}
+
+bool CompressedSnapshotReader::IsOpen() {
+    return cow_ != nullptr;
+}
+
+bool CompressedSnapshotReader::Flush() {
+    return true;
+}
+
 }  // namespace snapshot
 }  // namespace android
diff --git a/fs_mgr/libsnapshot/snapshot_reader.h b/fs_mgr/libsnapshot/snapshot_reader.h
index 1f2ffe2..a3e7e22 100644
--- a/fs_mgr/libsnapshot/snapshot_reader.h
+++ b/fs_mgr/libsnapshot/snapshot_reader.h
@@ -16,7 +16,11 @@
 
 #pragma once
 
+#include <optional>
+#include <vector>
+
 #include <android-base/file.h>
+#include <libsnapshot/cow_reader.h>
 #include <payload_consumer/file_descriptor.h>
 
 namespace android {
@@ -46,5 +50,36 @@
     android::base::unique_fd fd_;
 };
 
+class CompressedSnapshotReader : public ReadOnlyFileDescriptor {
+  public:
+    bool SetCow(std::unique_ptr<CowReader>&& cow);
+    void SetSourceDevice(const std::string& source_device);
+    void SetBlockDeviceSize(uint64_t block_device_size);
+
+    ssize_t Read(void* buf, size_t count) override;
+    off64_t Seek(off64_t offset, int whence) override;
+    uint64_t BlockDevSize() override;
+    bool Close() override;
+    bool IsSettingErrno() override;
+    bool IsOpen() override;
+    bool Flush() override;
+
+  private:
+    ssize_t ReadBlock(uint64_t chunk, IByteSink* sink, size_t start_offset,
+                      const std::optional<uint64_t>& max_bytes = {});
+    android::base::borrowed_fd GetSourceFd();
+
+    std::unique_ptr<CowReader> cow_;
+    std::unique_ptr<ICowOpIter> op_iter_;
+    uint32_t block_size_ = 0;
+
+    std::optional<std::string> source_device_;
+    android::base::unique_fd source_fd_;
+    uint64_t block_device_size_ = 0;
+    off64_t offset_ = 0;
+
+    std::vector<const CowOperation*> ops_;
+};
+
 }  // namespace snapshot
 }  // namespace android
diff --git a/fs_mgr/libsnapshot/snapshot_reader_test.cpp b/fs_mgr/libsnapshot/snapshot_reader_test.cpp
new file mode 100644
index 0000000..f54eefa
--- /dev/null
+++ b/fs_mgr/libsnapshot/snapshot_reader_test.cpp
@@ -0,0 +1,165 @@
+// Copyright (C) 2018 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <libsnapshot/snapshot.h>
+
+#include <unordered_set>
+
+#include <android-base/file.h>
+#include <gtest/gtest.h>
+#include <libsnapshot/cow_writer.h>
+#include <payload_consumer/file_descriptor.h>
+
+namespace android {
+namespace snapshot {
+
+using android::base::unique_fd;
+using chromeos_update_engine::FileDescriptor;
+
+static constexpr uint32_t kBlockSize = 4096;
+static constexpr size_t kBlockCount = 10;
+
+class OfflineSnapshotTest : public ::testing::Test {
+  protected:
+    virtual void SetUp() override {
+        base_ = std::make_unique<TemporaryFile>();
+        ASSERT_GE(base_->fd, 0) << strerror(errno);
+
+        cow_ = std::make_unique<TemporaryFile>();
+        ASSERT_GE(cow_->fd, 0) << strerror(errno);
+
+        WriteBaseDevice();
+    }
+
+    virtual void TearDown() override {
+        base_ = nullptr;
+        cow_ = nullptr;
+        base_blocks_ = {};
+    }
+
+    void WriteBaseDevice() {
+        unique_fd random(open("/dev/urandom", O_RDONLY));
+        ASSERT_GE(random, 0);
+
+        for (size_t i = 0; i < kBlockCount; i++) {
+            std::string block(kBlockSize, 0);
+            ASSERT_TRUE(android::base::ReadFully(random, block.data(), block.size()));
+            ASSERT_TRUE(android::base::WriteFully(base_->fd, block.data(), block.size()));
+            base_blocks_.emplace_back(std::move(block));
+        }
+        ASSERT_EQ(fsync(base_->fd), 0);
+    }
+
+    void WriteCow(ISnapshotWriter* writer) {
+        std::string new_block = MakeNewBlockString();
+
+        ASSERT_TRUE(writer->AddCopy(3, 0));
+        ASSERT_TRUE(writer->AddRawBlocks(5, new_block.data(), new_block.size()));
+        ASSERT_TRUE(writer->AddZeroBlocks(7, 2));
+        ASSERT_TRUE(writer->Finalize());
+    }
+
+    void TestBlockReads(ISnapshotWriter* writer) {
+        auto reader = writer->OpenReader();
+        ASSERT_NE(reader, nullptr);
+
+        // Test that unchanged blocks are not modified.
+        std::unordered_set<size_t> changed_blocks = {3, 5, 7, 8};
+        for (size_t i = 0; i < kBlockCount; i++) {
+            if (changed_blocks.count(i)) {
+                continue;
+            }
+
+            std::string block(kBlockSize, 0);
+            ASSERT_EQ(reader->Seek(i * kBlockSize, SEEK_SET), i * kBlockSize);
+            ASSERT_EQ(reader->Read(block.data(), block.size()), kBlockSize);
+            ASSERT_EQ(block, base_blocks_[i]);
+        }
+
+        // Test that we can read back our modified blocks.
+        std::string block(kBlockSize, 0);
+        ASSERT_EQ(reader->Seek(3 * kBlockSize, SEEK_SET), 3 * kBlockSize);
+        ASSERT_EQ(reader->Read(block.data(), block.size()), kBlockSize);
+        ASSERT_EQ(block, base_blocks_[0]);
+
+        ASSERT_EQ(reader->Seek(5 * kBlockSize, SEEK_SET), 5 * kBlockSize);
+        ASSERT_EQ(reader->Read(block.data(), block.size()), kBlockSize);
+        ASSERT_EQ(block, MakeNewBlockString());
+
+        std::string two_blocks(kBlockSize * 2, 0x7f);
+        std::string zeroes(kBlockSize * 2, 0);
+        ASSERT_EQ(reader->Seek(7 * kBlockSize, SEEK_SET), 7 * kBlockSize);
+        ASSERT_EQ(reader->Read(two_blocks.data(), two_blocks.size()), two_blocks.size());
+        ASSERT_EQ(two_blocks, zeroes);
+    }
+
+    void TestByteReads(ISnapshotWriter* writer) {
+        auto reader = writer->OpenReader();
+        ASSERT_NE(reader, nullptr);
+
+        std::string blob(kBlockSize * 3, 'x');
+
+        // Test that we can read in the middle of a block.
+        static constexpr size_t kOffset = 970;
+        off64_t offset = 3 * kBlockSize + kOffset;
+        ASSERT_EQ(reader->Seek(0, SEEK_SET), 0);
+        ASSERT_EQ(reader->Seek(offset, SEEK_CUR), offset);
+        ASSERT_EQ(reader->Read(blob.data(), blob.size()), blob.size());
+        ASSERT_EQ(blob.substr(0, 100), base_blocks_[0].substr(kOffset, 100));
+        ASSERT_EQ(blob.substr(kBlockSize - kOffset, kBlockSize), base_blocks_[4]);
+        ASSERT_EQ(blob.substr(kBlockSize * 2 - kOffset, 100), MakeNewBlockString().substr(0, 100));
+        ASSERT_EQ(blob.substr(blob.size() - kOffset), base_blocks_[6].substr(0, kOffset));
+
+        // Pull a random byte from the compressed block.
+        char value;
+        offset = 5 * kBlockSize + 1000;
+        ASSERT_EQ(reader->Seek(offset, SEEK_SET), offset);
+        ASSERT_EQ(reader->Read(&value, sizeof(value)), sizeof(value));
+        ASSERT_EQ(value, MakeNewBlockString()[1000]);
+    }
+
+    void TestReads(ISnapshotWriter* writer) {
+        ASSERT_NO_FATAL_FAILURE(TestBlockReads(writer));
+        ASSERT_NO_FATAL_FAILURE(TestByteReads(writer));
+    }
+
+    std::string MakeNewBlockString() {
+        std::string new_block = "This is a new block";
+        new_block.resize(kBlockSize / 2, '*');
+        new_block.resize(kBlockSize, '!');
+        return new_block;
+    }
+
+    std::unique_ptr<TemporaryFile> base_;
+    std::unique_ptr<TemporaryFile> cow_;
+    std::vector<std::string> base_blocks_;
+};
+
+TEST_F(OfflineSnapshotTest, CompressedSnapshot) {
+    CowOptions options;
+    options.compression = "gz";
+    options.max_blocks = {kBlockCount};
+
+    unique_fd cow_fd(dup(cow_->fd));
+    ASSERT_GE(cow_fd, 0);
+
+    auto writer = std::make_unique<CompressedSnapshotWriter>(options);
+    writer->SetSourceDevice(base_->path);
+    ASSERT_TRUE(writer->SetCowDevice(std::move(cow_fd)));
+    ASSERT_NO_FATAL_FAILURE(WriteCow(writer.get()));
+    ASSERT_NO_FATAL_FAILURE(TestReads(writer.get()));
+}
+
+}  // namespace snapshot
+}  // namespace android
diff --git a/fs_mgr/libsnapshot/snapshot_writer.cpp b/fs_mgr/libsnapshot/snapshot_writer.cpp
index 9b1ab97..d6ecbfd 100644
--- a/fs_mgr/libsnapshot/snapshot_writer.cpp
+++ b/fs_mgr/libsnapshot/snapshot_writer.cpp
@@ -68,7 +68,31 @@
 }
 
 std::unique_ptr<FileDescriptor> CompressedSnapshotWriter::OpenReader() {
-    return nullptr;
+    unique_fd cow_fd(dup(cow_device_.get()));
+    if (cow_fd < 0) {
+        PLOG(ERROR) << "dup COW device";
+        return nullptr;
+    }
+
+    auto cow = std::make_unique<CowReader>();
+    if (!cow->Parse(std::move(cow_fd))) {
+        LOG(ERROR) << "Unable to read COW";
+        return nullptr;
+    }
+
+    auto reader = std::make_unique<CompressedSnapshotReader>();
+    if (!reader->SetCow(std::move(cow))) {
+        LOG(ERROR) << "Unable to initialize COW reader";
+        return nullptr;
+    }
+    if (source_device_) {
+        reader->SetSourceDevice(*source_device_);
+    }
+
+    const auto& cow_options = options();
+    reader->SetBlockDeviceSize(*cow_options.max_blocks * cow_options.block_size);
+
+    return reader;
 }
 
 bool CompressedSnapshotWriter::EmitCopy(uint64_t new_block, uint64_t old_block) {