diff --git a/services/surfaceflinger/CompositionEngine/src/Output.cpp b/services/surfaceflinger/CompositionEngine/src/Output.cpp
index 248933e..7929beb 100644
--- a/services/surfaceflinger/CompositionEngine/src/Output.cpp
+++ b/services/surfaceflinger/CompositionEngine/src/Output.cpp
@@ -913,9 +913,8 @@
 
     const nsecs_t renderEngineStart = systemTime();
     status_t status =
-            renderEngine.drawLayers(clientCompositionDisplay, clientCompositionLayerPointers,
-                                    buf->getNativeBuffer(), /*useFramebufferCache=*/true,
-                                    std::move(fd), &readyFence);
+            renderEngine.drawLayers(clientCompositionDisplay, clientCompositionLayerPointers, buf,
+                                    /*useFramebufferCache=*/true, std::move(fd), &readyFence);
 
     if (status != NO_ERROR && mClientCompositionRequestCache) {
         // If rendering was not successful, remove the request from the cache.
diff --git a/services/surfaceflinger/SurfaceFlinger.cpp b/services/surfaceflinger/SurfaceFlinger.cpp
index ef59b55..8f66928 100644
--- a/services/surfaceflinger/SurfaceFlinger.cpp
+++ b/services/surfaceflinger/SurfaceFlinger.cpp
@@ -5558,8 +5558,8 @@
 
 void SurfaceFlinger::renderScreenImplLocked(const RenderArea& renderArea,
                                             TraverseLayersFunction traverseLayers,
-                                            ANativeWindowBuffer* buffer, bool useIdentityTransform,
-                                            int* outSyncFd) {
+                                            const sp<GraphicBuffer>& buffer,
+                                            bool useIdentityTransform, int* outSyncFd) {
     ATRACE_CALL();
 
     const auto reqWidth = renderArea.getReqWidth();
@@ -5684,7 +5684,7 @@
 
 status_t SurfaceFlinger::captureScreenImplLocked(const RenderArea& renderArea,
                                                  TraverseLayersFunction traverseLayers,
-                                                 ANativeWindowBuffer* buffer,
+                                                 const sp<GraphicBuffer>& buffer,
                                                  bool useIdentityTransform, bool forSystem,
                                                  int* outSyncFd, bool& outCapturedSecureLayers) {
     ATRACE_CALL();
diff --git a/services/surfaceflinger/SurfaceFlinger.h b/services/surfaceflinger/SurfaceFlinger.h
index 09a3430..5f76260 100644
--- a/services/surfaceflinger/SurfaceFlinger.h
+++ b/services/surfaceflinger/SurfaceFlinger.h
@@ -692,7 +692,7 @@
     using TraverseLayersFunction = std::function<void(const LayerVector::Visitor&)>;
 
     void renderScreenImplLocked(const RenderArea& renderArea, TraverseLayersFunction traverseLayers,
-                                ANativeWindowBuffer* buffer, bool useIdentityTransform,
+                                const sp<GraphicBuffer>& buffer, bool useIdentityTransform,
                                 int* outSyncFd);
     status_t captureScreenCommon(RenderArea& renderArea, TraverseLayersFunction traverseLayers,
                                  sp<GraphicBuffer>* outBuffer, const ui::PixelFormat reqPixelFormat,
@@ -704,7 +704,7 @@
     const sp<DisplayDevice> getDisplayByLayerStack(uint64_t layerStack);
     status_t captureScreenImplLocked(const RenderArea& renderArea,
                                      TraverseLayersFunction traverseLayers,
-                                     ANativeWindowBuffer* buffer, bool useIdentityTransform,
+                                     const sp<GraphicBuffer>& buffer, bool useIdentityTransform,
                                      bool forSystem, int* outSyncFd, bool& outCapturedSecureLayers);
     void traverseLayersInDisplay(const sp<const DisplayDevice>& display,
                                  const LayerVector::Visitor& visitor);
diff --git a/services/surfaceflinger/tests/unittests/CompositionTest.cpp b/services/surfaceflinger/tests/unittests/CompositionTest.cpp
index 06ef8e7..867bb82 100644
--- a/services/surfaceflinger/tests/unittests/CompositionTest.cpp
+++ b/services/surfaceflinger/tests/unittests/CompositionTest.cpp
@@ -328,7 +328,7 @@
                 .WillRepeatedly(
                         [](const renderengine::DisplaySettings& displaySettings,
                            const std::vector<const renderengine::LayerSettings*>&,
-                           ANativeWindowBuffer*, const bool, base::unique_fd&&,
+                           const sp<GraphicBuffer>&, const bool, base::unique_fd&&,
                            base::unique_fd*) -> status_t {
                             EXPECT_EQ(DEFAULT_DISPLAY_MAX_LUMINANCE, displaySettings.maxLuminance);
                             EXPECT_EQ(Rect(DEFAULT_DISPLAY_WIDTH, DEFAULT_DISPLAY_HEIGHT),
@@ -378,7 +378,7 @@
                 .WillRepeatedly(
                         [](const renderengine::DisplaySettings& displaySettings,
                            const std::vector<const renderengine::LayerSettings*>&,
-                           ANativeWindowBuffer*, const bool, base::unique_fd&&,
+                           const sp<GraphicBuffer>&, const bool, base::unique_fd&&,
                            base::unique_fd*) -> status_t {
                             EXPECT_EQ(DEFAULT_DISPLAY_MAX_LUMINANCE, displaySettings.maxLuminance);
                             EXPECT_EQ(Rect(DEFAULT_DISPLAY_WIDTH, DEFAULT_DISPLAY_HEIGHT),
@@ -629,7 +629,7 @@
         EXPECT_CALL(*test->mRenderEngine, drawLayers)
                 .WillOnce([](const renderengine::DisplaySettings& displaySettings,
                              const std::vector<const renderengine::LayerSettings*>& layerSettings,
-                             ANativeWindowBuffer*, const bool, base::unique_fd&&,
+                             const sp<GraphicBuffer>&, const bool, base::unique_fd&&,
                              base::unique_fd*) -> status_t {
                     EXPECT_EQ(DEFAULT_DISPLAY_MAX_LUMINANCE, displaySettings.maxLuminance);
                     EXPECT_EQ(Rect(DEFAULT_DISPLAY_WIDTH, DEFAULT_DISPLAY_HEIGHT),
@@ -678,7 +678,7 @@
         EXPECT_CALL(*test->mRenderEngine, drawLayers)
                 .WillOnce([](const renderengine::DisplaySettings& displaySettings,
                              const std::vector<const renderengine::LayerSettings*>& layerSettings,
-                             ANativeWindowBuffer*, const bool, base::unique_fd&&,
+                             const sp<GraphicBuffer>&, const bool, base::unique_fd&&,
                              base::unique_fd*) -> status_t {
                     EXPECT_EQ(DEFAULT_DISPLAY_MAX_LUMINANCE, displaySettings.maxLuminance);
                     EXPECT_EQ(Rect(DEFAULT_DISPLAY_WIDTH, DEFAULT_DISPLAY_HEIGHT),
@@ -756,7 +756,7 @@
         EXPECT_CALL(*test->mRenderEngine, drawLayers)
                 .WillOnce([](const renderengine::DisplaySettings& displaySettings,
                              const std::vector<const renderengine::LayerSettings*>& layerSettings,
-                             ANativeWindowBuffer*, const bool, base::unique_fd&&,
+                             const sp<GraphicBuffer>&, const bool, base::unique_fd&&,
                              base::unique_fd*) -> status_t {
                     EXPECT_EQ(DEFAULT_DISPLAY_MAX_LUMINANCE, displaySettings.maxLuminance);
                     EXPECT_EQ(Rect(DEFAULT_DISPLAY_WIDTH, DEFAULT_DISPLAY_HEIGHT),
diff --git a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
index 3a4f349..557feaf 100644
--- a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
+++ b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
@@ -314,9 +314,10 @@
 
     auto onMessageReceived(int32_t what) { return mFlinger->onMessageReceived(what); }
 
-    auto captureScreenImplLocked(
-            const RenderArea& renderArea, SurfaceFlinger::TraverseLayersFunction traverseLayers,
-            ANativeWindowBuffer* buffer, bool useIdentityTransform, bool forSystem, int* outSyncFd) {
+    auto captureScreenImplLocked(const RenderArea& renderArea,
+                                 SurfaceFlinger::TraverseLayersFunction traverseLayers,
+                                 const sp<GraphicBuffer>& buffer, bool useIdentityTransform,
+                                 bool forSystem, int* outSyncFd) {
         bool ignored;
         return mFlinger->captureScreenImplLocked(renderArea, traverseLayers, buffer,
                                                  useIdentityTransform, forSystem, outSyncFd,