Merge "Remove from top layer list if re-parenting top layer."
diff --git a/cmds/installd/InstalldNativeService.cpp b/cmds/installd/InstalldNativeService.cpp
index e10fb6c..58786f6 100644
--- a/cmds/installd/InstalldNativeService.cpp
+++ b/cmds/installd/InstalldNativeService.cpp
@@ -436,7 +436,7 @@
             // profiles.
             int shared_app_gid = multiuser_get_shared_gid(0, appId);
             if ((shared_app_gid != -1) && fs_prepare_dir_strict(
-                    ref_profile_path.c_str(), 0700, shared_app_gid, shared_app_gid) != 0) {
+                    ref_profile_path.c_str(), 0701, shared_app_gid, shared_app_gid) != 0) {
                 return error("Failed to prepare " + ref_profile_path);
             }
         }
@@ -1833,6 +1833,29 @@
     return ok();
 }
 
+binder::Status InstalldNativeService::snapshotProfile(int32_t appId, const std::string& packageName,
+        const std::string& codePath, bool* _aidl_return) {
+    ENFORCE_UID(AID_SYSTEM);
+    CHECK_ARGUMENT_PACKAGE_NAME(packageName);
+    std::lock_guard<std::recursive_mutex> lock(mLock);
+
+    *_aidl_return = snapshot_profile(appId, packageName, codePath);
+    return ok();
+}
+
+binder::Status InstalldNativeService::destroyProfileSnapshot(const std::string& packageName,
+        const std::string& codePath) {
+    ENFORCE_UID(AID_SYSTEM);
+    CHECK_ARGUMENT_PACKAGE_NAME(packageName);
+    std::lock_guard<std::recursive_mutex> lock(mLock);
+
+    std::string snapshot = create_snapshot_profile_path(packageName, codePath);
+    if ((unlink(snapshot.c_str()) != 0) && (errno != ENOENT)) {
+        return error("Failed to destroy profile snapshot for " + packageName + ":" + codePath);
+    }
+    return ok();
+}
+
 binder::Status InstalldNativeService::dexopt(const std::string& apkPath, int32_t uid,
         const std::unique_ptr<std::string>& packageName, const std::string& instructionSet,
         int32_t dexoptNeeded, const std::unique_ptr<std::string>& outputPath, int32_t dexFlags,
diff --git a/cmds/installd/InstalldNativeService.h b/cmds/installd/InstalldNativeService.h
index 5167a14..eb7231c 100644
--- a/cmds/installd/InstalldNativeService.h
+++ b/cmds/installd/InstalldNativeService.h
@@ -96,6 +96,11 @@
     binder::Status clearAppProfiles(const std::string& packageName);
     binder::Status destroyAppProfiles(const std::string& packageName);
 
+    binder::Status snapshotProfile(int32_t appId, const std::string& packageName,
+            const std::string& codePath, bool* _aidl_return);
+    binder::Status destroyProfileSnapshot(const std::string& packageName,
+            const std::string& codePath);
+
     binder::Status idmap(const std::string& targetApkPath, const std::string& overlayApkPath,
             int32_t uid);
     binder::Status removeIdmap(const std::string& overlayApkPath);
diff --git a/cmds/installd/binder/android/os/IInstalld.aidl b/cmds/installd/binder/android/os/IInstalld.aidl
index 0f9862d..16ba2af 100644
--- a/cmds/installd/binder/android/os/IInstalld.aidl
+++ b/cmds/installd/binder/android/os/IInstalld.aidl
@@ -62,6 +62,9 @@
     void clearAppProfiles(@utf8InCpp String packageName);
     void destroyAppProfiles(@utf8InCpp String packageName);
 
+    boolean snapshotProfile(int appId, @utf8InCpp String packageName, @utf8InCpp String codePath);
+    void destroyProfileSnapshot(@utf8InCpp String packageName, @utf8InCpp String codePath);
+
     void idmap(@utf8InCpp String targetApkPath, @utf8InCpp String overlayApkPath, int uid);
     void removeIdmap(@utf8InCpp String overlayApkPath);
     void rmPackageDir(@utf8InCpp String packageDir);
diff --git a/cmds/installd/dexopt.cpp b/cmds/installd/dexopt.cpp
index 6f160b9..beffe00 100644
--- a/cmds/installd/dexopt.cpp
+++ b/cmds/installd/dexopt.cpp
@@ -566,14 +566,12 @@
     }
 }
 
-static bool create_profile(int uid, const std::string& profile) {
-    unique_fd fd(TEMP_FAILURE_RETRY(open(profile.c_str(), O_CREAT | O_NOFOLLOW, 0600)));
+static unique_fd create_profile(uid_t uid, const std::string& profile, int32_t flags) {
+    unique_fd fd(TEMP_FAILURE_RETRY(open(profile.c_str(), flags, 0600)));
     if (fd.get() < 0) {
-        if (errno == EEXIST) {
-            return true;
-        } else {
+        if (errno != EEXIST) {
             PLOG(ERROR) << "Failed to create profile " << profile;
-            return false;
+            return invalid_unique_fd();
         }
     }
     // Profiles should belong to the app; make sure of that by giving ownership to
@@ -581,27 +579,26 @@
     // since dex2oat/profman will fail with SElinux denials.
     if (fchown(fd.get(), uid, uid) < 0) {
         PLOG(ERROR) << "Could not chwon profile " << profile;
-        return false;
+        return invalid_unique_fd();
     }
-    return true;
+    return fd;
 }
 
-static unique_fd open_profile(int uid, const std::string& profile, bool read_write) {
-    // Check if we need to open the profile for a read-write operation. If so, we
-    // might need to create the profile since the file might not be there. Reference
-    // profiles are created on the fly so they might not exist beforehand.
-    if (read_write) {
-        if (!create_profile(uid, profile)) {
-            return invalid_unique_fd();
-        }
-    }
-    int flags = read_write ? O_RDWR : O_RDONLY;
+static unique_fd open_profile(uid_t uid, const std::string& profile, int32_t flags) {
     // Do not follow symlinks when opening a profile:
     //   - primary profiles should not contain symlinks in their paths
     //   - secondary dex paths should have been already resolved and validated
     flags |= O_NOFOLLOW;
 
-    unique_fd fd(TEMP_FAILURE_RETRY(open(profile.c_str(), flags)));
+    // Check if we need to create the profile
+    // Reference profiles and snapshots are created on the fly; so they might not exist beforehand.
+    unique_fd fd;
+    if ((flags & O_CREAT) != 0) {
+        fd = create_profile(uid, profile, flags);
+    } else {
+        fd.reset(TEMP_FAILURE_RETRY(open(profile.c_str(), flags)));
+    }
+
     if (fd.get() < 0) {
         if (errno != ENOENT) {
             // Profiles might be missing for various reasons. For example, in a
@@ -621,13 +618,19 @@
 static unique_fd open_current_profile(uid_t uid, userid_t user, const std::string& location,
         bool is_secondary_dex) {
     std::string profile = create_current_profile_path(user, location, is_secondary_dex);
-    return open_profile(uid, profile, /*read_write*/false);
+    return open_profile(uid, profile, O_RDONLY);
 }
 
 static unique_fd open_reference_profile(uid_t uid, const std::string& location, bool read_write,
         bool is_secondary_dex) {
     std::string profile = create_reference_profile_path(location, is_secondary_dex);
-    return open_profile(uid, profile, read_write);
+    return open_profile(uid, profile, read_write ? (O_CREAT | O_RDWR) : O_RDONLY);
+}
+
+static unique_fd open_spnashot_profile(uid_t uid, const std::string& package_name,
+        const std::string& code_path) {
+    std::string profile = create_snapshot_profile_path(package_name, code_path);
+    return open_profile(uid, profile, O_CREAT | O_RDWR | O_TRUNC);
 }
 
 static void open_profile_files(uid_t uid, const std::string& location, bool is_secondary_dex,
@@ -2396,5 +2399,42 @@
     }
 }
 
+bool snapshot_profile(int32_t app_id, const std::string& package_name,
+        const std::string& code_path) {
+    int app_shared_gid = multiuser_get_shared_gid(/*user_id*/ 0, app_id);
+
+    unique_fd snapshot_fd = open_spnashot_profile(AID_SYSTEM, package_name, code_path);
+    if (snapshot_fd < 0) {
+        return false;
+    }
+
+    std::vector<unique_fd> profiles_fd;
+    unique_fd reference_profile_fd;
+    open_profile_files(app_shared_gid, package_name, /*is_secondary_dex*/ false, &profiles_fd,
+            &reference_profile_fd);
+    if (profiles_fd.empty() || (reference_profile_fd.get() < 0)) {
+        return false;
+    }
+
+    profiles_fd.push_back(std::move(reference_profile_fd));
+
+    pid_t pid = fork();
+    if (pid == 0) {
+        /* child -- drop privileges before continuing */
+        drop_capabilities(app_shared_gid);
+        run_profman_merge(profiles_fd, snapshot_fd);
+        exit(42);   /* only get here on exec failure */
+    }
+
+    /* parent */
+    int return_code = wait_child(pid);
+    if (!WIFEXITED(return_code)) {
+        LOG(WARNING) << "profman failed for " << package_name << ":" << code_path;
+        return false;
+    }
+
+    return true;
+}
+
 }  // namespace installd
 }  // namespace android
diff --git a/cmds/installd/dexopt.h b/cmds/installd/dexopt.h
index 798e211..8a9ffd2 100644
--- a/cmds/installd/dexopt.h
+++ b/cmds/installd/dexopt.h
@@ -50,6 +50,17 @@
 // the reference profiles accessible with open_reference_profile().
 bool analyze_primary_profiles(uid_t uid, const std::string& pkgname);
 
+// Create a snapshot of the profile information for the given package and code path.
+// The profile snapshot is the aggregation of all existing profiles (all current user
+// profiles & the reference profile) and is meant to capture the all the profile information
+// without performing a merge into the reference profile which might impact future dex2oat
+// compilations.
+// The snapshot is created next to the reference profile of the package and the
+// ownership is assigned to AID_SYSTEM.
+// The snapshot location is reference_profile_location.snapshot. If a snapshot is already
+// there, it will be truncated and overwritten.
+bool snapshot_profile(int32_t app_id, const std::string& package, const std::string& code_path);
+
 bool dump_profiles(int32_t uid, const std::string& pkgname, const char* code_paths);
 
 bool copy_system_profile(const std::string& system_profile,
diff --git a/cmds/installd/tests/installd_dexopt_test.cpp b/cmds/installd/tests/installd_dexopt_test.cpp
index 821b96f..19b42b5 100644
--- a/cmds/installd/tests/installd_dexopt_test.cpp
+++ b/cmds/installd/tests/installd_dexopt_test.cpp
@@ -14,20 +14,33 @@
  * limitations under the License.
  */
 
+#include <fcntl.h>
 #include <stdlib.h>
 #include <string.h>
+#include <sys/types.h>
+#include <sys/stat.h>
 
+#include <android-base/file.h>
 #include <android-base/logging.h>
 #include <android-base/stringprintf.h>
+#include <android-base/unique_fd.h>
+
 #include <cutils/properties.h>
+
 #include <gtest/gtest.h>
 
+#include <selinux/android.h>
+#include <selinux/avc.h>
+
 #include "dexopt.h"
 #include "InstalldNativeService.h"
 #include "globals.h"
 #include "tests/test_utils.h"
 #include "utils.h"
 
+using android::base::ReadFully;
+using android::base::unique_fd;
+
 namespace android {
 namespace installd {
 
@@ -76,6 +89,42 @@
     ::chmod(path.c_str(), mode);
 }
 
+static int log_callback(int type, const char *fmt, ...) { // NOLINT
+    va_list ap;
+    int priority;
+
+    switch (type) {
+        case SELINUX_WARNING:
+            priority = ANDROID_LOG_WARN;
+            break;
+        case SELINUX_INFO:
+            priority = ANDROID_LOG_INFO;
+            break;
+        default:
+            priority = ANDROID_LOG_ERROR;
+            break;
+    }
+    va_start(ap, fmt);
+    LOG_PRI_VA(priority, "SELinux", fmt, ap);
+    va_end(ap);
+    return 0;
+}
+
+static bool init_selinux() {
+    int selinux_enabled = (is_selinux_enabled() > 0);
+
+    union selinux_callback cb;
+    cb.func_log = log_callback;
+    selinux_set_callback(SELINUX_CB_LOG, cb);
+
+    if (selinux_enabled && selinux_status_open(true) < 0) {
+        LOG(ERROR) << "Could not open selinux status; exiting";
+        return false;
+    }
+
+    return true;
+}
+
 // Base64 encoding of a simple dex files with 2 methods.
 static const char kDexFile[] =
     "UEsDBBQAAAAIAOiOYUs9y6BLCgEAABQCAAALABwAY2xhc3Nlcy5kZXhVVAkAA/Ns+lkOHv1ZdXgL"
@@ -97,6 +146,7 @@
     static constexpr int32_t kAppDataFlags = FLAG_STORAGE_CE | FLAG_STORAGE_DE;
     static constexpr uid_t kTestAppUid = 19999;
     static constexpr gid_t kTestAppGid = 19999;
+    static constexpr uid_t kTestAppId = kTestAppUid;
     static constexpr int32_t kTestUserId = 0;
 
     InstalldNativeService* service_;
@@ -116,15 +166,18 @@
     virtual void SetUp() {
         setenv("ANDROID_LOG_TAGS", "*:v", 1);
         android::base::InitLogging(nullptr);
-
+        // Initialize the globals holding the file system main paths (/data/, /system/ etc..).
+        // This is needed in order to compute the application and profile paths.
+        ASSERT_TRUE(init_globals_from_data_and_root());
+        // Initialize selinux log callbacks.
+        // This ensures that selinux is up and running and re-directs the selinux messages
+        // to logcat (in order to make it easier to investigate test results).
+        ASSERT_TRUE(init_selinux());
         service_ = new InstalldNativeService();
 
         volume_uuid_ = nullptr;
         package_name_ = "com.installd.test.dexopt";
         se_info_ = "default";
-
-        init_globals_from_data_and_root();
-
         app_apk_dir_ = android_app_dir + package_name_;
 
         create_mock_app();
@@ -183,14 +236,14 @@
     }
 
 
-    std::string get_secondary_dex_artifact(const std::string& path, const std::string& type) {
+    std::string GetSecondaryDexArtifact(const std::string& path, const std::string& type) {
         std::string::size_type end = path.rfind('.');
         std::string::size_type start = path.rfind('/', end);
         return path.substr(0, start) + "/oat/" + kRuntimeIsa + "/" +
                 path.substr(start + 1, end - start) + type;
     }
 
-    void compile_secondary_dex(const std::string& path, int32_t dex_storage_flag,
+    void CompileSecondaryDex(const std::string& path, int32_t dex_storage_flag,
             bool should_binder_call_succeed, bool should_dex_be_compiled = true,
             int uid = kTestAppUid) {
         std::unique_ptr<std::string> package_name_ptr(new std::string(package_name_));
@@ -216,9 +269,9 @@
                                                  downgrade);
         ASSERT_EQ(should_binder_call_succeed, result.isOk());
         int expected_access = should_dex_be_compiled ? 0 : -1;
-        std::string odex = get_secondary_dex_artifact(path, "odex");
-        std::string vdex = get_secondary_dex_artifact(path, "vdex");
-        std::string art = get_secondary_dex_artifact(path, "art");
+        std::string odex = GetSecondaryDexArtifact(path, "odex");
+        std::string vdex = GetSecondaryDexArtifact(path, "vdex");
+        std::string art = GetSecondaryDexArtifact(path, "art");
         ASSERT_EQ(expected_access, access(odex.c_str(), R_OK));
         ASSERT_EQ(expected_access, access(vdex.c_str(), R_OK));
         ASSERT_EQ(-1, access(art.c_str(), R_OK));  // empty profiles do not generate an image.
@@ -243,56 +296,64 @@
         ASSERT_EQ(should_dex_exist, out_secondary_dex_exists);
 
         int expected_access = should_dex_be_deleted ? -1 : 0;
-        std::string odex = get_secondary_dex_artifact(path, "odex");
-        std::string vdex = get_secondary_dex_artifact(path, "vdex");
-        std::string art = get_secondary_dex_artifact(path, "art");
+        std::string odex = GetSecondaryDexArtifact(path, "odex");
+        std::string vdex = GetSecondaryDexArtifact(path, "vdex");
+        std::string art = GetSecondaryDexArtifact(path, "art");
         ASSERT_EQ(expected_access, access(odex.c_str(), F_OK));
         ASSERT_EQ(expected_access, access(vdex.c_str(), F_OK));
         ASSERT_EQ(-1, access(art.c_str(), R_OK));  // empty profiles do not generate an image.
     }
+
+    void CheckFileAccess(const std::string& file, uid_t uid, gid_t gid, mode_t mode) {
+        struct stat st;
+        ASSERT_EQ(0, stat(file.c_str(), &st));
+        ASSERT_EQ(uid, st.st_uid);
+        ASSERT_EQ(gid, st.st_gid);
+        ASSERT_EQ(mode, st.st_mode);
+    }
 };
 
 
 TEST_F(DexoptTest, DexoptSecondaryCe) {
     LOG(INFO) << "DexoptSecondaryCe";
-    compile_secondary_dex(secondary_dex_ce_, DEXOPT_STORAGE_CE,
+    CompileSecondaryDex(secondary_dex_ce_, DEXOPT_STORAGE_CE,
         /*binder_ok*/ true, /*compile_ok*/ true);
 }
 
 TEST_F(DexoptTest, DexoptSecondaryCeLink) {
     LOG(INFO) << "DexoptSecondaryCeLink";
-    compile_secondary_dex(secondary_dex_ce_link_, DEXOPT_STORAGE_CE,
+    CompileSecondaryDex(secondary_dex_ce_link_, DEXOPT_STORAGE_CE,
         /*binder_ok*/ true, /*compile_ok*/ true);
 }
 
 TEST_F(DexoptTest, DexoptSecondaryDe) {
     LOG(INFO) << "DexoptSecondaryDe";
-    compile_secondary_dex(secondary_dex_de_, DEXOPT_STORAGE_DE,
+    CompileSecondaryDex(secondary_dex_de_, DEXOPT_STORAGE_DE,
         /*binder_ok*/ true, /*compile_ok*/ true);
 }
 
 TEST_F(DexoptTest, DexoptSecondaryDoesNotExist) {
     LOG(INFO) << "DexoptSecondaryDoesNotExist";
     // If the file validates but does not exist we do not treat it as an error.
-    compile_secondary_dex(secondary_dex_ce_ + "not.there", DEXOPT_STORAGE_CE,
+    CompileSecondaryDex(secondary_dex_ce_ + "not.there", DEXOPT_STORAGE_CE,
         /*binder_ok*/ true,  /*compile_ok*/ false);
 }
 
 TEST_F(DexoptTest, DexoptSecondaryStorageValidationError) {
     LOG(INFO) << "DexoptSecondaryStorageValidationError";
-    compile_secondary_dex(secondary_dex_ce_, DEXOPT_STORAGE_DE,
+    CompileSecondaryDex(secondary_dex_ce_, DEXOPT_STORAGE_DE,
         /*binder_ok*/ false,  /*compile_ok*/ false);
 }
 
 TEST_F(DexoptTest, DexoptSecondaryAppOwnershipValidationError) {
     LOG(INFO) << "DexoptSecondaryAppOwnershipValidationError";
-    compile_secondary_dex("/data/data/random.app/secondary.jar", DEXOPT_STORAGE_CE,
+    CompileSecondaryDex("/data/data/random.app/secondary.jar", DEXOPT_STORAGE_CE,
         /*binder_ok*/ false,  /*compile_ok*/ false);
 }
 
 TEST_F(DexoptTest, DexoptSecondaryAcessViaDifferentUidError) {
     LOG(INFO) << "DexoptSecondaryAcessViaDifferentUidError";
-    compile_secondary_dex(secondary_dex_ce_, DEXOPT_STORAGE_CE,
+    CompileSecondaryDex(secondary_dex_ce_, DEXOPT_STORAGE_CE,
         /*binder_ok*/ false,  /*compile_ok*/ false, kSystemUid);
 }
 
@@ -300,9 +361,9 @@
 class ReconcileTest : public DexoptTest {
     virtual void SetUp() {
         DexoptTest::SetUp();
-        compile_secondary_dex(secondary_dex_ce_, DEXOPT_STORAGE_CE,
+        CompileSecondaryDex(secondary_dex_ce_, DEXOPT_STORAGE_CE,
             /*binder_ok*/ true, /*compile_ok*/ true);
-        compile_secondary_dex(secondary_dex_de_, DEXOPT_STORAGE_DE,
+        CompileSecondaryDex(secondary_dex_de_, DEXOPT_STORAGE_DE,
             /*binder_ok*/ true, /*compile_ok*/ true);
     }
 };
@@ -355,5 +416,145 @@
         /*binder_ok*/ true, /*dex_ok */ false, /*odex_deleted*/ false, kSystemUid);
 }
 
+class ProfileTest : public DexoptTest {
+  protected:
+    std::string cur_profile_;
+    std::string ref_profile_;
+    std::string snap_profile_;
+
+    virtual void SetUp() {
+        DexoptTest::SetUp();
+        cur_profile_ = create_current_profile_path(
+                kTestUserId, package_name_, /*is_secondary_dex*/ false);
+        ref_profile_ = create_reference_profile_path(package_name_, /*is_secondary_dex*/ false);
+        snap_profile_ = create_snapshot_profile_path(package_name_, "base.jar");
+    }
+
+    void SetupProfile(const std::string& path, uid_t uid, gid_t gid, mode_t mode, int32_t seed) {
+        run_cmd("profman --generate-test-profile-seed=" + std::to_string(seed) +
+                " --generate-test-profile-num-dex=2 --generate-test-profile=" + path);
+        ::chmod(path.c_str(), mode);
+        ::chown(path.c_str(), uid, gid);
+    }
+
+    void SetupProfiles(bool setup_ref) {
+        SetupProfile(cur_profile_, kTestAppUid, kTestAppGid, 0600, 1);
+        if (setup_ref) {
+            SetupProfile(ref_profile_, kTestAppUid, kTestAppGid, 0060, 2);
+        }
+    }
+
+    void SnapshotProfile(int32_t appid, const std::string& package_name, bool expected_result) {
+        bool result;
+        binder::Status binder_result = service_->snapshotProfile(
+                appid, package_name, "base.jar", &result);
+        ASSERT_TRUE(binder_result.isOk());
+        ASSERT_EQ(expected_result, result);
+
+        if (!expected_result) {
+            // Do not check the files if we expect to fail.
+            return;
+        }
+
+        // Check that the snapshot was created witht he expected acess flags.
+        CheckFileAccess(snap_profile_, kSystemUid, kSystemGid, 0600 | S_IFREG);
+
+        // The snapshot should be equivalent to the merge of profiles.
+        std::string expected_profile_content = snap_profile_ + ".expected";
+        run_cmd("rm -f " + expected_profile_content);
+        run_cmd("touch " + expected_profile_content);
+        run_cmd("profman --profile-file=" + cur_profile_ +
+                " --profile-file=" + ref_profile_ +
+                " --reference-profile-file=" + expected_profile_content);
+
+        ASSERT_TRUE(AreFilesEqual(expected_profile_content, snap_profile_));
+
+        pid_t pid = fork();
+        if (pid == 0) {
+            /* child */
+            TransitionToSystemServer();
+
+            // System server should be able to open the the spanshot.
+            unique_fd fd(open(snap_profile_.c_str(), O_RDONLY));
+            ASSERT_TRUE(fd > -1) << "Failed to open profile as kSystemUid: " << strerror(errno);
+            _exit(0);
+        }
+        /* parent */
+        ASSERT_TRUE(WIFEXITED(wait_child(pid)));
+    }
+
+  private:
+    void TransitionToSystemServer() {
+        ASSERT_TRUE(DropCapabilities(kSystemUid, kSystemGid));
+        int32_t res = selinux_android_setcontext(
+                kSystemUid, true, se_info_.c_str(), "system_server");
+        ASSERT_EQ(0, res) << "Failed to setcon " << strerror(errno);
+    }
+
+    bool AreFilesEqual(const std::string& file1, const std::string& file2) {
+        std::vector<uint8_t> content1;
+        std::vector<uint8_t> content2;
+
+        if (!ReadAll(file1, &content1)) return false;
+        if (!ReadAll(file2, &content2)) return false;
+        return content1 == content2;
+    }
+
+    bool ReadAll(const std::string& file, std::vector<uint8_t>* content) {
+        unique_fd fd(open(file.c_str(), O_RDONLY));
+        if (fd < 0) {
+            PLOG(ERROR) << "Failed to open " << file;
+            return false;
+        }
+        struct stat st;
+        if (fstat(fd, &st) != 0) {
+            PLOG(ERROR) << "Failed to stat " << file;
+            return false;
+        }
+        content->resize(st.st_size);
+        bool result = ReadFully(fd, content->data(), content->size());
+        if (!result) {
+            PLOG(ERROR) << "Failed to read " << file;
+        }
+        return result;
+    }
+};
+
+TEST_F(ProfileTest, ProfileSnapshotOk) {
+    LOG(INFO) << "ProfileSnapshotOk";
+
+    SetupProfiles(/*setup_ref*/ true);
+    SnapshotProfile(kTestAppId, package_name_, /*expected_result*/ true);
+}
+
+// The reference profile is created on the fly. We need to be able to
+// snapshot without one.
+TEST_F(ProfileTest, ProfileSnapshotOkNoReference) {
+    LOG(INFO) << "ProfileSnapshotOkNoReference";
+
+    SetupProfiles(/*setup_ref*/ false);
+    SnapshotProfile(kTestAppId, package_name_, /*expected_result*/ true);
+}
+
+TEST_F(ProfileTest, ProfileSnapshotFailWrongPackage) {
+    LOG(INFO) << "ProfileSnapshotFailWrongPackage";
+
+    SetupProfiles(/*setup_ref*/ true);
+    SnapshotProfile(kTestAppId, "not.there", /*expected_result*/ false);
+}
+
+TEST_F(ProfileTest, ProfileSnapshotDestroySnapshot) {
+    LOG(INFO) << "ProfileSnapshotDestroySnapshot";
+
+    SetupProfiles(/*setup_ref*/ true);
+    SnapshotProfile(kTestAppId, package_name_, /*expected_result*/ true);
+
+    binder::Status binder_result = service_->destroyProfileSnapshot(package_name_, "base.jar");
+    ASSERT_TRUE(binder_result.isOk());
+    struct stat st;
+    ASSERT_EQ(-1, stat(snap_profile_.c_str(), &st));
+    ASSERT_EQ(ENOENT, errno);
+}
+
 }  // namespace installd
 }  // namespace android
diff --git a/cmds/installd/tests/installd_utils_test.cpp b/cmds/installd/tests/installd_utils_test.cpp
index a6a4451..49da85d 100644
--- a/cmds/installd/tests/installd_utils_test.cpp
+++ b/cmds/installd/tests/installd_utils_test.cpp
@@ -350,6 +350,12 @@
             create_reference_profile_path("com.example", /*is_secondary*/false));
 }
 
+TEST_F(UtilsTest, CreateProfileSnapshot) {
+    std::string expected =
+        create_primary_reference_profile_package_dir_path("com.example") + "/primary.prof.snapshot";
+    EXPECT_EQ(expected, create_snapshot_profile_path("com.example", "base.apk"));
+}
+
 TEST_F(UtilsTest, CreateSecondaryCurrentProfile) {
     EXPECT_EQ("/data/user/0/com.example/oat/secondary.dex.cur.prof",
             create_current_profile_path(/*user*/0,
diff --git a/cmds/installd/tests/test_utils.h b/cmds/installd/tests/test_utils.h
index 7d1162e..b8785c6 100644
--- a/cmds/installd/tests/test_utils.h
+++ b/cmds/installd/tests/test_utils.h
@@ -1,6 +1,9 @@
-#include <android-base/logging.h>
 #include <stdlib.h>
 #include <string.h>
+#include <sys/capability.h>
+
+#include <android-base/logging.h>
+#include <selinux/android.h>
 
 uint8_t kBase64Map[256] = {
     255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
@@ -105,3 +108,27 @@
     }
     return true;
 }
+
+// TODO(calin): fix dexopt drop_capabilities and move to general utils (b/69678790).
+bool DropCapabilities(uid_t uid, gid_t gid) {
+    if (setgid(gid) != 0) {
+        PLOG(ERROR) << "setgid failed: " <<  gid;
+        return false;
+    }
+    if (setuid(uid) != 0) {
+        PLOG(ERROR) << "setuid failed: " <<  uid;
+        return false;
+    }
+    // drop capabilities
+    struct __user_cap_header_struct capheader;
+    struct __user_cap_data_struct capdata[2];
+    memset(&capheader, 0, sizeof(capheader));
+    memset(&capdata, 0, sizeof(capdata));
+    capheader.version = _LINUX_CAPABILITY_VERSION_3;
+    if (capset(&capheader, &capdata[0]) < 0) {
+        PLOG(ERROR) << "capset failed";
+        return false;
+    }
+
+    return true;
+}
diff --git a/cmds/installd/utils.cpp b/cmds/installd/utils.cpp
index 7dca7c6..61c9c8f 100644
--- a/cmds/installd/utils.cpp
+++ b/cmds/installd/utils.cpp
@@ -236,6 +236,7 @@
 const std::string PROFILE_EXT = ".prof";
 const std::string CURRENT_PROFILE_EXT = ".cur";
 const std::string PRIMARY_PROFILE_NAME = "primary" + PROFILE_EXT;
+const std::string SNAPSHOT_PROFILE_EXT = ".snapshot";
 
 // Gets the parent directory and the file name for the given secondary dex path.
 // Returns true on success, false on failure (if the dex_path does not have the expected
@@ -289,6 +290,14 @@
     }
 }
 
+std::string create_snapshot_profile_path(const std::string& package,
+        const std::string& code_path ATTRIBUTE_UNUSED) {
+    // TODD(calin): code_path is ignored for now. It will be used when each split gets its own
+    // profile file.
+    std::string ref_profile = create_reference_profile_path(package, /*is_secondary_dex*/ false);
+    return ref_profile + SNAPSHOT_PROFILE_EXT;
+}
+
 std::vector<userid_t> get_known_users(const char* volume_uuid) {
     std::vector<userid_t> users;
 
diff --git a/cmds/installd/utils.h b/cmds/installd/utils.h
index b90caf9..5391061 100644
--- a/cmds/installd/utils.h
+++ b/cmds/installd/utils.h
@@ -83,6 +83,7 @@
         userid_t user, const std::string& package_name, bool is_secondary_dex);
 std::string create_reference_profile_path(
         const std::string& package_name, bool is_secondary_dex);
+std::string create_snapshot_profile_path(const std::string& package, const std::string& code_path);
 
 std::vector<userid_t> get_known_users(const char* volume_uuid);
 
diff --git a/libs/gui/ISurfaceComposer.cpp b/libs/gui/ISurfaceComposer.cpp
index c21c5e3..5de84ec 100644
--- a/libs/gui/ISurfaceComposer.cpp
+++ b/libs/gui/ISurfaceComposer.cpp
@@ -100,17 +100,13 @@
         remote()->transact(BnSurfaceComposer::BOOT_FINISHED, data, &reply);
     }
 
-    virtual status_t captureScreen(const sp<IBinder>& display,
-            const sp<IGraphicBufferProducer>& producer,
-            Rect sourceCrop, uint32_t reqWidth, uint32_t reqHeight,
-            int32_t minLayerZ, int32_t maxLayerZ,
-            bool useIdentityTransform,
-            ISurfaceComposer::Rotation rotation)
-    {
+    virtual status_t captureScreen(const sp<IBinder>& display, sp<GraphicBuffer>* outBuffer,
+                                   Rect sourceCrop, uint32_t reqWidth, uint32_t reqHeight,
+                                   int32_t minLayerZ, int32_t maxLayerZ, bool useIdentityTransform,
+                                   ISurfaceComposer::Rotation rotation) {
         Parcel data, reply;
         data.writeInterfaceToken(ISurfaceComposer::getInterfaceDescriptor());
         data.writeStrongBinder(display);
-        data.writeStrongBinder(IInterface::asBinder(producer));
         data.write(sourceCrop);
         data.writeUint32(reqWidth);
         data.writeUint32(reqHeight);
@@ -118,21 +114,45 @@
         data.writeInt32(maxLayerZ);
         data.writeInt32(static_cast<int32_t>(useIdentityTransform));
         data.writeInt32(static_cast<int32_t>(rotation));
-        remote()->transact(BnSurfaceComposer::CAPTURE_SCREEN, data, &reply);
-        return reply.readInt32();
+        status_t err = remote()->transact(BnSurfaceComposer::CAPTURE_SCREEN, data, &reply);
+
+        if (err != NO_ERROR) {
+            return err;
+        }
+
+        err = reply.readInt32();
+        if (err != NO_ERROR) {
+            return err;
+        }
+
+        *outBuffer = new GraphicBuffer();
+        reply.read(**outBuffer);
+        return err;
     }
 
     virtual status_t captureLayers(const sp<IBinder>& layerHandleBinder,
-                                   const sp<IGraphicBufferProducer>& producer,
-                                   const Rect& sourceCrop, float frameScale) {
+                                   sp<GraphicBuffer>* outBuffer, const Rect& sourceCrop,
+                                   float frameScale) {
         Parcel data, reply;
         data.writeInterfaceToken(ISurfaceComposer::getInterfaceDescriptor());
         data.writeStrongBinder(layerHandleBinder);
-        data.writeStrongBinder(IInterface::asBinder(producer));
         data.write(sourceCrop);
         data.writeFloat(frameScale);
-        remote()->transact(BnSurfaceComposer::CAPTURE_LAYERS, data, &reply);
-        return reply.readInt32();
+        status_t err = remote()->transact(BnSurfaceComposer::CAPTURE_LAYERS, data, &reply);
+
+        if (err != NO_ERROR) {
+            return err;
+        }
+
+        err = reply.readInt32();
+        if (err != NO_ERROR) {
+            return err;
+        }
+
+        *outBuffer = new GraphicBuffer();
+        reply.read(**outBuffer);
+
+        return err;
     }
 
     virtual bool authenticateSurfaceTexture(
@@ -583,8 +603,7 @@
         case CAPTURE_SCREEN: {
             CHECK_INTERFACE(ISurfaceComposer, data, reply);
             sp<IBinder> display = data.readStrongBinder();
-            sp<IGraphicBufferProducer> producer =
-                    interface_cast<IGraphicBufferProducer>(data.readStrongBinder());
+            sp<GraphicBuffer> outBuffer;
             Rect sourceCrop(Rect::EMPTY_RECT);
             data.read(sourceCrop);
             uint32_t reqWidth = data.readUint32();
@@ -594,24 +613,28 @@
             bool useIdentityTransform = static_cast<bool>(data.readInt32());
             int32_t rotation = data.readInt32();
 
-            status_t res = captureScreen(display, producer,
-                    sourceCrop, reqWidth, reqHeight, minLayerZ, maxLayerZ,
-                    useIdentityTransform,
-                    static_cast<ISurfaceComposer::Rotation>(rotation));
+            status_t res = captureScreen(display, &outBuffer, sourceCrop, reqWidth, reqHeight,
+                                         minLayerZ, maxLayerZ, useIdentityTransform,
+                                         static_cast<ISurfaceComposer::Rotation>(rotation));
             reply->writeInt32(res);
+            if (res == NO_ERROR) {
+                reply->write(*outBuffer);
+            }
             return NO_ERROR;
         }
         case CAPTURE_LAYERS: {
             CHECK_INTERFACE(ISurfaceComposer, data, reply);
             sp<IBinder> layerHandleBinder = data.readStrongBinder();
-            sp<IGraphicBufferProducer> producer =
-                    interface_cast<IGraphicBufferProducer>(data.readStrongBinder());
+            sp<GraphicBuffer> outBuffer;
             Rect sourceCrop(Rect::EMPTY_RECT);
             data.read(sourceCrop);
             float frameScale = data.readFloat();
 
-            status_t res = captureLayers(layerHandleBinder, producer, sourceCrop, frameScale);
+            status_t res = captureLayers(layerHandleBinder, &outBuffer, sourceCrop, frameScale);
             reply->writeInt32(res);
+            if (res == NO_ERROR) {
+                reply->write(*outBuffer);
+            }
             return NO_ERROR;
         }
         case AUTHENTICATE_SURFACE: {
diff --git a/libs/gui/Surface.cpp b/libs/gui/Surface.cpp
index e939383..80216bc 100644
--- a/libs/gui/Surface.cpp
+++ b/libs/gui/Surface.cpp
@@ -1760,4 +1760,25 @@
     return OK;
 }
 
+status_t Surface::attachAndQueueBuffer(Surface* surface, sp<GraphicBuffer> buffer) {
+    if (buffer == nullptr) {
+        return BAD_VALUE;
+    }
+    int err = static_cast<ANativeWindow*>(surface)->perform(surface, NATIVE_WINDOW_API_CONNECT,
+                                                            NATIVE_WINDOW_API_CPU);
+    if (err != OK) {
+        return err;
+    }
+    err = surface->attachBuffer(buffer->getNativeBuffer());
+    if (err != OK) {
+        return err;
+    }
+    err = static_cast<ANativeWindow*>(surface)->queueBuffer(surface, buffer->getNativeBuffer(), -1);
+    if (err != OK) {
+        return err;
+    }
+    err = surface->disconnect(NATIVE_WINDOW_API_CPU);
+    return err;
+}
+
 }; // namespace android
diff --git a/libs/gui/SurfaceComposerClient.cpp b/libs/gui/SurfaceComposerClient.cpp
index 2466d25..0dbe786 100644
--- a/libs/gui/SurfaceComposerClient.cpp
+++ b/libs/gui/SurfaceComposerClient.cpp
@@ -717,181 +717,27 @@
 
 // ----------------------------------------------------------------------------
 
-status_t ScreenshotClient::capture(
-        const sp<IBinder>& display,
-        const sp<IGraphicBufferProducer>& producer,
-        Rect sourceCrop, uint32_t reqWidth, uint32_t reqHeight,
-        int32_t minLayerZ, int32_t maxLayerZ, bool useIdentityTransform) {
+status_t ScreenshotClient::capture(const sp<IBinder>& display, Rect sourceCrop, uint32_t reqWidth,
+                                   uint32_t reqHeight, int32_t minLayerZ, int32_t maxLayerZ,
+                                   bool useIdentityTransform, uint32_t rotation,
+                                   sp<GraphicBuffer>* outBuffer) {
     sp<ISurfaceComposer> s(ComposerService::getComposerService());
     if (s == NULL) return NO_INIT;
-    return s->captureScreen(display, producer, sourceCrop,
-            reqWidth, reqHeight, minLayerZ, maxLayerZ, useIdentityTransform);
-}
-
-status_t ScreenshotClient::captureToBuffer(const sp<IBinder>& display,
-        Rect sourceCrop, uint32_t reqWidth, uint32_t reqHeight,
-        int32_t minLayerZ, int32_t maxLayerZ, bool useIdentityTransform,
-        uint32_t rotation,
-        sp<GraphicBuffer>* outBuffer) {
-    sp<ISurfaceComposer> s(ComposerService::getComposerService());
-    if (s == NULL) return NO_INIT;
-
-    sp<IGraphicBufferConsumer> gbpConsumer;
-    sp<IGraphicBufferProducer> producer;
-    BufferQueue::createBufferQueue(&producer, &gbpConsumer);
-    sp<BufferItemConsumer> consumer(new BufferItemConsumer(gbpConsumer,
-           GRALLOC_USAGE_HW_TEXTURE | GRALLOC_USAGE_SW_READ_NEVER | GRALLOC_USAGE_SW_WRITE_NEVER,
-           1, true));
-
-    status_t ret = s->captureScreen(display, producer, sourceCrop, reqWidth, reqHeight,
-            minLayerZ, maxLayerZ, useIdentityTransform,
-            static_cast<ISurfaceComposer::Rotation>(rotation));
+    status_t ret = s->captureScreen(display, outBuffer, sourceCrop, reqWidth, reqHeight, minLayerZ,
+                                    maxLayerZ, useIdentityTransform,
+                                    static_cast<ISurfaceComposer::Rotation>(rotation));
     if (ret != NO_ERROR) {
         return ret;
     }
-    BufferItem b;
-    consumer->acquireBuffer(&b, 0, true);
-    *outBuffer = b.mGraphicBuffer;
     return ret;
 }
 
-status_t ScreenshotClient::captureLayers(const sp<IBinder>& layerHandle,
-                                         const sp<IGraphicBufferProducer>& producer,
-                                         Rect sourceCrop, float frameScale) {
+status_t ScreenshotClient::captureLayers(const sp<IBinder>& layerHandle, Rect sourceCrop,
+                                         float frameScale, sp<GraphicBuffer>* outBuffer) {
     sp<ISurfaceComposer> s(ComposerService::getComposerService());
     if (s == NULL) return NO_INIT;
-    return s->captureLayers(layerHandle, producer, sourceCrop, frameScale);
-}
-
-status_t ScreenshotClient::captureLayersToBuffer(const sp<IBinder>& layerHandle, Rect sourceCrop,
-                                                 float frameScale, sp<GraphicBuffer>* outBuffer) {
-    sp<ISurfaceComposer> s(ComposerService::getComposerService());
-    if (s == NULL) return NO_INIT;
-
-    sp<IGraphicBufferConsumer> gbpConsumer;
-    sp<IGraphicBufferProducer> producer;
-    BufferQueue::createBufferQueue(&producer, &gbpConsumer);
-    sp<BufferItemConsumer> consumer(new BufferItemConsumer(gbpConsumer,
-                                                           GRALLOC_USAGE_HW_TEXTURE |
-                                                                   GRALLOC_USAGE_SW_READ_NEVER |
-                                                                   GRALLOC_USAGE_SW_WRITE_NEVER,
-                                                           1, true));
-
-    status_t ret = s->captureLayers(layerHandle, producer, sourceCrop, frameScale);
-    if (ret != NO_ERROR) {
-        return ret;
-    }
-    BufferItem b;
-    consumer->acquireBuffer(&b, 0, true);
-    *outBuffer = b.mGraphicBuffer;
+    status_t ret = s->captureLayers(layerHandle, outBuffer, sourceCrop, frameScale);
     return ret;
 }
-
-ScreenshotClient::ScreenshotClient()
-    : mHaveBuffer(false) {
-    memset(&mBuffer, 0, sizeof(mBuffer));
-}
-
-ScreenshotClient::~ScreenshotClient() {
-    ScreenshotClient::release();
-}
-
-sp<CpuConsumer> ScreenshotClient::getCpuConsumer() const {
-    if (mCpuConsumer == NULL) {
-        sp<IGraphicBufferConsumer> consumer;
-        BufferQueue::createBufferQueue(&mProducer, &consumer);
-        mCpuConsumer = new CpuConsumer(consumer, 1);
-        mCpuConsumer->setName(String8("ScreenshotClient"));
-    }
-    return mCpuConsumer;
-}
-
-status_t ScreenshotClient::update(const sp<IBinder>& display,
-        Rect sourceCrop, uint32_t reqWidth, uint32_t reqHeight,
-        int32_t minLayerZ, int32_t maxLayerZ,
-        bool useIdentityTransform, uint32_t rotation) {
-    sp<ISurfaceComposer> s(ComposerService::getComposerService());
-    if (s == NULL) return NO_INIT;
-    sp<CpuConsumer> cpuConsumer = getCpuConsumer();
-
-    if (mHaveBuffer) {
-        mCpuConsumer->unlockBuffer(mBuffer);
-        memset(&mBuffer, 0, sizeof(mBuffer));
-        mHaveBuffer = false;
-    }
-
-    status_t err = s->captureScreen(display, mProducer, sourceCrop,
-            reqWidth, reqHeight, minLayerZ, maxLayerZ, useIdentityTransform,
-            static_cast<ISurfaceComposer::Rotation>(rotation));
-
-    if (err == NO_ERROR) {
-        err = mCpuConsumer->lockNextBuffer(&mBuffer);
-        if (err == NO_ERROR) {
-            mHaveBuffer = true;
-        }
-    }
-    return err;
-}
-
-status_t ScreenshotClient::update(const sp<IBinder>& display,
-        Rect sourceCrop, uint32_t reqWidth, uint32_t reqHeight,
-        int32_t minLayerZ, int32_t maxLayerZ,
-        bool useIdentityTransform) {
-
-    return ScreenshotClient::update(display, sourceCrop, reqWidth, reqHeight,
-            minLayerZ, maxLayerZ, useIdentityTransform, ISurfaceComposer::eRotateNone);
-}
-
-status_t ScreenshotClient::update(const sp<IBinder>& display, Rect sourceCrop,
-        bool useIdentityTransform) {
-    return ScreenshotClient::update(display, sourceCrop, 0, 0,
-            INT32_MIN, INT32_MAX,
-            useIdentityTransform, ISurfaceComposer::eRotateNone);
-}
-
-status_t ScreenshotClient::update(const sp<IBinder>& display, Rect sourceCrop,
-        uint32_t reqWidth, uint32_t reqHeight, bool useIdentityTransform) {
-    return ScreenshotClient::update(display, sourceCrop, reqWidth, reqHeight,
-            INT32_MIN, INT32_MAX,
-            useIdentityTransform, ISurfaceComposer::eRotateNone);
-}
-
-void ScreenshotClient::release() {
-    if (mHaveBuffer) {
-        mCpuConsumer->unlockBuffer(mBuffer);
-        memset(&mBuffer, 0, sizeof(mBuffer));
-        mHaveBuffer = false;
-    }
-    mCpuConsumer.clear();
-}
-
-void const* ScreenshotClient::getPixels() const {
-    return mBuffer.data;
-}
-
-uint32_t ScreenshotClient::getWidth() const {
-    return mBuffer.width;
-}
-
-uint32_t ScreenshotClient::getHeight() const {
-    return mBuffer.height;
-}
-
-PixelFormat ScreenshotClient::getFormat() const {
-    return mBuffer.format;
-}
-
-uint32_t ScreenshotClient::getStride() const {
-    return mBuffer.stride;
-}
-
-size_t ScreenshotClient::getSize() const {
-    return mBuffer.stride * mBuffer.height * bytesPerPixel(mBuffer.format);
-}
-
-android_dataspace ScreenshotClient::getDataSpace() const {
-    return mBuffer.dataSpace;
-}
-
 // ----------------------------------------------------------------------------
 }; // namespace android
diff --git a/libs/gui/include/gui/ISurfaceComposer.h b/libs/gui/include/gui/ISurfaceComposer.h
index 1e4c329..e26e332 100644
--- a/libs/gui/include/gui/ISurfaceComposer.h
+++ b/libs/gui/include/gui/ISurfaceComposer.h
@@ -29,6 +29,7 @@
 
 #include <ui/FrameStats.h>
 #include <ui/PixelFormat.h>
+#include <ui/GraphicBuffer.h>
 
 #include <vector>
 
@@ -167,16 +168,14 @@
     /* Capture the specified screen. requires READ_FRAME_BUFFER permission
      * This function will fail if there is a secure window on screen.
      */
-    virtual status_t captureScreen(const sp<IBinder>& display,
-            const sp<IGraphicBufferProducer>& producer,
-            Rect sourceCrop, uint32_t reqWidth, uint32_t reqHeight,
-            int32_t minLayerZ, int32_t maxLayerZ,
-            bool useIdentityTransform,
-            Rotation rotation = eRotateNone) = 0;
+    virtual status_t captureScreen(const sp<IBinder>& display, sp<GraphicBuffer>* outBuffer,
+                                   Rect sourceCrop, uint32_t reqWidth, uint32_t reqHeight,
+                                   int32_t minLayerZ, int32_t maxLayerZ, bool useIdentityTransform,
+                                   Rotation rotation = eRotateNone) = 0;
 
     virtual status_t captureLayers(const sp<IBinder>& layerHandleBinder,
-                                   const sp<IGraphicBufferProducer>& producer,
-                                   const Rect& sourceCrop, float frameScale = 1.0) = 0;
+                                   sp<GraphicBuffer>* outBuffer, const Rect& sourceCrop,
+                                   float frameScale = 1.0) = 0;
 
     /* Clears the frame statistics for animations.
      *
diff --git a/libs/gui/include/gui/Surface.h b/libs/gui/include/gui/Surface.h
index 3fe29d9..354f23a 100644
--- a/libs/gui/include/gui/Surface.h
+++ b/libs/gui/include/gui/Surface.h
@@ -283,6 +283,8 @@
 
     android_dataspace_t getBuffersDataSpace();
 
+    static status_t attachAndQueueBuffer(Surface* surface, sp<GraphicBuffer> buffer);
+
 protected:
     enum { NUM_BUFFER_SLOTS = BufferQueueDefs::NUM_BUFFER_SLOTS };
     enum { DEFAULT_FORMAT = PIXEL_FORMAT_RGBA_8888 };
diff --git a/libs/gui/include/gui/SurfaceComposerClient.h b/libs/gui/include/gui/SurfaceComposerClient.h
index b0fa922..3f13946 100644
--- a/libs/gui/include/gui/SurfaceComposerClient.h
+++ b/libs/gui/include/gui/SurfaceComposerClient.h
@@ -279,74 +279,16 @@
 
 // ---------------------------------------------------------------------------
 
-class ScreenshotClient
-{
+class ScreenshotClient {
 public:
     // if cropping isn't required, callers may pass in a default Rect, e.g.:
     //   capture(display, producer, Rect(), reqWidth, ...);
-    static status_t capture(
-            const sp<IBinder>& display,
-            const sp<IGraphicBufferProducer>& producer,
-            Rect sourceCrop, uint32_t reqWidth, uint32_t reqHeight,
-            int32_t minLayerZ, int32_t maxLayerZ,
-            bool useIdentityTransform);
-    static status_t captureToBuffer(
-            const sp<IBinder>& display,
-            Rect sourceCrop, uint32_t reqWidth, uint32_t reqHeight,
-            int32_t minLayerZ, int32_t maxLayerZ,
-            bool useIdentityTransform,
-            uint32_t rotation,
-            sp<GraphicBuffer>* outbuffer);
-
-    static status_t captureLayers(const sp<IBinder>& layerHandle,
-                                  const sp<IGraphicBufferProducer>& producer, Rect sourceCrop,
-                                        float frameScale);
-    static status_t captureLayersToBuffer(const sp<IBinder>& layerHandle, Rect sourceCrop,
-                                          float frameScale, sp<GraphicBuffer>* outBuffer);
-
-private:
-    mutable sp<CpuConsumer> mCpuConsumer;
-    mutable sp<IGraphicBufferProducer> mProducer;
-    CpuConsumer::LockedBuffer mBuffer;
-    bool mHaveBuffer;
-
-public:
-    ScreenshotClient();
-    ~ScreenshotClient();
-
-    // frees the previous screenshot and captures a new one
-    // if cropping isn't required, callers may pass in a default Rect, e.g.:
-    //   update(display, Rect(), useIdentityTransform);
-    status_t update(const sp<IBinder>& display,
-            Rect sourceCrop, bool useIdentityTransform);
-    status_t update(const sp<IBinder>& display,
-            Rect sourceCrop, uint32_t reqWidth, uint32_t reqHeight,
-            bool useIdentityTransform);
-    status_t update(const sp<IBinder>& display,
-            Rect sourceCrop, uint32_t reqWidth, uint32_t reqHeight,
-            int32_t minLayerZ, int32_t maxLayerZ,
-            bool useIdentityTransform);
-    status_t update(const sp<IBinder>& display,
-            Rect sourceCrop, uint32_t reqWidth, uint32_t reqHeight,
-            int32_t minLayerZ, int32_t maxLayerZ,
-            bool useIdentityTransform, uint32_t rotation);
-
-    sp<CpuConsumer> getCpuConsumer() const;
-
-    // release memory occupied by the screenshot
-    void release();
-
-    // pixels are valid until this object is freed or
-    // release() or update() is called
-    void const* getPixels() const;
-
-    uint32_t getWidth() const;
-    uint32_t getHeight() const;
-    PixelFormat getFormat() const;
-    uint32_t getStride() const;
-    // size of allocated memory in bytes
-    size_t getSize() const;
-    android_dataspace getDataSpace() const;
+    static status_t capture(const sp<IBinder>& display, Rect sourceCrop, uint32_t reqWidth,
+                            uint32_t reqHeight, int32_t minLayerZ, int32_t maxLayerZ,
+                            bool useIdentityTransform, uint32_t rotation,
+                            sp<GraphicBuffer>* outBuffer);
+    static status_t captureLayers(const sp<IBinder>& layerHandle, Rect sourceCrop, float fameScale,
+                                  sp<GraphicBuffer>* outBuffer);
 };
 
 // ---------------------------------------------------------------------------
diff --git a/libs/gui/tests/Surface_test.cpp b/libs/gui/tests/Surface_test.cpp
index ac4eb0c..470a338 100644
--- a/libs/gui/tests/Surface_test.cpp
+++ b/libs/gui/tests/Surface_test.cpp
@@ -116,14 +116,11 @@
     sp<ANativeWindow> anw(mSurface);
 
     // Verify the screenshot works with no protected buffers.
-    sp<IGraphicBufferProducer> producer;
-    sp<IGraphicBufferConsumer> consumer;
-    BufferQueue::createBufferQueue(&producer, &consumer);
-    sp<CpuConsumer> cpuConsumer = new CpuConsumer(consumer, 1);
     sp<ISurfaceComposer> sf(ComposerService::getComposerService());
     sp<IBinder> display(sf->getBuiltInDisplay(
             ISurfaceComposer::eDisplayIdMain));
-    ASSERT_EQ(NO_ERROR, sf->captureScreen(display, producer, Rect(),
+    sp<GraphicBuffer> outBuffer;
+    ASSERT_EQ(NO_ERROR, sf->captureScreen(display, &outBuffer, Rect(),
             64, 64, 0, 0x7fffffff, false));
 
     ASSERT_EQ(NO_ERROR, native_window_api_connect(anw.get(),
@@ -154,7 +151,7 @@
                 &buf));
         ASSERT_EQ(NO_ERROR, anw->queueBuffer(anw.get(), buf, -1));
     }
-    ASSERT_EQ(NO_ERROR, sf->captureScreen(display, producer, Rect(),
+    ASSERT_EQ(NO_ERROR, sf->captureScreen(display, &outBuffer, Rect(),
             64, 64, 0, 0x7fffffff, false));
 }
 
@@ -524,13 +521,13 @@
     status_t setActiveColorMode(const sp<IBinder>& /*display*/,
             android_color_mode_t /*colorMode*/) override { return NO_ERROR; }
     status_t captureScreen(const sp<IBinder>& /*display*/,
-            const sp<IGraphicBufferProducer>& /*producer*/,
+            sp<GraphicBuffer>* /*outBuffer*/,
             Rect /*sourceCrop*/, uint32_t /*reqWidth*/, uint32_t /*reqHeight*/,
             int32_t /*minLayerZ*/, int32_t /*maxLayerZ*/,
             bool /*useIdentityTransform*/,
             Rotation /*rotation*/) override { return NO_ERROR; }
     virtual status_t captureLayers(const sp<IBinder>& /*parentHandle*/,
-                                   const sp<IGraphicBufferProducer>& /*producer*/,
+                                   sp<GraphicBuffer>* /*outBuffer*/,
                                    const Rect& /*sourceCrop*/, float /*frameScale*/) override {
         return NO_ERROR;
     }
diff --git a/services/surfaceflinger/SurfaceFlinger.cpp b/services/surfaceflinger/SurfaceFlinger.cpp
index f29d980..59ac0ec 100644
--- a/services/surfaceflinger/SurfaceFlinger.cpp
+++ b/services/surfaceflinger/SurfaceFlinger.cpp
@@ -4343,44 +4343,10 @@
     const int mApi;
 };
 
-static status_t getWindowBuffer(ANativeWindow* window, uint32_t requestedWidth,
-                                uint32_t requestedHeight, bool hasWideColorDisplay,
-                                bool renderEngineUsesWideColor, ANativeWindowBuffer** outBuffer) {
-    const uint32_t usage = GRALLOC_USAGE_SW_READ_OFTEN | GRALLOC_USAGE_SW_WRITE_OFTEN |
-            GRALLOC_USAGE_HW_RENDER | GRALLOC_USAGE_HW_TEXTURE;
-
-    int err = 0;
-    err = native_window_set_buffers_dimensions(window, requestedWidth, requestedHeight);
-    err |= native_window_set_scaling_mode(window, NATIVE_WINDOW_SCALING_MODE_SCALE_TO_WINDOW);
-    err |= native_window_set_buffers_format(window, HAL_PIXEL_FORMAT_RGBA_8888);
-    err |= native_window_set_usage(window, usage);
-
-    if (hasWideColorDisplay) {
-        err |= native_window_set_buffers_data_space(window,
-                                                    renderEngineUsesWideColor
-                                                            ? HAL_DATASPACE_DISPLAY_P3
-                                                            : HAL_DATASPACE_V0_SRGB);
-    }
-
-    if (err != NO_ERROR) {
-        return BAD_VALUE;
-    }
-
-    /* TODO: Once we have the sync framework everywhere this can use
-     * server-side waits on the fence that dequeueBuffer returns.
-     */
-    err = native_window_dequeue_buffer_and_wait(window, outBuffer);
-    if (err != NO_ERROR) {
-        return err;
-    }
-
-    return NO_ERROR;
-}
-
-status_t SurfaceFlinger::captureScreen(const sp<IBinder>& display,
-                                       const sp<IGraphicBufferProducer>& producer, Rect sourceCrop,
-                                       uint32_t reqWidth, uint32_t reqHeight, int32_t minLayerZ,
-                                       int32_t maxLayerZ, bool useIdentityTransform,
+status_t SurfaceFlinger::captureScreen(const sp<IBinder>& display, sp<GraphicBuffer>* outBuffer,
+                                       Rect sourceCrop, uint32_t reqWidth, uint32_t reqHeight,
+                                       int32_t minLayerZ, int32_t maxLayerZ,
+                                       bool useIdentityTransform,
                                        ISurfaceComposer::Rotation rotation) {
     ATRACE_CALL();
 
@@ -4391,18 +4357,18 @@
 
     auto traverseLayers = std::bind(std::mem_fn(&SurfaceFlinger::traverseLayersInDisplay), this,
                                     device, minLayerZ, maxLayerZ, std::placeholders::_1);
-    return captureScreenCommon(renderArea, traverseLayers, producer, useIdentityTransform);
+    return captureScreenCommon(renderArea, traverseLayers, outBuffer, useIdentityTransform);
 }
 
 status_t SurfaceFlinger::captureLayers(const sp<IBinder>& layerHandleBinder,
-                                       const sp<IGraphicBufferProducer>& producer,
-                                       const Rect& sourceCrop, float frameScale) {
+                                       sp<GraphicBuffer>* outBuffer, const Rect& sourceCrop, 
+                                       float frameScale) {
     ATRACE_CALL();
 
     class LayerRenderArea : public RenderArea {
     public:
         LayerRenderArea(const sp<Layer>& layer, const Rect crop, int32_t reqWidth,
-            int32_t reqHeight)
+                        int32_t reqHeight)
               : RenderArea(reqHeight, reqWidth), mLayer(layer), mCrop(crop) {}
         const Transform& getTransform() const override {
             // Make the top level transform the inverse the transform and it's parent so it sets
@@ -4465,51 +4431,21 @@
             visitor(layer);
         });
     };
-    return captureScreenCommon(renderArea, traverseLayers, producer, false);
+    return captureScreenCommon(renderArea, traverseLayers, outBuffer, false);
 }
 
 status_t SurfaceFlinger::captureScreenCommon(RenderArea& renderArea,
                                              TraverseLayersFunction traverseLayers,
-                                             const sp<IGraphicBufferProducer>& producer,
+                                             sp<GraphicBuffer>* outBuffer,
                                              bool useIdentityTransform) {
     ATRACE_CALL();
 
-    if (CC_UNLIKELY(producer == 0))
-        return BAD_VALUE;
-
     renderArea.updateDimensions();
 
-    // if we have secure windows on this display, never allow the screen capture
-    // unless the producer interface is local (i.e.: we can take a screenshot for
-    // ourselves).
-    bool isLocalScreenshot = IInterface::asBinder(producer)->localBinder();
-
-    // create a surface (because we're a producer, and we need to
-    // dequeue/queue a buffer)
-    sp<Surface> surface = new Surface(producer, false);
-
-    // Put the screenshot Surface into async mode so that
-    // Layer::headFenceHasSignaled will always return true and we'll latch the
-    // first buffer regardless of whether or not its acquire fence has
-    // signaled. This is needed to avoid a race condition in the rotation
-    // animation. See b/30209608
-    surface->setAsyncMode(true);
-
-    ANativeWindow* window = surface.get();
-
-    status_t result = native_window_api_connect(window, NATIVE_WINDOW_API_EGL);
-    if (result != NO_ERROR) {
-        return result;
-    }
-    WindowDisconnector disconnector(window, NATIVE_WINDOW_API_EGL);
-
-    ANativeWindowBuffer* buffer = nullptr;
-    result = getWindowBuffer(window, renderArea.getReqWidth(), renderArea.getReqHeight(),
-                             hasWideColorDisplay && !mForceNativeColorMode,
-                             getRenderEngine().usesWideColor(), &buffer);
-    if (result != NO_ERROR) {
-        return result;
-    }
+    const uint32_t usage = GRALLOC_USAGE_SW_READ_OFTEN | GRALLOC_USAGE_SW_WRITE_OFTEN |
+            GRALLOC_USAGE_HW_RENDER | GRALLOC_USAGE_HW_TEXTURE;
+    *outBuffer = new GraphicBuffer(renderArea.getReqWidth(), renderArea.getReqHeight(),
+                                   HAL_PIXEL_FORMAT_RGBA_8888, 1, usage, "screenshot");
 
     // This mutex protects syncFd and captureResult for communication of the return values from the
     // main thread back to this Binder thread
@@ -4534,8 +4470,8 @@
         int fd = -1;
         {
             Mutex::Autolock _l(mStateLock);
-            result = captureScreenImplLocked(renderArea, traverseLayers, buffer,
-                                             useIdentityTransform, isLocalScreenshot, &fd);
+            result = captureScreenImplLocked(renderArea, traverseLayers, (*outBuffer).get(),
+                                             useIdentityTransform, &fd);
         }
 
         {
@@ -4546,7 +4482,7 @@
         }
     });
 
-    result = postMessageAsync(message);
+    status_t result = postMessageAsync(message);
     if (result == NO_ERROR) {
         captureCondition.wait(captureLock, [&]() { return captureResult; });
         while (*captureResult == EAGAIN) {
@@ -4561,9 +4497,10 @@
     }
 
     if (result == NO_ERROR) {
-        // queueBuffer takes ownership of syncFd
-        result = window->queueBuffer(window, buffer, syncFd);
+        sync_wait(syncFd, -1);
+        close(syncFd);
     }
+
     return result;
 }
 
@@ -4606,7 +4543,8 @@
     }
 
     engine.setWideColor(renderArea.getWideColorSupport() && !mForceNativeColorMode);
-    engine.setColorMode(mForceNativeColorMode ? HAL_COLOR_MODE_NATIVE : renderArea.getActiveColorMode());
+    engine.setColorMode(mForceNativeColorMode ? HAL_COLOR_MODE_NATIVE
+                                              : renderArea.getActiveColorMode());
 
     // make sure to clear all GL error flags
     engine.checkErrors();
@@ -4649,7 +4587,7 @@
 status_t SurfaceFlinger::captureScreenImplLocked(const RenderArea& renderArea,
                                                  TraverseLayersFunction traverseLayers,
                                                  ANativeWindowBuffer* buffer,
-                                                 bool useIdentityTransform, bool isLocalScreenshot,
+                                                 bool useIdentityTransform,
                                                  int* outSyncFd) {
     ATRACE_CALL();
 
@@ -4659,7 +4597,7 @@
         secureLayerIsVisible = secureLayerIsVisible || (layer->isVisible() && layer->isSecure());
     });
 
-    if (!isLocalScreenshot && secureLayerIsVisible) {
+    if (secureLayerIsVisible) {
         ALOGW("FB is protected: PERMISSION_DENIED");
         return PERMISSION_DENIED;
     }
@@ -4808,4 +4746,4 @@
 
 #if defined(__gl2_h_)
 #error "don't include gl2/gl2.h in this file"
-#endif
+#endif
\ No newline at end of file
diff --git a/services/surfaceflinger/SurfaceFlinger.h b/services/surfaceflinger/SurfaceFlinger.h
index f61dc75..a01d353 100644
--- a/services/surfaceflinger/SurfaceFlinger.h
+++ b/services/surfaceflinger/SurfaceFlinger.h
@@ -291,13 +291,11 @@
             std::vector<FrameEvent>* outSupported) const;
     virtual sp<IDisplayEventConnection> createDisplayEventConnection(
             ISurfaceComposer::VsyncSource vsyncSource = eVsyncSourceApp);
-    virtual status_t captureScreen(const sp<IBinder>& display,
-            const sp<IGraphicBufferProducer>& producer,
-            Rect sourceCrop, uint32_t reqWidth, uint32_t reqHeight,
-            int32_t minLayerZ, int32_t maxLayerZ,
-            bool useIdentityTransform, ISurfaceComposer::Rotation rotation);
-    virtual status_t captureLayers(const sp<IBinder>& parentHandle,
-                                   const sp<IGraphicBufferProducer>& producer,
+    virtual status_t captureScreen(const sp<IBinder>& display, sp<GraphicBuffer>* outBuffer,
+                                   Rect sourceCrop, uint32_t reqWidth, uint32_t reqHeight,
+                                   int32_t minLayerZ, int32_t maxLayerZ, bool useIdentityTransform,
+                                   ISurfaceComposer::Rotation rotation);
+    virtual status_t captureLayers(const sp<IBinder>& parentHandle, sp<GraphicBuffer>* outBuffer,
                                    const Rect& sourceCrop, float frameScale);
     virtual status_t getDisplayStats(const sp<IBinder>& display,
             DisplayStatInfo* stats);
@@ -436,15 +434,13 @@
 
     void renderScreenImplLocked(const RenderArea& renderArea, TraverseLayersFunction traverseLayers,
                                 bool yswap, bool useIdentityTransform);
-
     status_t captureScreenCommon(RenderArea& renderArea, TraverseLayersFunction traverseLayers,
-                                 const sp<IGraphicBufferProducer>& producer,
+                                 sp<GraphicBuffer>* outBuffer,
                                  bool useIdentityTransform);
-
     status_t captureScreenImplLocked(const RenderArea& renderArea,
                                      TraverseLayersFunction traverseLayers,
                                      ANativeWindowBuffer* buffer, bool useIdentityTransform,
-                                     bool isLocalScreenshot, int* outSyncFd);
+                                     int* outSyncFd);
     void traverseLayersInDisplay(const sp<const DisplayDevice>& display, int32_t minLayerZ,
                                  int32_t maxLayerZ, const LayerVector::Visitor& visitor);
 
diff --git a/services/surfaceflinger/tests/Transaction_test.cpp b/services/surfaceflinger/tests/Transaction_test.cpp
index 39aef81..ff81dc9 100644
--- a/services/surfaceflinger/tests/Transaction_test.cpp
+++ b/services/surfaceflinger/tests/Transaction_test.cpp
@@ -104,20 +104,22 @@
 }
 
 // Check if a region has the specified color.
-void expectBufferColor(const CpuConsumer::LockedBuffer& buffer, const Rect& rect,
+void expectBufferColor(const sp<GraphicBuffer>& outBuffer, uint8_t* pixels, const Rect& rect,
                        const Color& color, uint8_t tolerance) {
     int32_t x = rect.left;
     int32_t y = rect.top;
     int32_t width = rect.right - rect.left;
     int32_t height = rect.bottom - rect.top;
 
-    if (x + width > int32_t(buffer.width)) {
-        x = std::min(x, int32_t(buffer.width));
-        width = buffer.width - x;
+    int32_t bufferWidth = int32_t(outBuffer->getWidth());
+    int32_t bufferHeight = int32_t(outBuffer->getHeight());
+    if (x + width > bufferWidth) {
+        x = std::min(x, bufferWidth);
+        width = bufferWidth - x;
     }
-    if (y + height > int32_t(buffer.height)) {
-        y = std::min(y, int32_t(buffer.height));
-        height = buffer.height - y;
+    if (y + height > bufferHeight) {
+        y = std::min(y, bufferHeight);
+        height = bufferHeight - y;
     }
 
     auto colorCompare = [tolerance](uint8_t a, uint8_t b) {
@@ -125,8 +127,7 @@
         return tmp <= tolerance;
     };
     for (int32_t j = 0; j < height; j++) {
-        const uint8_t* src =
-                static_cast<const uint8_t*>(buffer.data) + (buffer.stride * (y + j) + x) * 4;
+        const uint8_t* src = pixels + (outBuffer->getStride() * (y + j) + x) * 4;
         for (int32_t i = 0; i < width; i++) {
             const uint8_t expected[4] = {color.r, color.g, color.b, color.a};
             EXPECT_TRUE(std::equal(src, src + 4, expected, colorCompare))
@@ -170,30 +171,38 @@
 public:
     static void captureScreen(sp<ScreenCapture>* sc, int32_t minLayerZ = 0,
                               int32_t maxLayerZ = std::numeric_limits<int32_t>::max()) {
-        sp<IGraphicBufferProducer> producer;
-        sp<IGraphicBufferConsumer> consumer;
-        BufferQueue::createBufferQueue(&producer, &consumer);
-        sp<CpuConsumer> cpuConsumer = new CpuConsumer(consumer, 1);
         sp<ISurfaceComposer> sf(ComposerService::getComposerService());
         sp<IBinder> display(sf->getBuiltInDisplay(ISurfaceComposer::eDisplayIdMain));
         SurfaceComposerClient::Transaction().apply(true);
 
+        sp<GraphicBuffer> outBuffer;
         ASSERT_EQ(NO_ERROR,
-                  sf->captureScreen(display, producer, Rect(), 0, 0, minLayerZ, maxLayerZ, false));
-        *sc = new ScreenCapture(cpuConsumer);
+                  sf->captureScreen(display, &outBuffer, Rect(), 0, 0, minLayerZ, maxLayerZ,
+                                    false));
+        *sc = new ScreenCapture(outBuffer);
+    }
+
+    static void captureLayers(std::unique_ptr<ScreenCapture>* sc, sp<IBinder>& parentHandle,
+                              Rect crop = Rect::EMPTY_RECT, float frameScale = 1.0) {
+        sp<ISurfaceComposer> sf(ComposerService::getComposerService());
+        SurfaceComposerClient::Transaction().apply(true);
+
+        sp<GraphicBuffer> outBuffer;
+        ASSERT_EQ(NO_ERROR, sf->captureLayers(parentHandle, &outBuffer, crop, frameScale));
+        *sc = std::make_unique<ScreenCapture>(outBuffer);
     }
 
     void expectColor(const Rect& rect, const Color& color, uint8_t tolerance = 0) {
-        ASSERT_EQ(HAL_PIXEL_FORMAT_RGBA_8888, mBuf.format);
-        expectBufferColor(mBuf, rect, color, tolerance);
+        ASSERT_EQ(HAL_PIXEL_FORMAT_RGBA_8888, mOutBuffer->getPixelFormat());
+        expectBufferColor(mOutBuffer, mPixels, rect, color, tolerance);
     }
 
     void expectBorder(const Rect& rect, const Color& color, uint8_t tolerance = 0) {
-        ASSERT_EQ(HAL_PIXEL_FORMAT_RGBA_8888, mBuf.format);
+        ASSERT_EQ(HAL_PIXEL_FORMAT_RGBA_8888, mOutBuffer->getPixelFormat());
         const bool leftBorder = rect.left > 0;
         const bool topBorder = rect.top > 0;
-        const bool rightBorder = rect.right < int32_t(mBuf.width);
-        const bool bottomBorder = rect.bottom < int32_t(mBuf.height);
+        const bool rightBorder = rect.right < int32_t(mOutBuffer->getWidth());
+        const bool bottomBorder = rect.bottom < int32_t(mOutBuffer->getHeight());
 
         if (topBorder) {
             Rect top(rect.left, rect.top - 1, rect.right, rect.top);
@@ -246,9 +255,8 @@
     }
 
     void checkPixel(uint32_t x, uint32_t y, uint8_t r, uint8_t g, uint8_t b) {
-        ASSERT_EQ(HAL_PIXEL_FORMAT_RGBA_8888, mBuf.format);
-        const uint8_t* img = static_cast<const uint8_t*>(mBuf.data);
-        const uint8_t* pixel = img + (4 * (y * mBuf.stride + x));
+        ASSERT_EQ(HAL_PIXEL_FORMAT_RGBA_8888, mOutBuffer->getPixelFormat());
+        const uint8_t* pixel = mPixels + (4 * (y * mOutBuffer->getStride() + x));
         if (r != pixel[0] || g != pixel[1] || b != pixel[2]) {
             String8 err(String8::format("pixel @ (%3d, %3d): "
                                         "expected [%3d, %3d, %3d], got [%3d, %3d, %3d]",
@@ -263,58 +271,15 @@
 
     void expectChildColor(uint32_t x, uint32_t y) { checkPixel(x, y, 200, 200, 200); }
 
-private:
-    ScreenCapture(const sp<CpuConsumer>& cc) : mCC(cc) {
-        EXPECT_EQ(NO_ERROR, mCC->lockNextBuffer(&mBuf));
+    ScreenCapture(const sp<GraphicBuffer>& outBuffer) : mOutBuffer(outBuffer) {
+        mOutBuffer->lock(GRALLOC_USAGE_SW_READ_OFTEN, reinterpret_cast<void**>(&mPixels));
     }
 
-    ~ScreenCapture() { mCC->unlockBuffer(mBuf); }
-
-    sp<CpuConsumer> mCC;
-    CpuConsumer::LockedBuffer mBuf;
-};
-
-class CaptureLayer {
-public:
-    static void captureScreen(std::unique_ptr<CaptureLayer>* sc, sp<IBinder>& parentHandle,
-                              Rect crop = Rect::EMPTY_RECT, float frameScale = 1.0) {
-        sp<IGraphicBufferProducer> producer;
-        sp<IGraphicBufferConsumer> consumer;
-        BufferQueue::createBufferQueue(&producer, &consumer);
-        sp<CpuConsumer> cpuConsumer = new CpuConsumer(consumer, 1);
-        sp<ISurfaceComposer> sf(ComposerService::getComposerService());
-        SurfaceComposerClient::Transaction().apply(true);
-        ASSERT_EQ(NO_ERROR, sf->captureLayers(parentHandle, producer, crop, frameScale));
-        *sc = std::make_unique<CaptureLayer>(cpuConsumer);
-    }
-
-    void checkPixel(uint32_t x, uint32_t y, uint8_t r, uint8_t g, uint8_t b) {
-        ASSERT_EQ(HAL_PIXEL_FORMAT_RGBA_8888, mBuffer.format);
-        const uint8_t* img = static_cast<const uint8_t*>(mBuffer.data);
-        const uint8_t* pixel = img + (4 * (y * mBuffer.stride + x));
-        if (r != pixel[0] || g != pixel[1] || b != pixel[2]) {
-            String8 err(String8::format("pixel @ (%3d, %3d): "
-                                        "expected [%3d, %3d, %3d], got [%3d, %3d, %3d]",
-                                        x, y, r, g, b, pixel[0], pixel[1], pixel[2]));
-            EXPECT_EQ(String8(), err) << err.string();
-        }
-    }
-
-    void expectFGColor(uint32_t x, uint32_t y) { checkPixel(x, y, 195, 63, 63); }
-
-    void expectBGColor(uint32_t x, uint32_t y) { checkPixel(x, y, 63, 63, 195); }
-
-    void expectChildColor(uint32_t x, uint32_t y) { checkPixel(x, y, 200, 200, 200); }
-
-    CaptureLayer(const sp<CpuConsumer>& cc) : mCC(cc) {
-        EXPECT_EQ(NO_ERROR, mCC->lockNextBuffer(&mBuffer));
-    }
-
-    ~CaptureLayer() { mCC->unlockBuffer(mBuffer); }
+    ~ScreenCapture() { mOutBuffer->unlock(); }
 
 private:
-    sp<CpuConsumer> mCC;
-    CpuConsumer::LockedBuffer mBuffer;
+    sp<GraphicBuffer> mOutBuffer;
+    uint8_t* mPixels = NULL;
 };
 
 class LayerTransactionTest : public ::testing::Test {
@@ -858,21 +823,17 @@
     ASSERT_NO_FATAL_FAILURE(fillLayerColor(layer, Color::RED));
 
     sp<ISurfaceComposer> composer = ComposerService::getComposerService();
-    sp<IGraphicBufferProducer> producer;
-    sp<IGraphicBufferConsumer> consumer;
-    BufferQueue::createBufferQueue(&producer, &consumer);
-    sp<CpuConsumer> cpuConsumer = new CpuConsumer(consumer, 1);
-
+    sp<GraphicBuffer> outBuffer;
     Transaction()
             .setFlags(layer, layer_state_t::eLayerSecure, layer_state_t::eLayerSecure)
             .apply(true);
     ASSERT_EQ(PERMISSION_DENIED,
-              composer->captureScreen(mDisplay, producer, Rect(), 0, 0, mLayerZBase, mLayerZBase,
+              composer->captureScreen(mDisplay, &outBuffer, Rect(), 0, 0, mLayerZBase, mLayerZBase,
                                       false));
 
     Transaction().setFlags(layer, 0, layer_state_t::eLayerSecure).apply(true);
     ASSERT_EQ(NO_ERROR,
-              composer->captureScreen(mDisplay, producer, Rect(), 0, 0, mLayerZBase, mLayerZBase,
+              composer->captureScreen(mDisplay, &outBuffer, Rect(), 0, 0, mLayerZBase, mLayerZBase,
                                       false));
 }
 
@@ -1504,7 +1465,7 @@
     }
 }
 
-class LayerUpdateTest : public ::testing::Test {
+class LayerUpdateTest : public LayerTransactionTest {
 protected:
     virtual void SetUp() {
         mComposerClient = new SurfaceComposerClient;
@@ -2318,12 +2279,12 @@
 
 class ScreenCaptureTest : public LayerUpdateTest {
 protected:
-    std::unique_ptr<CaptureLayer> mCapture;
+    std::unique_ptr<ScreenCapture> mCapture;
 };
 
 TEST_F(ScreenCaptureTest, CaptureSingleLayer) {
     auto bgHandle = mBGSurfaceControl->getHandle();
-    CaptureLayer::captureScreen(&mCapture, bgHandle);
+    ScreenCapture::captureLayers(&mCapture, bgHandle);
     mCapture->expectBGColor(0, 0);
     // Doesn't capture FG layer which is at 64, 64
     mCapture->expectBGColor(64, 64);
@@ -2340,7 +2301,7 @@
     SurfaceComposerClient::Transaction().show(child).apply(true);
 
     // Captures mFGSurfaceControl layer and its child.
-    CaptureLayer::captureScreen(&mCapture, fgHandle);
+    ScreenCapture::captureLayers(&mCapture, fgHandle);
     mCapture->expectFGColor(10, 10);
     mCapture->expectChildColor(0, 0);
 }
@@ -2365,7 +2326,7 @@
             .apply(true);
 
     // Captures mFGSurfaceControl, its child, and the grandchild.
-    CaptureLayer::captureScreen(&mCapture, fgHandle);
+    ScreenCapture::captureLayers(&mCapture, fgHandle);
     mCapture->expectFGColor(10, 10);
     mCapture->expectChildColor(0, 0);
     mCapture->checkPixel(5, 5, 50, 50, 50);
@@ -2381,7 +2342,7 @@
     SurfaceComposerClient::Transaction().setPosition(child, 5, 5).show(child).apply(true);
 
     // Captures only the child layer, and not the parent.
-    CaptureLayer::captureScreen(&mCapture, childHandle);
+    ScreenCapture::captureLayers(&mCapture, childHandle);
     mCapture->expectChildColor(0, 0);
     mCapture->expectChildColor(9, 9);
 }
@@ -2407,94 +2368,96 @@
     auto grandchildHandle = grandchild->getHandle();
 
     // Captures only the grandchild.
-    CaptureLayer::captureScreen(&mCapture, grandchildHandle);
+    ScreenCapture::captureLayers(&mCapture, grandchildHandle);
     mCapture->checkPixel(0, 0, 50, 50, 50);
     mCapture->checkPixel(4, 4, 50, 50, 50);
 }
 
 TEST_F(ScreenCaptureTest, CaptureCrop) {
-    sp<SurfaceControl> redLayer = mComposerClient->createSurface(
-        String8("Red surface"),
-        60, 60, PIXEL_FORMAT_RGBA_8888, 0);
-    sp<SurfaceControl> blueLayer = mComposerClient->createSurface(
-        String8("Blue surface"),
-        30, 30, PIXEL_FORMAT_RGBA_8888, 0, redLayer.get());
+    sp<SurfaceControl> redLayer = mComposerClient->createSurface(String8("Red surface"), 60, 60,
+                                                                 PIXEL_FORMAT_RGBA_8888, 0);
+    sp<SurfaceControl> blueLayer =
+            mComposerClient->createSurface(String8("Blue surface"), 30, 30, PIXEL_FORMAT_RGBA_8888,
+                                           0, redLayer.get());
 
-    fillSurfaceRGBA8(redLayer, 255, 0, 0);
-    fillSurfaceRGBA8(blueLayer, 0, 0, 255);
+    ASSERT_NO_FATAL_FAILURE(fillLayerColor(redLayer, Color::RED));
+    ASSERT_NO_FATAL_FAILURE(fillLayerColor(blueLayer, Color::BLUE));
 
     SurfaceComposerClient::Transaction()
-        .setLayer(redLayer, INT32_MAX-1)
-        .show(redLayer)
-        .show(blueLayer)
-        .apply(true);
+            .setLayer(redLayer, INT32_MAX - 1)
+            .show(redLayer)
+            .show(blueLayer)
+            .apply(true);
 
     auto redLayerHandle = redLayer->getHandle();
 
     // Capturing full screen should have both red and blue are visible.
-    CaptureLayer::captureScreen(&mCapture, redLayerHandle);
-    mCapture->checkPixel(29, 29, 0, 0, 255);
-    mCapture->checkPixel(30, 30, 255, 0, 0);
+    ScreenCapture::captureLayers(&mCapture, redLayerHandle);
+    mCapture->expectColor(Rect(0, 0, 29, 29), Color::BLUE);
+    // red area below the blue area
+    mCapture->expectColor(Rect(0, 30, 59, 59), Color::RED);
+    // red area to the right of the blue area
+    mCapture->expectColor(Rect(30, 0, 59, 59), Color::RED);
 
     Rect crop = Rect(0, 0, 30, 30);
-    CaptureLayer::captureScreen(&mCapture, redLayerHandle, crop);
+    ScreenCapture::captureLayers(&mCapture, redLayerHandle, crop);
     // Capturing the cropped screen, cropping out the shown red area, should leave only the blue
     // area visible.
-    mCapture->checkPixel(29, 29, 0, 0, 255);
+    mCapture->expectColor(Rect(0, 0, 29, 29), Color::BLUE);
     mCapture->checkPixel(30, 30, 0, 0, 0);
 }
 
 TEST_F(ScreenCaptureTest, CaptureSize) {
-    sp<SurfaceControl> redLayer = mComposerClient->createSurface(
-        String8("Red surface"),
-        60, 60, PIXEL_FORMAT_RGBA_8888, 0);
-    sp<SurfaceControl> blueLayer = mComposerClient->createSurface(
-        String8("Blue surface"),
-        30, 30, PIXEL_FORMAT_RGBA_8888, 0, redLayer.get());
+    sp<SurfaceControl> redLayer = mComposerClient->createSurface(String8("Red surface"), 60, 60,
+                                                                 PIXEL_FORMAT_RGBA_8888, 0);
+    sp<SurfaceControl> blueLayer =
+            mComposerClient->createSurface(String8("Blue surface"), 30, 30, PIXEL_FORMAT_RGBA_8888,
+                                           0, redLayer.get());
 
-    fillSurfaceRGBA8(redLayer, 255, 0, 0);
-    fillSurfaceRGBA8(blueLayer, 0, 0, 255);
+    ASSERT_NO_FATAL_FAILURE(fillLayerColor(redLayer, Color::RED));
+    ASSERT_NO_FATAL_FAILURE(fillLayerColor(blueLayer, Color::BLUE));
 
     SurfaceComposerClient::Transaction()
-        .setLayer(redLayer, INT32_MAX-1)
-        .show(redLayer)
-        .show(blueLayer)
-        .apply(true);
+            .setLayer(redLayer, INT32_MAX - 1)
+            .show(redLayer)
+            .show(blueLayer)
+            .apply(true);
 
     auto redLayerHandle = redLayer->getHandle();
 
     // Capturing full screen should have both red and blue are visible.
-    CaptureLayer::captureScreen(&mCapture, redLayerHandle);
-    mCapture->checkPixel(29, 29, 0, 0, 255);
-    mCapture->checkPixel(30, 30, 255, 0, 0);
+    ScreenCapture::captureLayers(&mCapture, redLayerHandle);
+    mCapture->expectColor(Rect(0, 0, 29, 29), Color::BLUE);
+    // red area below the blue area
+    mCapture->expectColor(Rect(0, 30, 59, 59), Color::RED);
+    // red area to the right of the blue area
+    mCapture->expectColor(Rect(30, 0, 59, 59), Color::RED);
 
-    CaptureLayer::captureScreen(&mCapture, redLayerHandle, Rect::EMPTY_RECT, 0.5);
+    ScreenCapture::captureLayers(&mCapture, redLayerHandle, Rect::EMPTY_RECT, 0.5);
     // Capturing the downsized area (30x30) should leave both red and blue but in a smaller area.
-    mCapture->checkPixel(14, 14, 0, 0, 255);
-    mCapture->checkPixel(15, 15, 255, 0, 0);
-    mCapture->checkPixel(29, 29, 255, 0, 0);
+    mCapture->expectColor(Rect(0, 0, 14, 14), Color::BLUE);
+    // red area below the blue area
+    mCapture->expectColor(Rect(0, 15, 29, 29), Color::RED);
+    // red area to the right of the blue area
+    mCapture->expectColor(Rect(15, 0, 29, 29), Color::RED);
     mCapture->checkPixel(30, 30, 0, 0, 0);
 }
 
 TEST_F(ScreenCaptureTest, CaptureInvalidLayer) {
-    sp<SurfaceControl> redLayer = mComposerClient->createSurface(
-        String8("Red surface"),
-        60, 60, PIXEL_FORMAT_RGBA_8888, 0);
+    sp<SurfaceControl> redLayer = mComposerClient->createSurface(String8("Red surface"), 60, 60,
+                                                                 PIXEL_FORMAT_RGBA_8888, 0);
 
-    fillSurfaceRGBA8(redLayer, 255, 0, 0);
+    ASSERT_NO_FATAL_FAILURE(fillLayerColor(redLayer, Color::RED));
 
     auto redLayerHandle = redLayer->getHandle();
     mComposerClient->destroySurface(redLayerHandle);
     SurfaceComposerClient::Transaction().apply(true);
 
-    sp<IGraphicBufferProducer> producer;
-    sp<IGraphicBufferConsumer> consumer;
-    BufferQueue::createBufferQueue(&producer, &consumer);
-    sp<CpuConsumer> cpuConsumer = new CpuConsumer(consumer, 1);
-    sp<ISurfaceComposer> sf(ComposerService::getComposerService());
+    sp<GraphicBuffer> outBuffer;
 
     // Layer was deleted so captureLayers should fail with NAME_NOT_FOUND
-    ASSERT_EQ(NAME_NOT_FOUND, sf->captureLayers(redLayerHandle, producer, Rect::EMPTY_RECT, 1.0));
+    sp<ISurfaceComposer> sf(ComposerService::getComposerService());
+    ASSERT_EQ(NAME_NOT_FOUND, sf->captureLayers(redLayerHandle, &outBuffer, Rect::EMPTY_RECT, 1.0));
 }
 
-}
\ No newline at end of file
+} // namespace android