[DO NOT MERGE] Close screenshot process on user switched

Currently, we keep the process up even if the user switches,
meaning that in some cases (if the user is switched while the
screenshot UI is up) we will save images to the wrong profile.
This change makes ScreenshotHelper listen for user switches and
close the screenshot service, so that a new screenshot is
guaranteed to be constructed with the correct user's context.

Bug: 170474245
Fix: 170474245
Test: manual -- verified bad state occurs if user switches within
the timeout period, ensured that screenshots work immediately
after switching with this change.

Change-Id: I9d32d0928e6c2bda161d04555438d0dd7afef0ba
(cherry picked from commit 7ef1a5dd1506075507412626f2533283d9520144)
(cherry picked from commit 05f7aef1b07665aa10aab535613a8457296279a3)

diff --git a/core/java/com/android/internal/util/ScreenshotHelper.java b/core/java/com/android/internal/util/ScreenshotHelper.java
index a23fc4b..7ee846e 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 @@
     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 @@
             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();
                     }
                 }
@@ -304,11 +319,7 @@
                             break;
                         case SCREENSHOT_MSG_PROCESS_COMPLETE:
                             synchronized (mScreenshotLock) {
-                                if (mScreenshotConnection != null) {
-                                    mContext.unbindService(mScreenshotConnection);
-                                    mScreenshotConnection = null;
-                                    mScreenshotService = null;
-                                }
+                                resetConnection();
                             }
                             break;
                     }
@@ -348,9 +359,7 @@
                     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);
@@ -383,6 +392,17 @@
     }
 
     /**
+     * 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() {