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.TAKE_SCREENSHOT_PROVIDED_IMAGE
24 import com.android.internal.logging.UiEventLogger
25 import com.android.internal.util.ScreenshotRequest
26 import com.android.systemui.dagger.SysUISingleton
27 import com.android.systemui.dagger.qualifiers.Application
28 import com.android.systemui.display.data.repository.DisplayRepository
29 import com.android.systemui.res.R
30 import com.android.systemui.screenshot.ScreenshotEvent.SCREENSHOT_CAPTURE_FAILED
31 import com.android.systemui.screenshot.ScreenshotEvent.SCREENSHOT_DISMISSED_OTHER
32 import com.android.systemui.screenshot.TakeScreenshotService.RequestCallback
33 import java.util.function.Consumer
34 import javax.inject.Inject
35 import kotlinx.coroutines.CoroutineScope
36 import kotlinx.coroutines.flow.first
37 import kotlinx.coroutines.launch
38 
39 interface TakeScreenshotExecutor {
40     suspend fun executeScreenshots(
41         screenshotRequest: ScreenshotRequest,
42         onSaved: (Uri?) -> Unit,
43         requestCallback: RequestCallback
44     )
45 
46     fun onCloseSystemDialogsReceived()
47 
48     fun removeWindows()
49 
50     fun onDestroy()
51 
52     fun executeScreenshotsAsync(
53         screenshotRequest: ScreenshotRequest,
54         onSaved: Consumer<Uri?>,
55         requestCallback: RequestCallback
56     )
57 }
58 
59 interface ScreenshotHandler {
handleScreenshotnull60     fun handleScreenshot(
61         screenshot: ScreenshotData,
62         finisher: Consumer<Uri?>,
63         requestCallback: RequestCallback
64     )
65 }
66 
67 /**
68  * Receives the signal to take a screenshot from [TakeScreenshotService], and calls back with the
69  * result.
70  *
71  * Captures a screenshot for each [Display] available.
72  */
73 @SysUISingleton
74 class TakeScreenshotExecutorImpl
75 @Inject
76 constructor(
77     private val screenshotControllerFactory: ScreenshotController.Factory,
78     displayRepository: DisplayRepository,
79     @Application private val mainScope: CoroutineScope,
80     private val screenshotRequestProcessor: ScreenshotRequestProcessor,
81     private val uiEventLogger: UiEventLogger,
82     private val screenshotNotificationControllerFactory: ScreenshotNotificationsController.Factory,
83     private val headlessScreenshotHandler: HeadlessScreenshotHandler,
84 ) : TakeScreenshotExecutor {
85     private val displays = displayRepository.displays
86     private var screenshotController: ScreenshotController? = null
87     private val notificationControllers = mutableMapOf<Int, ScreenshotNotificationsController>()
88 
89     /**
90      * Executes the [ScreenshotRequest].
91      *
92      * [onSaved] is invoked only on the default display result. [RequestCallback.onFinish] is
93      * invoked only when both screenshot UIs are removed.
94      */
95     override suspend fun executeScreenshots(
96         screenshotRequest: ScreenshotRequest,
97         onSaved: (Uri?) -> Unit,
98         requestCallback: RequestCallback
99     ) {
100         val displays = getDisplaysToScreenshot(screenshotRequest.type)
101         val resultCallbackWrapper = MultiResultCallbackWrapper(requestCallback)
102         displays.forEach { display ->
103             val displayId = display.displayId
104             var screenshotHandler: ScreenshotHandler =
105                 if (displayId == Display.DEFAULT_DISPLAY) {
106                     getScreenshotController(display)
107                 } else {
108                     headlessScreenshotHandler
109                 }
110             Log.d(TAG, "Executing screenshot for display $displayId")
111             dispatchToController(
112                 screenshotHandler,
113                 rawScreenshotData = ScreenshotData.fromRequest(screenshotRequest, displayId),
114                 onSaved =
115                     if (displayId == Display.DEFAULT_DISPLAY) {
116                         onSaved
117                     } else { _ -> },
118                 callback = resultCallbackWrapper.createCallbackForId(displayId)
119             )
120         }
121     }
122 
123     /** All logging should be triggered only by this method. */
124     private suspend fun dispatchToController(
125         screenshotHandler: ScreenshotHandler,
126         rawScreenshotData: ScreenshotData,
127         onSaved: (Uri?) -> Unit,
128         callback: RequestCallback
129     ) {
130         // Let's wait before logging "screenshot requested", as we should log the processed
131         // ScreenshotData.
132         val screenshotData =
133             runCatching { screenshotRequestProcessor.process(rawScreenshotData) }
134                 .onFailure {
135                     Log.e(TAG, "Failed to process screenshot request!", it)
136                     logScreenshotRequested(rawScreenshotData)
137                     onFailedScreenshotRequest(rawScreenshotData, callback)
138                 }
139                 .getOrNull() ?: return
140 
141         logScreenshotRequested(screenshotData)
142         Log.d(TAG, "Screenshot request: $screenshotData")
143         try {
144             screenshotHandler.handleScreenshot(screenshotData, onSaved, callback)
145         } catch (e: IllegalStateException) {
146             Log.e(TAG, "Error while ScreenshotController was handling ScreenshotData!", e)
147             onFailedScreenshotRequest(screenshotData, callback)
148             return // After a failure log, nothing else should run.
149         }
150     }
151 
152     /**
153      * This should be logged also in case of failed requests, before the [SCREENSHOT_CAPTURE_FAILED]
154      * event.
155      */
156     private fun logScreenshotRequested(screenshotData: ScreenshotData) {
157         uiEventLogger.log(
158             ScreenshotEvent.getScreenshotSource(screenshotData.source),
159             0,
160             screenshotData.packageNameString
161         )
162     }
163 
164     private fun onFailedScreenshotRequest(
165         screenshotData: ScreenshotData,
166         callback: RequestCallback
167     ) {
168         uiEventLogger.log(SCREENSHOT_CAPTURE_FAILED, 0, screenshotData.packageNameString)
169         getNotificationController(screenshotData.displayId)
170             .notifyScreenshotError(R.string.screenshot_failed_to_capture_text)
171         callback.reportError()
172     }
173 
174     private suspend fun getDisplaysToScreenshot(requestType: Int): List<Display> {
175         val allDisplays = displays.first()
176         return if (requestType == TAKE_SCREENSHOT_PROVIDED_IMAGE) {
177             // If this is a provided image, let's show the UI on the default display only.
178             allDisplays.filter { it.displayId == Display.DEFAULT_DISPLAY }
179         } else {
180             allDisplays.filter { it.type in ALLOWED_DISPLAY_TYPES }
181         }
182     }
183 
184     /** Propagates the close system dialog signal to the ScreenshotController. */
185     override fun onCloseSystemDialogsReceived() {
186         if (screenshotController?.isPendingSharedTransition == false) {
187             screenshotController?.requestDismissal(SCREENSHOT_DISMISSED_OTHER)
188         }
189     }
190 
191     /** Removes all screenshot related windows. */
192     override fun removeWindows() {
193         screenshotController?.removeWindow()
194     }
195 
196     /**
197      * Destroys the executor. Afterwards, this class is not expected to work as intended anymore.
198      */
199     override fun onDestroy() {
200         screenshotController?.onDestroy()
201         screenshotController = null
202     }
203 
204     private fun getNotificationController(id: Int): ScreenshotNotificationsController {
205         return notificationControllers.computeIfAbsent(id) {
206             screenshotNotificationControllerFactory.create(id)
207         }
208     }
209 
210     /** For java compatibility only. see [executeScreenshots] */
211     override fun executeScreenshotsAsync(
212         screenshotRequest: ScreenshotRequest,
213         onSaved: Consumer<Uri?>,
214         requestCallback: RequestCallback
215     ) {
216         mainScope.launch {
217             executeScreenshots(screenshotRequest, { uri -> onSaved.accept(uri) }, requestCallback)
218         }
219     }
220 
221     private fun getScreenshotController(display: Display): ScreenshotController {
222         val controller = screenshotController ?: screenshotControllerFactory.create(display, false)
223         screenshotController = controller
224         return controller
225     }
226 
227     /**
228      * Returns a [RequestCallback] that wraps [originalCallback].
229      *
230      * Each [RequestCallback] created with [createCallbackForId] is expected to be used with either
231      * [reportError] or [onFinish]. Once they are both called:
232      * - If any finished with an error, [reportError] of [originalCallback] is called
233      * - Otherwise, [onFinish] is called.
234      */
235     private class MultiResultCallbackWrapper(
236         private val originalCallback: RequestCallback,
237     ) {
238         private val idsPending = mutableSetOf<Int>()
239         private val idsWithErrors = mutableSetOf<Int>()
240 
241         /**
242          * Creates a callback for [id].
243          *
244          * [originalCallback]'s [onFinish] will be called only when this (and the other created)
245          * callback's [onFinish] have been called.
246          */
247         fun createCallbackForId(id: Int): RequestCallback {
248             Trace.asyncTraceForTrackBegin(Trace.TRACE_TAG_APP, TAG, "Waiting for id=$id", id)
249             idsPending += id
250             return object : RequestCallback {
251                 override fun reportError() {
252                     endTrace("reportError id=$id")
253                     idsWithErrors += id
254                     idsPending -= id
255                     reportToOriginalIfNeeded()
256                 }
257 
258                 override fun onFinish() {
259                     endTrace("onFinish id=$id")
260                     idsPending -= id
261                     reportToOriginalIfNeeded()
262                 }
263 
264                 private fun endTrace(reason: String) {
265                     Log.d(TAG, "Finished waiting for id=$id. $reason")
266                     Trace.asyncTraceForTrackEnd(Trace.TRACE_TAG_APP, TAG, id)
267                     Trace.instantForTrack(Trace.TRACE_TAG_APP, TAG, reason)
268                 }
269             }
270         }
271 
272         private fun reportToOriginalIfNeeded() {
273             if (idsPending.isNotEmpty()) return
274             if (idsWithErrors.isEmpty()) {
275                 originalCallback.onFinish()
276             } else {
277                 originalCallback.reportError()
278             }
279         }
280     }
281 
282     private companion object {
283         val TAG = LogConfig.logTag(TakeScreenshotService::class.java)
284 
285         val ALLOWED_DISPLAY_TYPES =
286             listOf(
287                 Display.TYPE_EXTERNAL,
288                 Display.TYPE_INTERNAL,
289                 Display.TYPE_OVERLAY,
290                 Display.TYPE_WIFI
291             )
292     }
293 }
294