5 files changed, 73 insertions, 8 deletions

diff --git a/aidl/android/gsi/IGsiService.aidl b/aidl/android/gsi/IGsiService.aidl
index 34b1441..6aed1c1 100644
--- a/aidl/android/gsi/IGsiService.aidl
+++ b/aidl/android/gsi/IGsiService.aidl
@@ -56,12 +56,21 @@ interface IGsiService {
     GsiProgress getInstallProgress();
 
     /**
-     * Write bytes from memory to the on-disk GSI.
+     * Set the file descriptor that points to a ashmem which will be used
+     * to fetch data during the commitGsiChunkFromAshmem.
      *
-     * @param bytes         Byte array.
+     * @param stream        fd that points to a ashmem
+     * @param size          size of the ashmem file
+     */
+    boolean setGsiAshmem(in ParcelFileDescriptor stream, long size);
+
+    /**
+     * Write bytes from ashmem previously set with setGsiAshmem to GSI partition
+     *
+     * @param bytes         Number of bytes to submit
      * @return              true on success, false otherwise.
      */
-    boolean commitGsiChunkFromMemory(in byte[] bytes);
+    boolean commitGsiChunkFromAshmem(long bytes);
 
     /**
      * Complete a GSI installation and mark it as bootable. The caller is
diff --git a/gsi_installer.cpp b/gsi_installer.cpp
index e08e642..27e6dd0 100644
--- a/gsi_installer.cpp
+++ b/gsi_installer.cpp
@@ -75,6 +75,9 @@ GsiInstaller::~GsiInstaller() {
 
         GsiService::RemoveGsiFiles(install_dir_, wipe_userdata_on_failure_);
     }
+    if (IsAshmemMapped()) {
+        UnmapAshmem();
+    }
 }
 
 void GsiInstaller::PostInstallCleanup() {
@@ -256,6 +259,14 @@ bool GsiInstaller::CommitGsiChunk(int stream_fd, int64_t bytes) {
     return true;
 }
 
+bool GsiInstaller::IsFinishedWriting() {
+    return gsi_bytes_written_ == gsi_size_;
+}
+
+bool GsiInstaller::IsAshmemMapped() {
+    return ashmem_data_ != MAP_FAILED;
+}
+
 bool GsiInstaller::CommitGsiChunk(const void* data, size_t bytes) {
     if (static_cast<uint64_t>(bytes) > gsi_size_ - gsi_bytes_written_) {
         // We cannot write past the end of the image file.
@@ -274,6 +285,33 @@ bool GsiInstaller::CommitGsiChunk(const void* data, size_t bytes) {
     return true;
 }
 
+bool GsiInstaller::MapAshmem(int fd, size_t size) {
+    ashmem_size_ = size;
+    ashmem_data_ = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
+    return ashmem_data_ != MAP_FAILED;
+}
+
+void GsiInstaller::UnmapAshmem() {
+    if (munmap(ashmem_data_, ashmem_size_) != 0) {
+        PLOG(ERROR) << "cannot munmap";
+        return;
+    }
+    ashmem_data_ = MAP_FAILED;
+    ashmem_size_ = -1;
+}
+
+bool GsiInstaller::CommitGsiChunk(size_t bytes) {
+    if (!IsAshmemMapped()) {
+        PLOG(ERROR) << "ashmem is not mapped";
+        return false;
+    }
+    bool success = CommitGsiChunk(ashmem_data_, bytes);
+    if (success && IsFinishedWriting()) {
+        UnmapAshmem();
+    }
+    return success;
+}
+
 bool GsiInstaller::SetBootMode(bool one_shot) {
     if (one_shot) {
         if (!android::base::WriteStringToFile("1", kDsuOneShotBootFile)) {
diff --git a/gsi_installer.h b/gsi_installer.h
index 9f9a94f..33bebfb 100644
--- a/gsi_installer.h
+++ b/gsi_installer.h
@@ -16,10 +16,12 @@
 #pragma once
 
 #include <stdint.h>
+#include <sys/mman.h>
 
 #include <memory>
 #include <string>
 
+#include <android-base/unique_fd.h>
 #include <android/gsi/IGsiService.h>
 #include <android/gsi/MappedImage.h>
 #include <libfiemap/image_manager.h>
@@ -45,6 +47,8 @@ class GsiInstaller final {
     int StartInstall();
     bool CommitGsiChunk(int stream_fd, int64_t bytes);
     bool CommitGsiChunk(const void* data, size_t bytes);
+    bool MapAshmem(int fd, size_t size);
+    bool CommitGsiChunk(size_t bytes);
     int SetGsiBootable(bool one_shot);
 
     // Methods for interacting with an existing install.
@@ -69,6 +73,9 @@ class GsiInstaller final {
     int CheckInstallState();
     bool CreateInstallStatusFile();
     bool SetBootMode(bool one_shot);
+    bool IsFinishedWriting();
+    bool IsAshmemMapped();
+    void UnmapAshmem();
 
     GsiService* service_;
 
@@ -81,6 +88,8 @@ class GsiInstaller final {
     // Remaining data we're waiting to receive for the GSI image.
     uint64_t gsi_bytes_written_ = 0;
     bool succeeded_ = false;
+    uint64_t ashmem_size_ = -1;
+    void* ashmem_data_ = MAP_FAILED;
 
     std::unique_ptr<MappedDevice> system_device_;
 };
diff --git a/gsi_service.cpp b/gsi_service.cpp
index f9805c4..dcf2d83 100644
--- a/gsi_service.cpp
+++ b/gsi_service.cpp
@@ -169,8 +169,7 @@ binder::Status GsiService::getInstallProgress(::android::gsi::GsiProgress* _aidl
     return binder::Status::ok();
 }
 
-binder::Status GsiService::commitGsiChunkFromMemory(const std::vector<uint8_t>& bytes,
-                                                    bool* _aidl_return) {
+binder::Status GsiService::commitGsiChunkFromAshmem(int64_t bytes, bool* _aidl_return) {
     ENFORCE_SYSTEM;
     std::lock_guard<std::mutex> guard(parent_->lock());
 
@@ -178,8 +177,17 @@ binder::Status GsiService::commitGsiChunkFromMemory(const std::vector<uint8_t>&
         *_aidl_return = false;
         return binder::Status::ok();
     }
+    *_aidl_return = installer_->CommitGsiChunk(bytes);
+    return binder::Status::ok();
+}
 
-    *_aidl_return = installer_->CommitGsiChunk(bytes.data(), bytes.size());
+binder::Status GsiService::setGsiAshmem(const ::android::os::ParcelFileDescriptor& ashmem,
+                                        int64_t size, bool* _aidl_return) {
+    if (!installer_) {
+        *_aidl_return = false;
+        return binder::Status::ok();
+    }
+    *_aidl_return = installer_->MapAshmem(ashmem.get(), size);
     return binder::Status::ok();
 }
 
diff --git a/gsi_service.h b/gsi_service.h
index 00043dd..1df4aeb 100644
--- a/gsi_service.h
+++ b/gsi_service.h
@@ -61,8 +61,9 @@ class GsiService : public BinderService<GsiService>, public BnGsiService {
     binder::Status commitGsiChunkFromStream(const ::android::os::ParcelFileDescriptor& stream,
                                             int64_t bytes, bool* _aidl_return) override;
     binder::Status getInstallProgress(::android::gsi::GsiProgress* _aidl_return) override;
-    binder::Status commitGsiChunkFromMemory(const ::std::vector<uint8_t>& bytes,
-                                            bool* _aidl_return) override;
+    binder::Status setGsiAshmem(const ::android::os::ParcelFileDescriptor& ashmem, int64_t size,
+                                bool* _aidl_return) override;
+    binder::Status commitGsiChunkFromAshmem(int64_t bytes, bool* _aidl_return) override;
     binder::Status cancelGsiInstall(bool* _aidl_return) override;
     binder::Status enableGsi(bool oneShot, int* _aidl_return) override;
     binder::Status isGsiEnabled(bool* _aidl_return) override;