1 files changed, 31 insertions, 9 deletions
diff --git a/core/java/com/android/internal/util/ScreenshotHelper.java b/core/java/com/android/internal/util/ScreenshotHelper.java
index 9bf05135c4c5..9ee7c3005b0e 100644
--- a/core/java/com/android/internal/util/ScreenshotHelper.java
+++ b/core/java/com/android/internal/util/ScreenshotHelper.java
@@ -1,12 +1,15 @@
 package com.android.internal.util;
 
+import static android.content.Intent.ACTION_USER_SWITCHED;
 import static android.view.WindowManager.ScreenshotSource.SCREENSHOT_OTHER;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.content.BroadcastReceiver;
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
+import android.content.IntentFilter;
 import android.content.ServiceConnection;
 import android.graphics.Insets;
 import android.graphics.Rect;
@@ -161,8 +164,21 @@ public class ScreenshotHelper {
     private ServiceConnection mScreenshotConnection = null;
     private final Context mContext;
 
+    private final BroadcastReceiver mBroadcastReceiver = new BroadcastReceiver() {
+        @Override
+        public void onReceive(Context context, Intent intent) {
+            synchronized (mScreenshotLock) {
+                if (ACTION_USER_SWITCHED.equals(intent.getAction())) {
+                    resetConnection();
+                }
+            }
+        }
+    };
+
     public ScreenshotHelper(Context context) {
         mContext = context;
+        IntentFilter filter = new IntentFilter(ACTION_USER_SWITCHED);
+        mContext.registerReceiver(mBroadcastReceiver, filter);
     }
 
     /**
@@ -279,9 +295,8 @@ public class ScreenshotHelper {
             final Runnable mScreenshotTimeout = () -> {
                 synchronized (mScreenshotLock) {
                     if (mScreenshotConnection != null) {
-                        mContext.unbindService(mScreenshotConnection);
-                        mScreenshotConnection = null;
-                        mScreenshotService = null;
+                        Log.e(TAG, "Timed out before getting screenshot capture response");
+                        resetConnection();
                         notifyScreenshotError();
                     }
                 }
@@ -305,9 +320,7 @@ public class ScreenshotHelper {
                         case SCREENSHOT_MSG_PROCESS_COMPLETE:
                             synchronized (mScreenshotLock) {
                                 if (myConn != null && mScreenshotConnection == myConn) {
-                                    mContext.unbindService(myConn);
-                                    mScreenshotConnection = null;
-                                    mScreenshotService = null;
+                                    resetConnection();
                                 }
                             }
                             break;
@@ -348,9 +361,7 @@ public class ScreenshotHelper {
                     public void onServiceDisconnected(ComponentName name) {
                         synchronized (mScreenshotLock) {
                             if (mScreenshotConnection != null) {
-                                mContext.unbindService(mScreenshotConnection);
-                                mScreenshotConnection = null;
-                                mScreenshotService = null;
+                                resetConnection();
                                 // only log an error if we're still within the timeout period
                                 if (handler.hasCallbacks(mScreenshotTimeout)) {
                                     handler.removeCallbacks(mScreenshotTimeout);
@@ -382,6 +393,17 @@ public class ScreenshotHelper {
     }
 
     /**
+     * Unbinds the current screenshot connection (if any).
+     */
+    private void resetConnection() {
+        if (mScreenshotConnection != null) {
+            mContext.unbindService(mScreenshotConnection);
+            mScreenshotConnection = null;
+            mScreenshotService = null;
+        }
+    }
+
+    /**
      * Notifies the screenshot service to show an error.
      */
     private void notifyScreenshotError() {