From 649ca6d383e1f5fa80cb3805a5430e6707e23f3c Mon Sep 17 00:00:00 2001
From: Mark Renouf <mrenouf@google.com>
Date: Wed, 27 Apr 2022 13:31:58 -0400
Subject: [PATCH] Fix an NPE on cancellation when connection is closed

Store the future locally so the callback doesn't lose
a reference to it. Even when the field is cleared, the
callback will always receive the original value, avoiding
NPEs. This change also adds a test for cancellation instead
of falling through to an ignored CancellationException.

Bug: 219965002
Test: manual; timing dependent, dismiss screenshot while
      scroll capture request is in flight
Change-Id: I88918bbe9379208cad9456b53ed79d600c73a60b
---
 .../systemui/screenshot/ScreenshotController.java    | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
index 009d4b9b48e60..4728c678f96c6 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
@@ -674,15 +674,21 @@ public class ScreenshotController {
         if (mLastScrollCaptureRequest != null) {
             mLastScrollCaptureRequest.cancel(true);
         }
-        mLastScrollCaptureRequest = mScrollCaptureClient.request(DEFAULT_DISPLAY);
+        final ListenableFuture<ScrollCaptureResponse> future =
+                mScrollCaptureClient.request(DEFAULT_DISPLAY);
+        mLastScrollCaptureRequest = future;
         mLastScrollCaptureRequest.addListener(() ->
-                onScrollCaptureResponseReady(mLastScrollCaptureRequest), mMainExecutor);
+                onScrollCaptureResponseReady(future), mMainExecutor);
     }
 
     private void onScrollCaptureResponseReady(Future<ScrollCaptureResponse> responseFuture) {
         try {
             if (mLastScrollCaptureResponse != null) {
                 mLastScrollCaptureResponse.close();
+                mLastScrollCaptureResponse = null;
+            }
+            if (responseFuture.isCancelled()) {
+                return;
             }
             mLastScrollCaptureResponse = responseFuture.get();
             if (!mLastScrollCaptureResponse.isConnected()) {
@@ -707,8 +713,6 @@ public class ScreenshotController {
                 // delay starting scroll capture to make sure the scrim is up before the app moves
                 mScreenshotView.post(() -> runBatchScrollCapture(response));
             });
-        } catch (CancellationException e) {
-            // Ignore
         } catch (InterruptedException | ExecutionException e) {
             Log.e(TAG, "requestScrollCapture failed", e);
         }