Merge "Add write support to SplitFiemap."
diff --git a/fs_mgr/libfiemap_writer/fiemap_writer.cpp b/fs_mgr/libfiemap_writer/fiemap_writer.cpp
index 3d41876..99a1a2f 100644
--- a/fs_mgr/libfiemap_writer/fiemap_writer.cpp
+++ b/fs_mgr/libfiemap_writer/fiemap_writer.cpp
@@ -618,7 +618,6 @@
 
     fmap->file_path_ = abs_path;
     fmap->bdev_path_ = bdev_path;
-    fmap->file_fd_ = std::move(file_fd);
     fmap->file_size_ = file_size;
     fmap->bdev_size_ = bdevsz;
     fmap->fs_type_ = fs_type;
diff --git a/fs_mgr/libfiemap_writer/fiemap_writer_test.cpp b/fs_mgr/libfiemap_writer/fiemap_writer_test.cpp
index ab4efae..66eb9ae 100644
--- a/fs_mgr/libfiemap_writer/fiemap_writer_test.cpp
+++ b/fs_mgr/libfiemap_writer/fiemap_writer_test.cpp
@@ -38,6 +38,9 @@
 
 #include "utility.h"
 
+namespace android {
+namespace fiemap_writer {
+
 using namespace std;
 using namespace std::string_literals;
 using namespace android::fiemap_writer;
@@ -234,6 +237,105 @@
     ASSERT_EQ(errno, ENOENT);
 }
 
+static string ReadSplitFiles(const std::string& base_path, size_t num_files) {
+    std::string result;
+    for (int i = 0; i < num_files; i++) {
+        std::string path = base_path + android::base::StringPrintf(".%04d", i);
+        std::string data;
+        if (!android::base::ReadFileToString(path, &data)) {
+            return {};
+        }
+        result += data;
+    }
+    return result;
+}
+
+TEST_F(SplitFiemapTest, WriteWholeFile) {
+    static constexpr size_t kChunkSize = 32768;
+    static constexpr size_t kSize = kChunkSize * 3;
+    auto ptr = SplitFiemap::Create(testfile, kSize, kChunkSize);
+    ASSERT_NE(ptr, nullptr);
+
+    auto buffer = std::make_unique<int[]>(kSize / sizeof(int));
+    for (size_t i = 0; i < kSize / sizeof(int); i++) {
+        buffer[i] = i;
+    }
+    ASSERT_TRUE(ptr->Write(buffer.get(), kSize));
+
+    std::string expected(reinterpret_cast<char*>(buffer.get()), kSize);
+    auto actual = ReadSplitFiles(testfile, 3);
+    ASSERT_EQ(expected.size(), actual.size());
+    EXPECT_EQ(memcmp(expected.data(), actual.data(), actual.size()), 0);
+}
+
+TEST_F(SplitFiemapTest, WriteFileInChunks1) {
+    static constexpr size_t kChunkSize = 32768;
+    static constexpr size_t kSize = kChunkSize * 3;
+    auto ptr = SplitFiemap::Create(testfile, kSize, kChunkSize);
+    ASSERT_NE(ptr, nullptr);
+
+    auto buffer = std::make_unique<int[]>(kSize / sizeof(int));
+    for (size_t i = 0; i < kSize / sizeof(int); i++) {
+        buffer[i] = i;
+    }
+
+    // Write in chunks of 1000 (so some writes straddle the boundary of two
+    // files).
+    size_t bytes_written = 0;
+    while (bytes_written < kSize) {
+        size_t to_write = std::min(kSize - bytes_written, (size_t)1000);
+        char* data = reinterpret_cast<char*>(buffer.get()) + bytes_written;
+        ASSERT_TRUE(ptr->Write(data, to_write));
+        bytes_written += to_write;
+    }
+
+    std::string expected(reinterpret_cast<char*>(buffer.get()), kSize);
+    auto actual = ReadSplitFiles(testfile, 3);
+    ASSERT_EQ(expected.size(), actual.size());
+    EXPECT_EQ(memcmp(expected.data(), actual.data(), actual.size()), 0);
+}
+
+TEST_F(SplitFiemapTest, WriteFileInChunks2) {
+    static constexpr size_t kChunkSize = 32768;
+    static constexpr size_t kSize = kChunkSize * 3;
+    auto ptr = SplitFiemap::Create(testfile, kSize, kChunkSize);
+    ASSERT_NE(ptr, nullptr);
+
+    auto buffer = std::make_unique<int[]>(kSize / sizeof(int));
+    for (size_t i = 0; i < kSize / sizeof(int); i++) {
+        buffer[i] = i;
+    }
+
+    // Write in chunks of 32KiB so every write is exactly at the end of the
+    // current file.
+    size_t bytes_written = 0;
+    while (bytes_written < kSize) {
+        size_t to_write = std::min(kSize - bytes_written, kChunkSize);
+        char* data = reinterpret_cast<char*>(buffer.get()) + bytes_written;
+        ASSERT_TRUE(ptr->Write(data, to_write));
+        bytes_written += to_write;
+    }
+
+    std::string expected(reinterpret_cast<char*>(buffer.get()), kSize);
+    auto actual = ReadSplitFiles(testfile, 3);
+    ASSERT_EQ(expected.size(), actual.size());
+    EXPECT_EQ(memcmp(expected.data(), actual.data(), actual.size()), 0);
+}
+
+TEST_F(SplitFiemapTest, WritePastEnd) {
+    static constexpr size_t kChunkSize = 32768;
+    static constexpr size_t kSize = kChunkSize * 3;
+    auto ptr = SplitFiemap::Create(testfile, kSize, kChunkSize);
+    ASSERT_NE(ptr, nullptr);
+
+    auto buffer = std::make_unique<int[]>(kSize / sizeof(int));
+    for (size_t i = 0; i < kSize / sizeof(int); i++) {
+        buffer[i] = i;
+    }
+    ASSERT_TRUE(ptr->Write(buffer.get(), kSize));
+    ASSERT_FALSE(ptr->Write(buffer.get(), kSize));
+}
+
 class VerifyBlockWritesExt4 : public ::testing::Test {
     // 2GB Filesystem and 4k block size by default
     static constexpr uint64_t block_size = 4096;
@@ -333,6 +435,11 @@
     return true;
 }
 
+}  // namespace fiemap_writer
+}  // namespace android
+
+using namespace android::fiemap_writer;
+
 int main(int argc, char** argv) {
     ::testing::InitGoogleTest(&argc, argv);
     if (argc <= 1) {
diff --git a/fs_mgr/libfiemap_writer/include/libfiemap_writer/fiemap_writer.h b/fs_mgr/libfiemap_writer/include/libfiemap_writer/fiemap_writer.h
index 831bc75..9486122 100644
--- a/fs_mgr/libfiemap_writer/include/libfiemap_writer/fiemap_writer.h
+++ b/fs_mgr/libfiemap_writer/include/libfiemap_writer/fiemap_writer.h
@@ -88,9 +88,6 @@
     // Block device on which we have created the file.
     std::string bdev_path_;
 
-    // File descriptors for the file and block device
-    ::android::base::unique_fd file_fd_;
-
     // Size in bytes of the file this class is writing
     uint64_t file_size_;
 
diff --git a/fs_mgr/libfiemap_writer/include/libfiemap_writer/split_fiemap_writer.h b/fs_mgr/libfiemap_writer/include/libfiemap_writer/split_fiemap_writer.h
index 765cc84..07f3c10 100644
--- a/fs_mgr/libfiemap_writer/include/libfiemap_writer/split_fiemap_writer.h
+++ b/fs_mgr/libfiemap_writer/include/libfiemap_writer/split_fiemap_writer.h
@@ -23,6 +23,8 @@
 #include <string>
 #include <vector>
 
+#include <android-base/unique_fd.h>
+
 #include "fiemap_writer.h"
 
 namespace android {
@@ -54,6 +56,16 @@
     // this returns true and does not report an error.
     static bool RemoveSplitFiles(const std::string& file_path, std::string* message = nullptr);
 
+    // Return whether all components of a split file still have pinned extents.
+    bool HasPinnedExtents() const;
+
+    // Helper method for writing data that spans files. Note there is no seek
+    // method (yet); this starts at 0 and increments the position by |bytes|.
+    bool Write(const void* data, uint64_t bytes);
+
+    // Flush all writes to all split files.
+    bool Flush();
+
     const std::vector<struct fiemap_extent>& extents();
     uint32_t block_size() const;
     uint64_t size() const { return total_size_; }
@@ -73,6 +85,11 @@
     std::vector<FiemapUniquePtr> files_;
     std::vector<struct fiemap_extent> extents_;
     uint64_t total_size_ = 0;
+
+    // Most recently open file and position for Write().
+    size_t cursor_index_ = 0;
+    uint64_t cursor_file_pos_ = 0;
+    android::base::unique_fd cursor_fd_;
 };
 
 }  // namespace fiemap_writer
diff --git a/fs_mgr/libfiemap_writer/split_fiemap_writer.cpp b/fs_mgr/libfiemap_writer/split_fiemap_writer.cpp
index 1f80370..dbb67a8 100644
--- a/fs_mgr/libfiemap_writer/split_fiemap_writer.cpp
+++ b/fs_mgr/libfiemap_writer/split_fiemap_writer.cpp
@@ -176,6 +176,15 @@
     return ok;
 }
 
+bool SplitFiemap::HasPinnedExtents() const {
+    for (const auto& file : files_) {
+        if (!FiemapWriter::HasPinnedExtents(file->file_path())) {
+            return false;
+        }
+    }
+    return true;
+}
+
 const std::vector<struct fiemap_extent>& SplitFiemap::extents() {
     if (extents_.empty()) {
         for (const auto& file : files_) {
@@ -186,6 +195,76 @@
     return extents_;
 }
 
+bool SplitFiemap::Write(const void* data, uint64_t bytes) {
+    // Open the current file.
+    FiemapWriter* file = files_[cursor_index_].get();
+
+    const uint8_t* data_ptr = reinterpret_cast<const uint8_t*>(data);
+    uint64_t bytes_remaining = bytes;
+    while (bytes_remaining) {
+        // How many bytes can we write into the current file?
+        uint64_t file_bytes_left = file->size() - cursor_file_pos_;
+        if (!file_bytes_left) {
+            if (cursor_index_ == files_.size() - 1) {
+                LOG(ERROR) << "write past end of file requested";
+                return false;
+            }
+
+            // No space left in the current file, but we have more files to
+            // use, so prep the next one.
+            cursor_fd_ = {};
+            cursor_file_pos_ = 0;
+            file = files_[++cursor_index_].get();
+            file_bytes_left = file->size();
+        }
+
+        // Open the current file if it's not open.
+        if (cursor_fd_ < 0) {
+            cursor_fd_.reset(open(file->file_path().c_str(), O_CLOEXEC | O_WRONLY));
+            if (cursor_fd_ < 0) {
+                PLOG(ERROR) << "open failed: " << file->file_path();
+                return false;
+            }
+            CHECK(cursor_file_pos_ == 0);
+        }
+
+        if (!FiemapWriter::HasPinnedExtents(file->file_path())) {
+            LOG(ERROR) << "file is no longer pinned: " << file->file_path();
+            return false;
+        }
+
+        uint64_t bytes_to_write = std::min(file_bytes_left, bytes_remaining);
+        if (!android::base::WriteFully(cursor_fd_, data_ptr, bytes_to_write)) {
+            PLOG(ERROR) << "write failed: " << file->file_path();
+            return false;
+        }
+        data_ptr += bytes_to_write;
+        bytes_remaining -= bytes_to_write;
+        cursor_file_pos_ += bytes_to_write;
+    }
+
+    // If we've reached the end of the current file, close it for sanity.
+    if (cursor_file_pos_ == file->size()) {
+        cursor_fd_ = {};
+    }
+    return true;
+}
+
+bool SplitFiemap::Flush() {
+    for (const auto& file : files_) {
+        unique_fd fd(open(file->file_path().c_str(), O_RDONLY | O_CLOEXEC));
+        if (fd < 0) {
+            PLOG(ERROR) << "open failed: " << file->file_path();
+            return false;
+        }
+        if (fsync(fd)) {
+            PLOG(ERROR) << "fsync failed: " << file->file_path();
+            return false;
+        }
+    }
+    return true;
+}
+
 SplitFiemap::~SplitFiemap() {
     if (!creating_) {
         return;