1 /*
2 * Copyright (c) 2023 Huawei Device Co., Ltd.
3 * Licensed under the Apache License, Version 2.0 (the "License");
4 * you may not use this file except in compliance with the License.
5 * You may obtain a copy of the License at
6 *
7 * http://www.apache.org/licenses/LICENSE-2.0
8 *
9 * Unless required by applicable law or agreed to in writing, software
10 * distributed under the License is distributed on an "AS IS" BASIS,
11 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 * See the License for the specific language governing permissions and
13 * limitations under the License.
14 */
15
16 #include <display_manager.h>
17 #include <iconsumer_surface.h>
18 #include <jpeglib.h>
19 #include <pixel_map.h>
20 #include <refbase.h>
21 #include <screen_manager.h>
22 #include <securec.h>
23 #include <csetjmp>
24 #include <surface.h>
25 #include <poll.h>
26 #include "common_utilities_hpp.h"
27 #include "screen_copy.h"
28
29 namespace OHOS::uitest {
30 using namespace std;
31 using namespace OHOS;
32 using namespace OHOS::Rosen;
33
34 using OnBufferAvailableHandler = function<void()>;
35 class BufferConsumerListener : public IBufferConsumerListener {
36 public:
BufferConsumerListener(OnBufferAvailableHandler handler)37 explicit BufferConsumerListener(OnBufferAvailableHandler handler): handler_(handler) {}
OnBufferAvailable()38 void OnBufferAvailable() override
39 {
40 if (handler_ != nullptr) {
41 handler_();
42 }
43 }
44
45 private:
46 OnBufferAvailableHandler handler_ = nullptr;
47 };
48
49 class ScreenCopy {
50 public:
51 ScreenCopy() = default;
52 virtual ~ScreenCopy();
53 bool Setup();
54 bool Run();
55 bool Pause();
56 void Destroy();
57 void HandleConsumerBuffer();
58 private:
59 void ConsumeFrameBuffer(const sptr<SurfaceBuffer> &buf);
60 sptr<IBufferConsumerListener> bufferListener_ = nullptr;
61 sptr<IConsumerSurface> consumerSurface_ = nullptr;
62 sptr<Surface> producerSurface_ = nullptr;
63 pair<uint32_t, uint32_t> screenSize_;
64 ScreenId mainScreenId_ = SCREEN_ID_INVALID;
65 ScreenId virtualScreenId_ = SCREEN_ID_INVALID;
66 bool workable_ = false;
67 };
68 static unique_ptr<ScreenCopy> g_screenCopy = nullptr;
69 static ScreenCopyHandler g_screenCopyHandler = nullptr;
70
~ScreenCopy()71 ScreenCopy::~ScreenCopy()
72 {
73 Destroy();
74 }
75
Destroy()76 void ScreenCopy::Destroy()
77 {
78 Pause();
79 producerSurface_ = nullptr;
80 if (consumerSurface_ != nullptr) {
81 consumerSurface_->UnregisterConsumerListener();
82 consumerSurface_ = nullptr;
83 }
84 }
85
Setup()86 bool ScreenCopy::Setup()
87 {
88 // create pruducer surface and consumer surface
89 consumerSurface_ = IConsumerSurface::Create();
90 if (consumerSurface_ == nullptr) {
91 LOG_E("Failed to create IConsumerSurface");
92 return false;
93 }
94 auto producer = consumerSurface_->GetProducer();
95 producerSurface_ = Surface::CreateSurfaceAsProducer(producer);
96 if (producerSurface_ == nullptr) {
97 LOG_E("Failed to CreateSurfaceAsProducer");
98 return false;
99 }
100 auto handler = [this]() {
101 this->HandleConsumerBuffer();
102 };
103 bufferListener_ = new BufferConsumerListener(handler);
104 if (consumerSurface_->RegisterConsumerListener(bufferListener_) != 0) {
105 LOG_E("Failed to RegisterConsumerListener");
106 return false;
107 }
108 // make screen mirror from main screen to accept frames with producer surface buffer
109 mainScreenId_ = static_cast<ScreenId>(DisplayManager::GetInstance().GetDefaultDisplayId());
110 auto mainScreen = ScreenManager::GetInstance().GetScreenById(mainScreenId_);
111 if (mainScreenId_ == SCREEN_ID_INVALID || mainScreen == nullptr) {
112 LOG_E("Get main screen failed!");
113 return false;
114 }
115 screenSize_.first = mainScreen->GetWidth();
116 screenSize_.second = mainScreen->GetHeight();
117 workable_ = true;
118 return true;
119 }
120
Run()121 bool ScreenCopy::Run()
122 {
123 if (!workable_) {
124 return false;
125 }
126 if (virtualScreenId_ != SCREEN_ID_INVALID) {
127 LOG_W("ScreenCopy already running!");
128 return false;
129 }
130 VirtualScreenOption option = {
131 .name_ = "virtualScreen",
132 .width_ = screenSize_.first, // * scale_,
133 .height_ = screenSize_.second, // * scale_,
134 .density_ = 2.0,
135 .surface_ = producerSurface_,
136 .flags_ = 0,
137 .isForShot_ = true,
138 };
139 virtualScreenId_ = ScreenManager::GetInstance().CreateVirtualScreen(option);
140 vector<ScreenId> mirrorIds;
141 mirrorIds.push_back(virtualScreenId_);
142 ScreenId screenGroupId = static_cast<ScreenId>(1);
143 auto ret = ScreenManager::GetInstance().MakeMirror(mainScreenId_, mirrorIds, screenGroupId);
144 if (ret != DMError::DM_OK) {
145 LOG_E("Make mirror screen for default screen failed");
146 return false;
147 }
148 return true;
149 }
150
Pause()151 bool ScreenCopy::Pause()
152 {
153 if (!workable_ || virtualScreenId_ == SCREEN_ID_INVALID) {
154 return true;
155 }
156 vector<ScreenId> mirrorIds;
157 mirrorIds.push_back(virtualScreenId_);
158 auto err0 = ScreenManager::GetInstance().StopMirror(mirrorIds);
159 auto err1 = ScreenManager::GetInstance().DestroyVirtualScreen(virtualScreenId_);
160 virtualScreenId_ = SCREEN_ID_INVALID;
161 if (err0 == DMError::DM_OK && err1 == DMError::DM_OK) {
162 return true;
163 } else {
164 LOG_E("Pause screenCopy failed, stopMirrorErr=%{public}d, destroyVirScrErr=%{public}d", err0, err1);
165 return false;
166 }
167 }
168
HandleConsumerBuffer()169 void ScreenCopy::HandleConsumerBuffer()
170 {
171 LOG_D("ScreenCopy::HandleConsumerBuffer");
172 sptr<SurfaceBuffer> consumerSuffer = nullptr;
173 int32_t fenceFd = -1;
174 int64_t tick = 0;
175 int32_t timeout = 3000;
176 Rect rect;
177 auto aquireRet = consumerSurface_->AcquireBuffer(consumerSuffer, fenceFd, tick, rect);
178 if (aquireRet != 0 || consumerSuffer == nullptr) {
179 LOG_E("AcquireBuffer failed");
180 return;
181 }
182 int pollRet = -1;
183 struct pollfd pfd = {0};
184 pfd.fd = fenceFd;
185 pfd.events = POLLIN;
186 do {
187 pollRet = poll(&pfd, 1, timeout);
188 } while (pollRet == -1 && (errno == EINTR || errno == EAGAIN));
189
190 if (pollRet == 0) {
191 pollRet = -1;
192 errno = ETIME;
193 } else if (pollRet > 0) {
194 pollRet = 0;
195 if (pfd.revents & (POLLERR | POLLNVAL)) {
196 pollRet = -1;
197 errno = EINVAL;
198 }
199 }
200 if (pollRet < 0) {
201 LOG_E("Poll on fenceFd failed: %{public}s", strerror(errno));
202 return;
203 }
204 LOG_D("start ConsumerFrameBuffer");
205 ConsumeFrameBuffer(consumerSuffer);
206 LOG_D("end ConsumeFrameBuffer");
207 if (consumerSurface_->ReleaseBuffer(consumerSuffer, -1) != 0) {
208 LOG_E("ReleaseBuffer failed");
209 return;
210 }
211 }
212
ConvertRGBA2RGB(const uint8_t * input,uint8_t * output,int32_t pixelNum)213 static void ConvertRGBA2RGB(const uint8_t* input, uint8_t* output, int32_t pixelNum)
214 {
215 constexpr uint8_t BLUE_INDEX = 0;
216 constexpr uint8_t GREEN_INDEX = 1;
217 constexpr uint8_t RED_INDEX = 2;
218 constexpr uint32_t RGB_PIXEL_BYTES = 3;
219 constexpr uint8_t SHIFT_8_BIT = 8;
220 constexpr uint8_t SHIFT_16_BIT = 16;
221 constexpr uint32_t RGBA_MASK_BLUE = 0x000000FF;
222 constexpr uint32_t RGBA_MASK_GREEN = 0x0000FF00;
223 constexpr uint32_t RGBA_MASK_RED = 0x00FF0000;
224 DCHECK(input != nullptr && output != nullptr && pixelNum > 0);
225 auto pRgba = reinterpret_cast<const uint32_t*>(input);
226 for (int32_t index = 0; index < pixelNum; index++) {
227 output[index * RGB_PIXEL_BYTES + RED_INDEX] = (pRgba[index] & RGBA_MASK_RED) >> SHIFT_16_BIT;
228 output[index * RGB_PIXEL_BYTES + GREEN_INDEX] = (pRgba[index] & RGBA_MASK_GREEN) >> SHIFT_8_BIT;
229 output[index * RGB_PIXEL_BYTES + BLUE_INDEX] = (pRgba[index] & RGBA_MASK_BLUE);
230 }
231 }
232
233 struct MissionErrorMgr : public jpeg_error_mgr {
234 jmp_buf setjmp_buffer;
235 };
236
AdaptJpegSize(jpeg_compress_struct & jpeg,uint32_t width,uint32_t height)237 static void AdaptJpegSize(jpeg_compress_struct &jpeg, uint32_t width, uint32_t height)
238 {
239 constexpr int32_t VIRTUAL_SCREEN_SCALE = 32;
240 if (width % VIRTUAL_SCREEN_SCALE == 0) {
241 jpeg.image_width = width;
242 } else {
243 LOG_D("The width need to be adapted!");
244 jpeg.image_width = ceil((double)width / (double)VIRTUAL_SCREEN_SCALE) * VIRTUAL_SCREEN_SCALE;
245 }
246 jpeg.image_height = height;
247 }
248
ConsumeFrameBuffer(const sptr<SurfaceBuffer> & buf)249 void ScreenCopy::ConsumeFrameBuffer(const sptr<SurfaceBuffer> &buf)
250 {
251 auto bufHdl = buf->GetBufferHandle();
252 if (bufHdl == nullptr) {
253 LOG_E("GetBufferHandle failed");
254 return;
255 }
256 if (g_screenCopyHandler == nullptr) {
257 LOG_W("Consumer handler is nullptr, ignore this frame");
258 return;
259 }
260 LOG_I("ConsumeFrameBuffer_BeginEncodeFrameToJPEG");
261 auto width = static_cast<uint32_t>(bufHdl->width);
262 auto height = static_cast<uint32_t>(bufHdl->height);
263 auto stride = static_cast<uint32_t>(bufHdl->stride);
264 auto data = (uint8_t *)buf->GetVirAddr();
265 constexpr int32_t RGBA_PIXEL_BYTES = 4;
266 constexpr int32_t RGB_PIXEL_BYTES = 3;
267 int32_t rgbSize = stride * height * RGB_PIXEL_BYTES / RGBA_PIXEL_BYTES;
268 auto rgb = new uint8_t[rgbSize];
269 ConvertRGBA2RGB(data, rgb, rgbSize / RGB_PIXEL_BYTES);
270
271 jpeg_compress_struct jpeg;
272 MissionErrorMgr jerr;
273 jpeg.err = jpeg_std_error(&jerr);
274 jpeg_create_compress(&jpeg);
275 AdaptJpegSize(jpeg, width, height);
276 jpeg.input_components = RGB_PIXEL_BYTES;
277 jpeg.in_color_space = JCS_RGB;
278 jpeg_set_defaults(&jpeg);
279 constexpr int32_t COMPRESS_QUALITY = 75;
280 jpeg_set_quality(&jpeg, COMPRESS_QUALITY, 1);
281 uint8_t *imgBuf = nullptr;
282 unsigned long imgSize = 0;
283 jpeg_mem_dest(&jpeg, &imgBuf, &imgSize);
284 jpeg_start_compress(&jpeg, 1);
285 JSAMPROW rowPointer[1];
286 for (uint32_t rowIndex = 0; rowIndex < jpeg.image_height; rowIndex++) {
287 rowPointer[0] = const_cast<uint8_t *>(rgb + rowIndex * jpeg.image_width * RGB_PIXEL_BYTES);
288 (void)jpeg_write_scanlines(&jpeg, rowPointer, 1);
289 }
290 jpeg_finish_compress(&jpeg);
291 jpeg_destroy_compress(&jpeg);
292 free(rgb);
293 LOG_I("ConsumeFrameBuffer_EndEncodeFrameToJPEG");
294 if (g_screenCopyHandler != nullptr) {
295 g_screenCopyHandler(imgBuf, imgSize);
296 } else {
297 free(imgBuf);
298 }
299 }
300
StartScreenCopy(float scale,ScreenCopyHandler handler)301 bool StartScreenCopy(float scale, ScreenCopyHandler handler)
302 {
303 if (scale <= 0 || scale > 1 || handler == nullptr) {
304 LOG_E("Illegal arguments");
305 return false;
306 }
307 g_screenCopyHandler = handler;
308 if (g_screenCopy == nullptr) {
309 g_screenCopy = make_unique<ScreenCopy>();
310 g_screenCopy->Setup();
311 }
312 return g_screenCopy->Run();
313 }
314
StopScreenCopy()315 void StopScreenCopy()
316 {
317 if (g_screenCopy != nullptr) {
318 g_screenCopy->Pause();
319 }
320 }
321 }