1 /*
<lambda>null2  * Copyright (C) 2024 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.systemui.screenshot
18 
19 import android.net.Uri
20 import android.os.Trace
21 import android.util.Log
22 import android.view.Display
23 import android.view.WindowManager.ScreenshotSource
24 import android.view.WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE
25 import com.android.internal.logging.UiEventLogger
26 import com.android.internal.util.ScreenshotRequest
27 import com.android.systemui.Flags.screenshotMultidisplayFocusChange
28 import com.android.systemui.dagger.SysUISingleton
29 import com.android.systemui.dagger.qualifiers.Application
30 import com.android.systemui.display.data.repository.DisplayRepository
31 import com.android.systemui.display.data.repository.FocusedDisplayRepository
32 import com.android.systemui.res.R
33 import com.android.systemui.screenshot.ScreenshotEvent.SCREENSHOT_CAPTURE_FAILED
34 import com.android.systemui.screenshot.ScreenshotEvent.SCREENSHOT_DISMISSED_OTHER
35 import com.android.systemui.screenshot.TakeScreenshotService.RequestCallback
36 import java.util.function.Consumer
37 import javax.inject.Inject
38 import kotlinx.coroutines.CoroutineScope
39 import kotlinx.coroutines.flow.first
40 import com.android.app.tracing.coroutines.launchTraced as launch
41 
42 interface TakeScreenshotExecutor {
43     suspend fun executeScreenshots(
44         screenshotRequest: ScreenshotRequest,
45         onSaved: (Uri?) -> Unit,
46         requestCallback: RequestCallback,
47     )
48 
49     fun onCloseSystemDialogsReceived()
50 
51     fun removeWindows()
52 
53     fun onDestroy()
54 
55     fun executeScreenshotsAsync(
56         screenshotRequest: ScreenshotRequest,
57         onSaved: Consumer<Uri?>,
58         requestCallback: RequestCallback,
59     )
60 }
61 
62 interface ScreenshotHandler {
handleScreenshotnull63     fun handleScreenshot(
64         screenshot: ScreenshotData,
65         finisher: Consumer<Uri?>,
66         requestCallback: RequestCallback,
67     )
68 }
69 
70 /**
71  * Receives the signal to take a screenshot from [TakeScreenshotService], and calls back with the
72  * result.
73  *
74  * Captures a screenshot for each [Display] available.
75  */
76 @SysUISingleton
77 class TakeScreenshotExecutorImpl
78 @Inject
79 constructor(
80     private val interactiveScreenshotHandlerFactory: InteractiveScreenshotHandler.Factory,
81     private val displayRepository: DisplayRepository,
82     @Application private val mainScope: CoroutineScope,
83     private val screenshotRequestProcessor: ScreenshotRequestProcessor,
84     private val uiEventLogger: UiEventLogger,
85     private val screenshotNotificationControllerFactory: ScreenshotNotificationsController.Factory,
86     private val headlessScreenshotHandler: HeadlessScreenshotHandler,
87     private val focusedDisplayRepository: FocusedDisplayRepository,
88 ) : TakeScreenshotExecutor {
89     private val displays = displayRepository.displays
90     private var screenshotController: InteractiveScreenshotHandler? = null
91     private val notificationControllers = mutableMapOf<Int, ScreenshotNotificationsController>()
92 
93     /**
94      * Executes the [ScreenshotRequest].
95      *
96      * [onSaved] is invoked only on the default display result. [RequestCallback.onFinish] is
97      * invoked only when both screenshot UIs are removed.
98      */
99     override suspend fun executeScreenshots(
100         screenshotRequest: ScreenshotRequest,
101         onSaved: (Uri?) -> Unit,
102         requestCallback: RequestCallback,
103     ) {
104         if (screenshotMultidisplayFocusChange()) {
105             val display = getDisplayToScreenshot(screenshotRequest)
106             val screenshotHandler = getScreenshotController(display)
107             dispatchToController(
108                 screenshotHandler,
109                 ScreenshotData.fromRequest(screenshotRequest, display.displayId),
110                 onSaved,
111                 requestCallback,
112             )
113         } else {
114             val displays = getDisplaysToScreenshot(screenshotRequest.type)
115             val resultCallbackWrapper = MultiResultCallbackWrapper(requestCallback)
116             if (displays.isEmpty()) {
117                 Log.e(TAG, "No displays found for screenshot.")
118             }
119 
120             displays.forEach { display ->
121                 val displayId = display.displayId
122                 var screenshotHandler: ScreenshotHandler =
123                     if (displayId == Display.DEFAULT_DISPLAY) {
124                         getScreenshotController(display)
125                     } else {
126                         headlessScreenshotHandler
127                     }
128 
129                 Log.d(TAG, "Executing screenshot for display $displayId")
130                 dispatchToController(
131                     screenshotHandler,
132                     rawScreenshotData = ScreenshotData.fromRequest(screenshotRequest, displayId),
133                     onSaved =
134                         if (displayId == Display.DEFAULT_DISPLAY) {
135                             onSaved
136                         } else { _ -> },
137                     callback = resultCallbackWrapper.createCallbackForId(displayId),
138                 )
139             }
140         }
141     }
142 
143     /** All logging should be triggered only by this method. */
144     private suspend fun dispatchToController(
145         screenshotHandler: ScreenshotHandler,
146         rawScreenshotData: ScreenshotData,
147         onSaved: (Uri?) -> Unit,
148         callback: RequestCallback,
149     ) {
150         // Let's wait before logging "screenshot requested", as we should log the processed
151         // ScreenshotData.
152         val screenshotData =
153             runCatching { screenshotRequestProcessor.process(rawScreenshotData) }
154                 .onFailure {
155                     Log.e(TAG, "Failed to process screenshot request!", it)
156                     logScreenshotRequested(rawScreenshotData)
157                     onFailedScreenshotRequest(rawScreenshotData, callback)
158                 }
159                 .getOrNull() ?: return
160 
161         logScreenshotRequested(screenshotData)
162         Log.d(TAG, "Screenshot request: $screenshotData")
163         try {
164             screenshotHandler.handleScreenshot(screenshotData, onSaved, callback)
165         } catch (e: IllegalStateException) {
166             Log.e(TAG, "Error while ScreenshotController was handling ScreenshotData!", e)
167             onFailedScreenshotRequest(screenshotData, callback)
168             return // After a failure log, nothing else should run.
169         }
170     }
171 
172     /**
173      * This should be logged also in case of failed requests, before the [SCREENSHOT_CAPTURE_FAILED]
174      * event.
175      */
176     private fun logScreenshotRequested(screenshotData: ScreenshotData) {
177         uiEventLogger.log(
178             ScreenshotEvent.getScreenshotSource(screenshotData.source),
179             0,
180             screenshotData.packageNameString,
181         )
182     }
183 
184     private fun onFailedScreenshotRequest(
185         screenshotData: ScreenshotData,
186         callback: RequestCallback,
187     ) {
188         uiEventLogger.log(SCREENSHOT_CAPTURE_FAILED, 0, screenshotData.packageNameString)
189         getNotificationController(screenshotData.displayId)
190             .notifyScreenshotError(R.string.screenshot_failed_to_capture_text)
191         callback.reportError()
192     }
193 
194     private suspend fun getDisplaysToScreenshot(requestType: Int): List<Display> {
195         val allDisplays = displays.first()
196         return if (requestType == TAKE_SCREENSHOT_PROVIDED_IMAGE) {
197             // If this is a provided image just screenshot th default display
198             allDisplays.filter { it.displayId == Display.DEFAULT_DISPLAY }
199         } else {
200             allDisplays.filter { it.type in ALLOWED_DISPLAY_TYPES }
201         }
202     }
203 
204     // Return the single display to be screenshot based upon the request.
205     private suspend fun getDisplayToScreenshot(screenshotRequest: ScreenshotRequest): Display {
206         return when (screenshotRequest.source) {
207             ScreenshotSource.SCREENSHOT_OVERVIEW ->
208                 // Show on the display where overview was shown if available.
209                 displayRepository.getDisplay(screenshotRequest.displayId)
210                     ?: displayRepository.getDisplay(Display.DEFAULT_DISPLAY)
211                     ?: error("Can't find default display")
212 
213             // Key chord and vendor gesture occur on the device itself, so screenshot the device's
214             // display
215             ScreenshotSource.SCREENSHOT_KEY_CHORD,
216             ScreenshotSource.SCREENSHOT_VENDOR_GESTURE ->
217                 displayRepository.getDisplay(Display.DEFAULT_DISPLAY)
218                     ?: error("Can't find default display")
219 
220             // All other invocations use the focused display
221             else ->
222                 displayRepository.getDisplay(focusedDisplayRepository.focusedDisplayId.value)
223                     ?: displayRepository.getDisplay(Display.DEFAULT_DISPLAY)
224                     ?: error("Can't find default display")
225         }
226     }
227 
228     /** Propagates the close system dialog signal to the ScreenshotController. */
229     override fun onCloseSystemDialogsReceived() {
230         if (screenshotController?.isPendingSharedTransition() == false) {
231             screenshotController?.requestDismissal(SCREENSHOT_DISMISSED_OTHER)
232         }
233     }
234 
235     /** Removes all screenshot related windows. */
236     override fun removeWindows() {
237         screenshotController?.removeWindow()
238     }
239 
240     /**
241      * Destroys the executor. Afterwards, this class is not expected to work as intended anymore.
242      */
243     override fun onDestroy() {
244         screenshotController?.onDestroy()
245         screenshotController = null
246     }
247 
248     private fun getNotificationController(id: Int): ScreenshotNotificationsController {
249         return notificationControllers.computeIfAbsent(id) {
250             screenshotNotificationControllerFactory.create(id)
251         }
252     }
253 
254     /** For java compatibility only. see [executeScreenshots] */
255     override fun executeScreenshotsAsync(
256         screenshotRequest: ScreenshotRequest,
257         onSaved: Consumer<Uri?>,
258         requestCallback: RequestCallback,
259     ) {
260         mainScope.launch {
261             executeScreenshots(screenshotRequest, { uri -> onSaved.accept(uri) }, requestCallback)
262         }
263     }
264 
265     private fun getScreenshotController(display: Display): InteractiveScreenshotHandler {
266         val controller = screenshotController ?: interactiveScreenshotHandlerFactory.create(display)
267         screenshotController = controller
268         return controller
269     }
270 
271     /**
272      * Returns a [RequestCallback] that wraps [originalCallback].
273      *
274      * Each [RequestCallback] created with [createCallbackForId] is expected to be used with either
275      * [reportError] or [onFinish]. Once they are both called:
276      * - If any finished with an error, [reportError] of [originalCallback] is called
277      * - Otherwise, [onFinish] is called.
278      */
279     private class MultiResultCallbackWrapper(private val originalCallback: RequestCallback) {
280         private val idsPending = mutableSetOf<Int>()
281         private val idsWithErrors = mutableSetOf<Int>()
282 
283         /**
284          * Creates a callback for [id].
285          *
286          * [originalCallback]'s [onFinish] will be called only when this (and the other created)
287          * callback's [onFinish] have been called.
288          */
289         fun createCallbackForId(id: Int): RequestCallback {
290             Trace.asyncTraceForTrackBegin(Trace.TRACE_TAG_APP, TAG, "Waiting for id=$id", id)
291             idsPending += id
292             return object : RequestCallback {
293                 override fun reportError() {
294                     endTrace("reportError id=$id")
295                     idsWithErrors += id
296                     idsPending -= id
297                     reportToOriginalIfNeeded()
298                 }
299 
300                 override fun onFinish() {
301                     endTrace("onFinish id=$id")
302                     idsPending -= id
303                     reportToOriginalIfNeeded()
304                 }
305 
306                 private fun endTrace(reason: String) {
307                     Log.d(TAG, "Finished waiting for id=$id. $reason")
308                     Trace.asyncTraceForTrackEnd(Trace.TRACE_TAG_APP, TAG, id)
309                     Trace.instantForTrack(Trace.TRACE_TAG_APP, TAG, reason)
310                 }
311             }
312         }
313 
314         private fun reportToOriginalIfNeeded() {
315             if (idsPending.isNotEmpty()) return
316             if (idsWithErrors.isEmpty()) {
317                 originalCallback.onFinish()
318             } else {
319                 originalCallback.reportError()
320             }
321         }
322     }
323 
324     private companion object {
325         val TAG = LogConfig.logTag(TakeScreenshotService::class.java)
326 
327         val ALLOWED_DISPLAY_TYPES =
328             listOf(
329                 Display.TYPE_EXTERNAL,
330                 Display.TYPE_INTERNAL,
331                 Display.TYPE_OVERLAY,
332                 Display.TYPE_WIFI,
333             )
334     }
335 }
336