Add write support to SplitFiemap.

This adds few new methods to introduce write support to SplitFiemap:

(1) Write(), which has an internal cursor to stream data into the split
    files. As the end of one file is reached, the next is opened.
(2) Flush(), which calls fsync() on each internal FiemapWriter.
(3) HasPinnedExtents(), which calls the same on each internal
    FiemapWriter.

Included are some tests for edge cases in Write().

Bug: 126230649
Test: fiemap_writer_test gtest
Change-Id: I9fd509215975dbbb20a44b020315d3c1b287d1a0
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;