1/*
2 * Copyright (c) 2021-2022 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#include "snapshot_utils.h"
16
17#include <cerrno>
18#include <climits>
19#include <cstdio>
20#include <cstdlib>
21#include <ctime>
22#include <getopt.h>
23#include <hitrace_meter.h>
24#include <image_type.h>
25#include <iostream>
26#include <ostream>
27#include <csetjmp>
28#include <pixel_map.h>
29#include <securec.h>
30#include <string>
31#include <sys/time.h>
32
33#include "image_packer.h"
34#include "jpeglib.h"
35
36using namespace OHOS::Rosen;
37
38namespace OHOS {
39constexpr int MAX_TIME_STR_LEN = 40;
40constexpr int YEAR_SINCE = 1900;
41constexpr int32_t RGB565_PIXEL_BYTES = 2;
42constexpr int32_t RGB888_PIXEL_BYTES = 3;
43constexpr int32_t RGBA8888_PIXEL_BYTES = 4;
44constexpr uint8_t B_INDEX = 0;
45constexpr uint8_t G_INDEX = 1;
46constexpr uint8_t R_INDEX = 2;
47constexpr uint8_t SHIFT_2_BIT = 2;
48constexpr uint8_t SHIFT_3_BIT = 3;
49constexpr uint8_t SHIFT_5_BIT = 5;
50constexpr uint8_t SHIFT_8_BIT = 8;
51constexpr uint8_t SHIFT_11_BIT = 11;
52constexpr uint8_t SHIFT_16_BIT = 16;
53
54constexpr uint16_t RGB565_MASK_BLUE = 0xF800;
55constexpr uint16_t RGB565_MASK_GREEN = 0x07E0;
56constexpr uint16_t RGB565_MASK_RED = 0x001F;
57constexpr uint32_t RGBA8888_MASK_BLUE = 0x000000FF;
58constexpr uint32_t RGBA8888_MASK_GREEN = 0x0000FF00;
59constexpr uint32_t RGBA8888_MASK_RED = 0x00FF0000;
60
61constexpr uint8_t PNG_PACKER_QUALITY = 100;
62constexpr uint8_t PACKER_QUALITY = 75;
63constexpr uint32_t PACKER_SUCCESS = 0;
64struct MissionErrorMgr : public jpeg_error_mgr {
65    jmp_buf environment;
66};
67
68void mission_error_exit(j_common_ptr cinfo)
69{
70    if (cinfo == nullptr || cinfo->err == nullptr) {
71        std::cout << __func__ << ": param is invalid." << std::endl;
72        return;
73    }
74    auto err = reinterpret_cast<MissionErrorMgr*>(cinfo->err);
75    longjmp(err->environment, 1);
76}
77
78const char *VALID_SNAPSHOT_PATH = "/data/local/tmp";
79const char *DEFAULT_SNAPSHOT_PREFIX = "/snapshot";
80const char *VALID_SNAPSHOT_SUFFIX = ".jpeg";
81const char *VALID_SNAPSHOT_PNG_SUFFIX = ".png";
82
83void SnapShotUtils::PrintUsage(const std::string& cmdLine)
84{
85    std::cout << "usage: " << cmdLine.c_str() <<
86        " [-i displayId] [-f output_file] [-w width] [-h height] [-t type] [-m]" << std::endl;
87}
88
89std::string SnapShotUtils::GenerateFileName(std::string fileType, int offset)
90{
91    timeval tv;
92    std::string fileName = VALID_SNAPSHOT_PATH;
93
94    fileName += DEFAULT_SNAPSHOT_PREFIX;
95    if (gettimeofday(&tv, nullptr) == 0) {
96        tv.tv_sec += offset; // add offset second
97        struct tm *tmVal = localtime(&tv.tv_sec);
98        if (tmVal != nullptr) {
99            char timeStr[MAX_TIME_STR_LEN] = { 0 };
100            snprintf_s(timeStr, sizeof(timeStr), sizeof(timeStr) - 1,
101                "_%04d-%02d-%02d_%02d-%02d-%02d",
102                tmVal->tm_year + YEAR_SINCE, tmVal->tm_mon + 1, tmVal->tm_mday,
103                tmVal->tm_hour, tmVal->tm_min, tmVal->tm_sec);
104            fileName += timeStr;
105        }
106    }
107    fileName += (fileType == "png") ? VALID_SNAPSHOT_PNG_SUFFIX : VALID_SNAPSHOT_SUFFIX;
108    return fileName;
109}
110
111bool SnapShotUtils::CheckFileNameValid(const std::string& fileName, std::string fileType)
112{
113    std::cout << "fileType: " << fileType << std::endl;
114    size_t fileMinLength = (fileType == "png") ? strlen(VALID_SNAPSHOT_PNG_SUFFIX) : strlen(VALID_SNAPSHOT_SUFFIX);
115    if (fileName.length() <= fileMinLength) {
116        std::cout << "error: fileName " << fileName.c_str() << " invalid, file length too short!" << std::endl;
117        return false;
118    }
119    // check file path
120    std::string fileDir = fileName;
121    auto pos = fileDir.find_last_of("/");
122    if (pos != std::string::npos) {
123        fileDir.erase(pos + 1);
124    } else {
125        fileDir = "."; // current work dir
126    }
127    char resolvedPath[PATH_MAX] = { 0 };
128    char *realPath = realpath(fileDir.c_str(), resolvedPath);
129    if (realPath == nullptr) {
130        std::cout << "error: fileName " << fileName.c_str() << " invalid, realpath nullptr!" << std::endl;
131        return false;
132    }
133    if (strncmp(realPath, VALID_SNAPSHOT_PATH, strlen(VALID_SNAPSHOT_PATH)) != 0) {
134        std::cout << "error: fileName " << fileName.c_str() << " invalid, realpath "
135            << realPath << " must dump at dir: " << VALID_SNAPSHOT_PATH << std::endl;
136        return false;
137    }
138
139    // check file suffix
140    const char *fileNameSuffix = fileName.c_str() + (fileName.length() - fileMinLength);
141    const char *fileSuffix = (fileType == "png") ? VALID_SNAPSHOT_PNG_SUFFIX : VALID_SNAPSHOT_SUFFIX;
142    if (strncmp(fileNameSuffix, fileSuffix, fileMinLength) == 0) {
143        return true; // valid suffix
144    }
145    std::cout << "error: fileName " << fileName.c_str() << " invalid, suffix must be " << fileSuffix << std::endl;
146    return false;
147}
148
149bool SnapShotUtils::CheckWHValid(int32_t param)
150{
151    return (param > 0) && (param <= DisplayManager::MAX_RESOLUTION_SIZE_SCREENSHOT);
152}
153
154bool SnapShotUtils::CheckWidthAndHeightValid(int32_t w, int32_t h)
155{
156    return CheckWHValid(w) && CheckWHValid(h);
157}
158
159bool SnapShotUtils::CheckParamValid(const WriteToJpegParam& param)
160{
161    switch (param.format) {
162        case Media::PixelFormat::RGBA_8888:
163            if (param.stride != param.width * RGBA8888_PIXEL_BYTES) {
164                return false;
165            }
166            break;
167        case Media::PixelFormat::RGB_565:
168            if (param.stride != param.width * RGB565_PIXEL_BYTES) {
169                return false;
170            }
171            break;
172        case Media::PixelFormat::RGB_888:
173            if (param.stride != param.width * RGB888_PIXEL_BYTES) {
174                return false;
175            }
176            break;
177        default:
178            std::cout << __func__ << ": unsupported pixel format: " <<
179                static_cast<uint32_t>(param.format) << std::endl;
180            return false;
181    }
182    if (!CheckWidthAndHeightValid(param.width, param.height)) {
183        return false;
184    }
185    if (param.data == nullptr) {
186        return false;
187    }
188    return true;
189}
190
191bool SnapShotUtils::RGBA8888ToRGB888(const uint8_t* rgba8888Buf, uint8_t* rgb888Buf, int32_t size)
192{
193    if (rgba8888Buf == nullptr || rgb888Buf == nullptr || size <= 0) {
194        std::cout << __func__ << ": params are invalid." << std::endl;
195        return false;
196    }
197    const uint32_t* rgba8888 = reinterpret_cast<const uint32_t*>(rgba8888Buf);
198    for (int32_t i = 0; i < size; i++) {
199        rgb888Buf[i * RGB888_PIXEL_BYTES + R_INDEX] = (rgba8888[i] & RGBA8888_MASK_RED) >> SHIFT_16_BIT;
200        rgb888Buf[i * RGB888_PIXEL_BYTES + G_INDEX] = (rgba8888[i] & RGBA8888_MASK_GREEN) >> SHIFT_8_BIT;
201        rgb888Buf[i * RGB888_PIXEL_BYTES + B_INDEX] = rgba8888[i] & RGBA8888_MASK_BLUE;
202    }
203    return true;
204}
205
206bool SnapShotUtils::RGB565ToRGB888(const uint8_t* rgb565Buf, uint8_t* rgb888Buf, int32_t size)
207{
208    if (rgb565Buf == nullptr || rgb888Buf == nullptr || size <= 0) {
209        std::cout << __func__ << ": params are invalid." << std::endl;
210        return false;
211    }
212    const uint16_t* rgb565 = reinterpret_cast<const uint16_t*>(rgb565Buf);
213    for (int32_t i = 0; i < size; i++) {
214        rgb888Buf[i * RGB888_PIXEL_BYTES + R_INDEX] = (rgb565[i] & RGB565_MASK_RED);
215        rgb888Buf[i * RGB888_PIXEL_BYTES + G_INDEX] = (rgb565[i] & RGB565_MASK_GREEN) >> SHIFT_5_BIT;
216        rgb888Buf[i * RGB888_PIXEL_BYTES + B_INDEX] = (rgb565[i] & RGB565_MASK_BLUE) >> SHIFT_11_BIT;
217        rgb888Buf[i * RGB888_PIXEL_BYTES + R_INDEX] <<= SHIFT_3_BIT;
218        rgb888Buf[i * RGB888_PIXEL_BYTES + G_INDEX] <<= SHIFT_2_BIT;
219        rgb888Buf[i * RGB888_PIXEL_BYTES + B_INDEX] <<= SHIFT_3_BIT;
220    }
221    return true;
222}
223
224// The method will NOT release file.
225bool SnapShotUtils::WriteRgb888ToJpeg(FILE* file, uint32_t width, uint32_t height, const uint8_t* data)
226{
227    if (data == nullptr) {
228        std::cout << "error: data error, nullptr!" << std::endl;
229        return false;
230    }
231
232    if (file == nullptr) {
233        std::cout << "error: file is null" << std::endl;
234        return false;
235    }
236
237    struct jpeg_compress_struct jpeg;
238    struct MissionErrorMgr jerr;
239    jpeg.err = jpeg_std_error(&jerr);
240    jerr.error_exit = mission_error_exit;
241    if (setjmp(jerr.environment)) {
242        jpeg_destroy_compress(&jpeg);
243        std::cout << "error: lib jpeg exit with error!" << std::endl;
244        return false;
245    }
246
247    jpeg_create_compress(&jpeg);
248    jpeg.image_width = width;
249    jpeg.image_height = height;
250    jpeg.input_components = RGB888_PIXEL_BYTES;
251    jpeg.in_color_space = JCS_RGB;
252    jpeg_set_defaults(&jpeg);
253
254    constexpr int32_t quality = 75;
255    jpeg_set_quality(&jpeg, quality, TRUE);
256
257    jpeg_stdio_dest(&jpeg, file);
258    jpeg_start_compress(&jpeg, TRUE);
259    JSAMPROW rowPointer[1];
260    for (uint32_t i = 0; i < jpeg.image_height; i++) {
261        rowPointer[0] = const_cast<uint8_t *>(data + i * jpeg.image_width * RGB888_PIXEL_BYTES);
262        (void)jpeg_write_scanlines(&jpeg, rowPointer, 1);
263    }
264
265    jpeg_finish_compress(&jpeg);
266    jpeg_destroy_compress(&jpeg);
267    return true;
268}
269
270bool SnapShotUtils::WriteToJpeg(const std::string& fileName, const WriteToJpegParam& param)
271{
272    bool ret = false;
273    if (!CheckFileNameValid(fileName)) {
274        return ret;
275    }
276    if (!CheckParamValid(param)) {
277        std::cout << "error: invalid param." << std::endl;
278        return ret;
279    }
280    HITRACE_METER_FMT(HITRACE_TAG_WINDOW_MANAGER, "snapshot:WriteToJpeg(%s)", fileName.c_str());
281
282    FILE *file = fopen(fileName.c_str(), "wb");
283    if (file == nullptr) {
284        std::cout << "error: open file [" << fileName.c_str() << "] error, " << errno << "!" << std::endl;
285        return ret;
286    }
287    std::cout << "snapshot: pixel format is: " << static_cast<uint32_t>(param.format) << std::endl;
288    if (param.format == Media::PixelFormat::RGBA_8888) {
289        int32_t rgb888Size = param.stride * param.height * RGB888_PIXEL_BYTES / RGBA8888_PIXEL_BYTES;
290        uint8_t *rgb888 = new uint8_t[rgb888Size];
291        ret = RGBA8888ToRGB888(param.data, rgb888, rgb888Size / RGB888_PIXEL_BYTES);
292        if (ret) {
293            std::cout << "snapshot: convert rgba8888 to rgb888 successfully." << std::endl;
294            ret = WriteRgb888ToJpeg(file, param.width, param.height, rgb888);
295        }
296        delete[] rgb888;
297    } else if (param.format == Media::PixelFormat::RGB_565) {
298        int32_t rgb888Size = param.stride * param.height * RGB888_PIXEL_BYTES / RGB565_PIXEL_BYTES;
299        uint8_t *rgb888 = new uint8_t[rgb888Size];
300        ret = RGB565ToRGB888(param.data, rgb888, rgb888Size / RGB888_PIXEL_BYTES);
301        if (ret) {
302            std::cout << "snapshot: convert rgb565 to rgb888 successfully." << std::endl;
303            ret = WriteRgb888ToJpeg(file, param.width, param.height, rgb888);
304        }
305        delete[] rgb888;
306    } else if (param.format == Media::PixelFormat::RGB_888) {
307        ret = WriteRgb888ToJpeg(file, param.width, param.height, param.data);
308    } else {
309        std::cout << "snapshot: invalid pixel format." << std::endl;
310    }
311    if (fclose(file) != 0) {
312        std::cout << "error: close file failed!" << std::endl;
313        ret = false;
314    }
315    return ret;
316}
317
318bool SnapShotUtils::WriteToJpeg(int fd, const WriteToJpegParam& param)
319{
320    bool ret = false;
321    if (!CheckParamValid(param)) {
322        std::cout << "error: invalid param." << std::endl;
323        return ret;
324    }
325
326    FILE *file = fdopen(fd, "wb");
327    if (file == nullptr) {
328        return ret;
329    }
330    std::cout << "snapshot: pixel format is: " << static_cast<uint32_t>(param.format) << std::endl;
331    if (param.format == Media::PixelFormat::RGBA_8888) {
332        int32_t rgb888Size = param.stride * param.height * RGB888_PIXEL_BYTES / RGBA8888_PIXEL_BYTES;
333        uint8_t *rgb888 = new uint8_t[rgb888Size];
334        ret = RGBA8888ToRGB888(param.data, rgb888, rgb888Size / RGB888_PIXEL_BYTES);
335        if (ret) {
336            std::cout << "snapshot: convert rgba8888 to rgb888 successfully." << std::endl;
337            ret = WriteRgb888ToJpeg(file, param.width, param.height, rgb888);
338        }
339        delete[] rgb888;
340    } else if (param.format == Media::PixelFormat::RGB_565) {
341        int32_t rgb888Size = param.stride * param.height * RGB888_PIXEL_BYTES / RGB565_PIXEL_BYTES;
342        uint8_t *rgb888 = new uint8_t[rgb888Size];
343        ret = RGB565ToRGB888(param.data, rgb888, rgb888Size / RGB888_PIXEL_BYTES);
344        if (ret) {
345            std::cout << "snapshot: convert rgb565 to rgb888 successfully." << std::endl;
346            ret = WriteRgb888ToJpeg(file, param.width, param.height, rgb888);
347        }
348        delete[] rgb888;
349    } else if (param.format == Media::PixelFormat::RGB_888) {
350        ret = WriteRgb888ToJpeg(file, param.width, param.height, param.data);
351    } else {
352        std::cout << "snapshot: invalid pixel format." << std::endl;
353    }
354    if (fclose(file) != 0) {
355        std::cout << "error: close file failed!" << std::endl;
356        ret = false;
357    }
358    return ret;
359}
360
361bool SnapShotUtils::SaveSnapShot(const std::string& fileName, Media::PixelMap& pixelMap, std::string fileType)
362{
363    OHOS::Media::ImagePacker imagePacker;
364    OHOS::Media::PackOption option;
365    option.format = (fileType == "png") ? "image/png" : "image/jpeg";
366    option.quality = (fileType == "png") ? PNG_PACKER_QUALITY : PACKER_QUALITY;
367    option.numberHint = 1;
368    std::set<std::string> formats;
369    auto ret = imagePacker.GetSupportedFormats(formats);
370    if (ret) {
371        std::cout << "error: get supported formats error" << std::endl;
372        return false;
373    }
374
375    imagePacker.StartPacking(fileName, option);
376    imagePacker.AddImage(pixelMap);
377    int64_t packedSize = 0;
378    uint32_t res = imagePacker.FinalizePacking(packedSize);
379    if (res != PACKER_SUCCESS) {
380        std::cout << "error:FinalizePacking error" << std::endl;
381        return false;
382    }
383    return true;
384}
385
386bool SnapShotUtils::WriteToJpegWithPixelMap(const std::string& fileName, Media::PixelMap& pixelMap)
387{
388    if (pixelMap.GetAllocatorType() == Media::AllocatorType::DMA_ALLOC) {
389        return SaveSnapShot(fileName, pixelMap);
390    }
391    WriteToJpegParam param;
392    param.width = static_cast<uint32_t>(pixelMap.GetWidth());
393    param.height = static_cast<uint32_t>(pixelMap.GetHeight());
394    param.data = pixelMap.GetPixels();
395    param.stride = static_cast<uint32_t>(pixelMap.GetRowBytes());
396    param.format = pixelMap.GetPixelFormat();
397    return SnapShotUtils::WriteToJpeg(fileName, param);
398}
399
400bool SnapShotUtils::WriteToJpegWithPixelMap(int fd, Media::PixelMap& pixelMap)
401{
402    WriteToJpegParam param;
403    param.width = static_cast<uint32_t>(pixelMap.GetWidth());
404    param.height = static_cast<uint32_t>(pixelMap.GetHeight());
405    param.data = pixelMap.GetPixels();
406    param.stride = static_cast<uint32_t>(pixelMap.GetRowBytes());
407    param.format = pixelMap.GetPixelFormat();
408    return SnapShotUtils::WriteToJpeg(fd, param);
409}
410
411bool SnapShotUtils::ProcessDisplayId(Rosen::DisplayId& displayId, bool isDisplayIdSet)
412{
413    if (!isDisplayIdSet) {
414        displayId = DisplayManager::GetInstance().GetDefaultDisplayId();
415    } else {
416        bool validFlag = false;
417        auto displayIds = DisplayManager::GetInstance().GetAllDisplayIds();
418        for (auto id : displayIds) {
419            if (displayId == id) {
420                validFlag = true;
421                break;
422            }
423        }
424        if (!validFlag) {
425            std::cout << "error: displayId " << static_cast<int64_t>(displayId) << " invalid!" << std::endl;
426            std::cout << "tips: supported displayIds:" << std::endl;
427            for (auto dispId : displayIds) {
428                std::cout << "\t" << dispId << std::endl;
429            }
430            return false;
431        }
432    }
433    return true;
434}
435
436bool SnapShotUtils::ProcessArgs(int argc, char* const argv[], CmdArguments& cmdArguments)
437{
438    int opt = 0;
439    const struct option longOption[] = {
440        { "id", required_argument, nullptr, 'i' },
441        { "width", required_argument, nullptr, 'w' },
442        { "height", required_argument, nullptr, 'h' },
443        { "file", required_argument, nullptr, 'f' },
444        { "type", required_argument, nullptr, 't' },
445        { "help", required_argument, nullptr, 'm' },
446        { nullptr, 0, nullptr, 0 }
447    };
448    while ((opt = getopt_long(argc, argv, "i:w:h:f:t:m", longOption, nullptr)) != -1) {
449        switch (opt) {
450            case 'i': // display id
451                cmdArguments.displayId = static_cast<DisplayId>(atoll(optarg));
452                cmdArguments.isDisplayIdSet = true;
453                break;
454            case 'w': // output width
455                cmdArguments.width = atoi(optarg);
456                cmdArguments.isWidthSet = true;
457                break;
458            case 'h': // output height
459                cmdArguments.height = atoi(optarg);
460                cmdArguments.isHeightSet = true;
461                break;
462            case 'f': // output file name
463                cmdArguments.fileName = optarg;
464                break;
465            case 't': // output file type
466                cmdArguments.fileType = optarg;
467                break;
468            case 'm': // help
469            default:
470                SnapShotUtils::PrintUsage(argv[0]);
471                return false;
472        }
473    }
474
475    if (!ProcessDisplayId(cmdArguments.displayId, cmdArguments.isDisplayIdSet)) {
476        return false;
477    }
478
479    if (cmdArguments.fileName == "") {
480        cmdArguments.fileName = GenerateFileName(cmdArguments.fileType);
481        std::cout << "process: set filename to " << cmdArguments.fileName.c_str() << std::endl;
482    }
483
484    // check fileName
485    if (!SnapShotUtils::CheckFileNameValid(cmdArguments.fileName, cmdArguments.fileType)) {
486        std::cout << "error: filename " << cmdArguments.fileName.c_str() << " invalid!" << std::endl;
487        return false;
488    }
489    return true;
490}
491}
492