1 package com.android.internal.util;
2 
3 import static android.content.Intent.ACTION_USER_SWITCHED;
4 import static android.view.WindowManager.TAKE_SCREENSHOT_FULLSCREEN;
5 
6 import android.annotation.NonNull;
7 import android.annotation.Nullable;
8 import android.content.BroadcastReceiver;
9 import android.content.ComponentName;
10 import android.content.Context;
11 import android.content.Intent;
12 import android.content.IntentFilter;
13 import android.content.ServiceConnection;
14 import android.net.Uri;
15 import android.os.Handler;
16 import android.os.IBinder;
17 import android.os.Message;
18 import android.os.Messenger;
19 import android.os.RemoteException;
20 import android.os.UserHandle;
21 import android.util.Log;
22 import android.view.WindowManager.ScreenshotSource;
23 
24 import com.android.internal.annotations.VisibleForTesting;
25 
26 import java.util.function.Consumer;
27 
28 public class ScreenshotHelper {
29 
30     public static final int SCREENSHOT_MSG_URI = 1;
31     public static final int SCREENSHOT_MSG_PROCESS_COMPLETE = 2;
32 
33     private static final String TAG = "ScreenshotHelper";
34 
35     // Time until we give up on the screenshot & show an error instead.
36     private final int SCREENSHOT_TIMEOUT_MS = 10000;
37 
38     private final Object mScreenshotLock = new Object();
39     private IBinder mScreenshotService = null;
40     private ServiceConnection mScreenshotConnection = null;
41     private final Context mContext;
42 
43     private final BroadcastReceiver mBroadcastReceiver = new BroadcastReceiver() {
44         @Override
45         public void onReceive(Context context, Intent intent) {
46             synchronized (mScreenshotLock) {
47                 if (ACTION_USER_SWITCHED.equals(intent.getAction())) {
48                     resetConnection();
49                 }
50             }
51         }
52     };
53 
ScreenshotHelper(Context context)54     public ScreenshotHelper(Context context) {
55         mContext = context;
56         IntentFilter filter = new IntentFilter(ACTION_USER_SWITCHED);
57         mContext.registerReceiver(mBroadcastReceiver, filter, Context.RECEIVER_EXPORTED);
58     }
59 
60     /**
61      * Request a screenshot be taken.
62      * <p>
63      * Convenience method for taking a full screenshot with provided source.
64      *
65      * @param source             source of the screenshot request, defined by {@link
66      *                           ScreenshotSource}
67      * @param handler            used to process messages received from the screenshot service
68      * @param completionConsumer receives the URI of the captured screenshot, once saved or
69      *                           null if no screenshot was saved
70      */
takeScreenshot(@creenshotSource int source, @NonNull Handler handler, @Nullable Consumer<Uri> completionConsumer)71     public void takeScreenshot(@ScreenshotSource int source, @NonNull Handler handler,
72             @Nullable Consumer<Uri> completionConsumer) {
73         ScreenshotRequest request =
74                 new ScreenshotRequest.Builder(TAKE_SCREENSHOT_FULLSCREEN, source).build();
75         takeScreenshot(request, handler, completionConsumer);
76     }
77 
78     /**
79      * Request a screenshot be taken.
80      * <p>
81      *
82      * @param request            description of the screenshot request, either for taking a
83      *                           screenshot or
84      *                           providing a bitmap
85      * @param handler            used to process messages received from the screenshot service
86      * @param completionConsumer receives the URI of the captured screenshot, once saved or
87      *                           null if no screenshot was saved
88      */
takeScreenshot(ScreenshotRequest request, @NonNull Handler handler, @Nullable Consumer<Uri> completionConsumer)89     public void takeScreenshot(ScreenshotRequest request, @NonNull Handler handler,
90             @Nullable Consumer<Uri> completionConsumer) {
91         takeScreenshotInternal(request, handler, completionConsumer, SCREENSHOT_TIMEOUT_MS);
92     }
93 
94     /**
95      * Request a screenshot be taken.
96      * <p>
97      * Added to support reducing unit test duration; the method variant without a timeout argument
98      * is recommended for general use.
99      *
100      * @param request            description of the screenshot request, either for taking a
101      *                           screenshot or providing a bitmap
102      * @param handler            used to process messages received from the screenshot service
103      * @param timeoutMs          time limit for processing, intended only for testing
104      * @param completionConsumer receives the URI of the captured screenshot, once saved or
105      *                           null if no screenshot was saved
106      */
107     @VisibleForTesting
takeScreenshotInternal(ScreenshotRequest request, @NonNull Handler handler, @Nullable Consumer<Uri> completionConsumer, long timeoutMs)108     public void takeScreenshotInternal(ScreenshotRequest request, @NonNull Handler handler,
109             @Nullable Consumer<Uri> completionConsumer, long timeoutMs) {
110         synchronized (mScreenshotLock) {
111 
112             final Runnable mScreenshotTimeout = () -> {
113                 synchronized (mScreenshotLock) {
114                     if (mScreenshotConnection != null) {
115                         Log.e(TAG, "Timed out before getting screenshot capture response");
116                         resetConnection();
117                         notifyScreenshotError();
118                     }
119                 }
120                 if (completionConsumer != null) {
121                     completionConsumer.accept(null);
122                 }
123             };
124 
125             Message msg = Message.obtain(null, 0, request);
126 
127             Handler h = new Handler(handler.getLooper()) {
128                 @Override
129                 public void handleMessage(Message msg) {
130                     switch (msg.what) {
131                         case SCREENSHOT_MSG_URI:
132                             if (completionConsumer != null) {
133                                 completionConsumer.accept((Uri) msg.obj);
134                             }
135                             handler.removeCallbacks(mScreenshotTimeout);
136                             break;
137                         case SCREENSHOT_MSG_PROCESS_COMPLETE:
138                             synchronized (mScreenshotLock) {
139                                 resetConnection();
140                             }
141                             break;
142                     }
143                 }
144             };
145             msg.replyTo = new Messenger(h);
146 
147             if (mScreenshotConnection == null || mScreenshotService == null) {
148                 if (mScreenshotConnection != null) {
149                     resetConnection();
150                 }
151                 final ComponentName serviceComponent = ComponentName.unflattenFromString(
152                         mContext.getResources().getString(
153                                 com.android.internal.R.string.config_screenshotServiceComponent));
154                 final Intent serviceIntent = new Intent();
155 
156                 serviceIntent.setComponent(serviceComponent);
157                 ServiceConnection conn = new ServiceConnection() {
158                     @Override
159                     public void onServiceConnected(ComponentName name, IBinder service) {
160                         synchronized (mScreenshotLock) {
161                             if (mScreenshotConnection != this) {
162                                 return;
163                             }
164                             mScreenshotService = service;
165                             Messenger messenger = new Messenger(mScreenshotService);
166 
167                             try {
168                                 messenger.send(msg);
169                             } catch (RemoteException e) {
170                                 Log.e(TAG, "Couldn't take screenshot: " + e);
171                                 if (completionConsumer != null) {
172                                     completionConsumer.accept(null);
173                                 }
174                             }
175                         }
176                     }
177 
178                     @Override
179                     public void onServiceDisconnected(ComponentName name) {
180                         synchronized (mScreenshotLock) {
181                             if (mScreenshotConnection != null) {
182                                 resetConnection();
183                                 // only log an error if we're still within the timeout period
184                                 if (handler.hasCallbacks(mScreenshotTimeout)) {
185                                     Log.e(TAG, "Screenshot service disconnected");
186                                     handler.removeCallbacks(mScreenshotTimeout);
187                                     notifyScreenshotError();
188                                 }
189                             }
190                         }
191                     }
192                 };
193                 if (mContext.bindServiceAsUser(serviceIntent, conn,
194                         Context.BIND_AUTO_CREATE | Context.BIND_FOREGROUND_SERVICE,
195                         UserHandle.CURRENT)) {
196                     mScreenshotConnection = conn;
197                     handler.postDelayed(mScreenshotTimeout, timeoutMs);
198                 } else {
199                     mContext.unbindService(conn);
200                 }
201             } else {
202                 Messenger messenger = new Messenger(mScreenshotService);
203 
204                 try {
205                     messenger.send(msg);
206                 } catch (RemoteException e) {
207                     Log.e(TAG, "Couldn't take screenshot: " + e);
208                     if (completionConsumer != null) {
209                         completionConsumer.accept(null);
210                     }
211                 }
212                 handler.postDelayed(mScreenshotTimeout, timeoutMs);
213             }
214         }
215     }
216 
217     /**
218      * Unbinds the current screenshot connection (if any).
219      */
resetConnection()220     private void resetConnection() {
221         if (mScreenshotConnection != null) {
222             mContext.unbindService(mScreenshotConnection);
223             mScreenshotConnection = null;
224             mScreenshotService = null;
225         }
226     }
227 
228     /**
229      * Notifies the screenshot service to show an error.
230      */
notifyScreenshotError()231     private void notifyScreenshotError() {
232         // If the service process is killed, then ask it to clean up after itself
233         final ComponentName errorComponent = ComponentName.unflattenFromString(
234                 mContext.getResources().getString(
235                         com.android.internal.R.string.config_screenshotErrorReceiverComponent));
236         // Broadcast needs to have a valid action.  We'll just pick
237         // a generic one, since the receiver here doesn't care.
238         Intent errorIntent = new Intent(Intent.ACTION_USER_PRESENT);
239         errorIntent.setComponent(errorComponent);
240         errorIntent.addFlags(Intent.FLAG_RECEIVER_REGISTERED_ONLY_BEFORE_BOOT |
241                 Intent.FLAG_RECEIVER_FOREGROUND);
242         mContext.sendBroadcastAsUser(errorIntent, UserHandle.CURRENT);
243     }
244 }
245