1 /*
2 * 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 #include "host/frontend/webrtc/screenshot_handler.h"
18
19 #include <filesystem>
20 #include <fstream>
21
22 #include <SkData.h>
23 #include <SkImage.h>
24 #include <SkJpegEncoder.h>
25 #include <SkPngEncoder.h>
26 #include <SkRefCnt.h>
27 #include <SkStream.h>
28 #include <SkWebpEncoder.h>
29 #include <libyuv.h>
30
31 namespace cuttlefish {
32 namespace {
33
GetSkImage(const webrtc_streaming::VideoFrameBuffer & frame)34 Result<sk_sp<SkImage>> GetSkImage(
35 const webrtc_streaming::VideoFrameBuffer& frame) {
36 const int w = frame.width();
37 const int h = frame.height();
38
39 sk_sp<SkData> rgba_data = SkData::MakeUninitialized(w * h * 4);
40 const int rgba_stride = w * 4;
41
42 int ret = libyuv::I420ToABGR(
43 frame.DataY(), frame.StrideY(), //
44 frame.DataU(), frame.StrideU(), //
45 frame.DataV(), frame.StrideV(), //
46 reinterpret_cast<uint8_t*>(rgba_data->writable_data()), rgba_stride, //
47 w, h);
48 CF_EXPECT_EQ(ret, 0, "Failed to convert input frame to RGBA.");
49
50 const SkImageInfo& image_info =
51 SkImageInfo::Make(w, h, kRGBA_8888_SkColorType, kOpaque_SkAlphaType);
52
53 sk_sp<SkImage> image =
54 SkImages::RasterFromData(image_info, rgba_data, rgba_stride);
55 CF_EXPECT(image != nullptr, "Failed to raster RGBA data.");
56
57 return image;
58 }
59
60 } // namespace
61
Screenshot(std::uint32_t display_number,const std::string & screenshot_path)62 Result<void> ScreenshotHandler::Screenshot(std::uint32_t display_number,
63 const std::string& screenshot_path) {
64 SharedFrameFuture frame_future;
65 {
66 std::lock_guard<std::mutex> lock(pending_screenshot_displays_mutex_);
67
68 auto [it, inserted] = pending_screenshot_displays_.emplace(
69 display_number, SharedFramePromise{});
70 if (!inserted) {
71 return CF_ERRF("Screenshot already pending for display {}",
72 display_number);
73 }
74
75 frame_future = it->second.get_future().share();
76 }
77
78 static constexpr const int kScreenshotTimeoutSeconds = 5;
79 auto result =
80 frame_future.wait_for(std::chrono::seconds(kScreenshotTimeoutSeconds));
81 CF_EXPECT(result == std::future_status::ready,
82 "Failed to get screenshot from webrtc display handler within "
83 << kScreenshotTimeoutSeconds << " seconds.");
84
85 SharedFrame frame = frame_future.get();
86
87 sk_sp<SkImage> screenshot_image =
88 CF_EXPECT(GetSkImage(*frame), "Failed to get skia image from raw frame.");
89
90 sk_sp<SkData> screenshot_data;
91 if (screenshot_path.ends_with(".jpg")) {
92 screenshot_data =
93 SkJpegEncoder::Encode(nullptr, screenshot_image.get(), {});
94 CF_EXPECT(screenshot_data != nullptr, "Failed to encode to JPEG.");
95 } else if (screenshot_path.ends_with(".png")) {
96 screenshot_data = SkPngEncoder::Encode(nullptr, screenshot_image.get(), {});
97 CF_EXPECT(screenshot_data != nullptr, "Failed to encode to PNG.");
98 } else if (screenshot_path.ends_with(".webp")) {
99 screenshot_data =
100 SkWebpEncoder::Encode(nullptr, screenshot_image.get(), {});
101 CF_EXPECT(screenshot_data != nullptr, "Failed to encode to WEBP.");
102 } else {
103 return CF_ERR("Unsupport file format: " << screenshot_path);
104 }
105
106 SkFILEWStream screenshot_file(screenshot_path.c_str());
107 CF_EXPECT(screenshot_file.isValid(),
108 "Failed to open " << screenshot_path << " for writing.");
109
110 CF_EXPECT(
111 screenshot_file.write(screenshot_data->data(), screenshot_data->size()),
112 "Failed to fully write png content to " << screenshot_path << ".");
113
114 return {};
115 }
116
OnFrame(std::uint32_t display_number,SharedFrame & frame)117 void ScreenshotHandler::OnFrame(std::uint32_t display_number,
118 SharedFrame& frame) {
119 std::lock_guard<std::mutex> lock(pending_screenshot_displays_mutex_);
120
121 auto pending_screenshot_it =
122 pending_screenshot_displays_.find(display_number);
123 if (pending_screenshot_it == pending_screenshot_displays_.end()) {
124 return;
125 }
126 SharedFramePromise& frame_promise = pending_screenshot_it->second;
127
128 frame_promise.set_value(frame);
129
130 pending_screenshot_displays_.erase(pending_screenshot_it);
131 }
132
133 } // namespace cuttlefish
134