1 /*
2  * Copyright (c) 2023-2024 Shenzhen Kaihong Digital Industry Development 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 "screen_capture_consumer.h"
17 #include "common/common_macro.h"
18 #include "common/reflect_registration.h"
19 #include "common/sharing_log.h"
20 #include "screen_capture_def.h"
21 
22 namespace OHOS {
23 namespace Sharing {
24 static std::chrono::system_clock::time_point start = std::chrono::system_clock::now();
OnFrame(const Frame::Ptr &frame)25 void ScreenCaptureConsumer::AudioEncoderReceiver::OnFrame(const Frame::Ptr &frame)
26 {
27     auto parent = parent_.lock();
28     if (parent && !parent->listener_.expired()) {
29         auto listener = parent->listener_.lock();
30         auto dispatcher = listener->GetDispatcher();
31         if (dispatcher) {
32             auto mediaData = dispatcher->RequestDataBuffer(MEDIA_TYPE_AUDIO, frame->Size());
33             mediaData->mediaType = MEDIA_TYPE_AUDIO;
34             mediaData->codecId = frame->GetCodecId();
35             mediaData->isRaw = false;
36             mediaData->keyFrame = false;
37             mediaData->pts = frame->Pts();
38             MEDIA_LOGD("recv audio encode data: codec:%{public}d size: %{public}d & put into dispatcher: %{public}u.",
39                        mediaData->codecId, frame->Size(), dispatcher->GetDispatcherId());
40             mediaData->buff = move(frame);
41             dispatcher->InputData(mediaData);
42         }
43     }
44 }
45 
OnFrameBufferUsed()46 void ScreenCaptureConsumer::OnFrameBufferUsed()
47 {
48     ReleaseScreenBuffer();
49 }
50 
HandleSpsFrame(BufferDispatcher::Ptr dispatcher, const Frame::Ptr &frame)51 void ScreenCaptureConsumer::HandleSpsFrame(BufferDispatcher::Ptr dispatcher, const Frame::Ptr &frame)
52 {
53     RETURN_IF_NULL(dispatcher);
54     auto spsOld = dispatcher->GetSPS();
55     if (spsOld != nullptr && spsOld->buff != nullptr) {
56         return;
57     }
58     // set sps into buffer dispathcer
59     auto sps = std::make_shared<MediaData>();
60     sps->buff = move(frame);
61     sps->mediaType = MEDIA_TYPE_VIDEO;
62     dispatcher->SetSpsNalu(sps);
63 }
64 
HandlePpsFrame(BufferDispatcher::Ptr dispatcher, const Frame::Ptr &frame)65 void ScreenCaptureConsumer::HandlePpsFrame(BufferDispatcher::Ptr dispatcher, const Frame::Ptr &frame)
66 {
67     RETURN_IF_NULL(dispatcher);
68     auto ppsOld = dispatcher->GetPPS();
69     if (ppsOld != nullptr && ppsOld->buff != nullptr) {
70         return;
71     }
72     // set pps into buffer dispather
73     auto pps = std::make_shared<MediaData>();
74     pps->buff = move(frame);
75     pps->mediaType = MEDIA_TYPE_VIDEO;
76     dispatcher->SetPpsNalu(pps);
77 }
78 
OnFrame(const Frame::Ptr &frame, FRAME_TYPE frameType, bool keyFrame)79 void ScreenCaptureConsumer::OnFrame(const Frame::Ptr &frame, FRAME_TYPE frameType, bool keyFrame)
80 {
81     SHARING_LOGD("trace.");
82     if (frame == nullptr) {
83         SHARING_LOGE("Screen Capture encodedBuf is null!");
84         return;
85     }
86 
87     if (listener_.expired() || IsPaused()) {
88         return;
89     }
90 
91     auto listener = listener_.lock();
92     BufferDispatcher::Ptr dispatcher = listener->GetDispatcher();
93     if (dispatcher) {
94         uint32_t len = frame->Size();
95         switch (frameType) {
96             case SPS_FRAME: {
97                 SHARING_LOGD("get sps frame, size:%{public}u.", len);
98                 HandleSpsFrame(dispatcher, frame);
99                 break;
100             }
101             case PPS_FRAME: {
102                 SHARING_LOGD("get pps frame, size:%{public}u.", len);
103                 HandlePpsFrame(dispatcher, frame);
104                 break;
105             }
106             case IDR_FRAME: {
107                 std::chrono::system_clock::time_point now = std::chrono::system_clock::now();
108                 uint64_t pts = std::chrono::duration_cast<std::chrono::milliseconds>(now - start).count();
109                 auto mediaData = dispatcher->RequestDataBuffer(MEDIA_TYPE_VIDEO, len);
110                 if (mediaData == nullptr) {
111                     SHARING_LOGE("RequestDataBuffer get null!");
112                     break;
113                 }
114                 mediaData->mediaType = MEDIA_TYPE_VIDEO;
115                 mediaData->isRaw = false;
116                 mediaData->keyFrame = keyFrame;
117                 mediaData->pts = pts;
118                 SHARING_LOGD("[%{public}" PRId64 "] capture a video into dispatcher:%{public}u, len:%{public}u.",
119                              mediaData->pts, dispatcher->GetDispatcherId(), len);
120                 mediaData->buff = move(frame);
121                 dispatcher->InputData(mediaData);
122                 break;
123             }
124             default:
125                 break;
126         }
127     }
128 }
129 
ScreenCaptureConsumer()130 ScreenCaptureConsumer::ScreenCaptureConsumer()
131 {
132     SHARING_LOGD("capture consumer Id: %{public}u.", GetId());
133 }
134 
~ScreenCaptureConsumer()135 ScreenCaptureConsumer::~ScreenCaptureConsumer()
136 {
137     SHARING_LOGD("~ capture consumer Id: %{public}u.", GetId());
138     Release();
139 }
140 
HandleProsumerInitState(SharingEvent &event)141 void ScreenCaptureConsumer::HandleProsumerInitState(SharingEvent &event)
142 {
143     SHARING_LOGD("trace.");
144     auto pPrivateMsg = std::make_shared<ScreenCaptureSessionEventMsg>();
145     pPrivateMsg->errorCode = ERR_GENERAL_ERROR;
146     pPrivateMsg->errorCode = ERR_OK;
147     pPrivateMsg->type = EVENT_AGENT_STATE_PROSUMER_INIT;
148     pPrivateMsg->toMgr = ModuleType::MODULE_CONTEXT;
149     pPrivateMsg->fromMgr = ModuleType::MODULE_MEDIACHANNEL;
150     pPrivateMsg->srcId = GetId();
151     pPrivateMsg->dstId = event.eventMsg->srcId;
152     pPrivateMsg->agentId = GetSinkAgentId();
153     pPrivateMsg->prosumerId = GetId();
154     pPrivateMsg->requestId = event.eventMsg->requestId;
155 
156     NotifyPrivateEvent(pPrivateMsg);
157 }
158 
HandleProsumerPlay(SharingEvent &event)159 void ScreenCaptureConsumer::HandleProsumerPlay(SharingEvent &event)
160 {
161     SHARING_LOGD("trace.");
162 
163     auto msg = ConvertEventMsg<ScreenCaptureConsumerEventMsg>(event);
164     if (msg) {
165         SHARING_LOGI("msg convert suc, vcodecId:%{public}d, acodecId:%{public}d.",
166                      msg->videoTrack.codecId, msg->audioTrack.codecId);
167         if (msg->audioTrack.codecId != CodecId::CODEC_NONE) {
168             audioTrack_ = msg->audioTrack;
169         }
170 
171         if (msg->videoTrack.codecId != CodecId::CODEC_NONE) {
172             videoTrack_ = msg->videoTrack;
173         }
174 
175         if (isInit_) {
176             SHARING_LOGD("Capture already inited!");
177             return;
178         }
179         if (!InitCapture(msg->screenId)) {
180             SHARING_LOGD("InitCapture failed!");
181             return;
182         }
183         if (!StartCapture()) {
184             SHARING_LOGD("StartCapture failed!");
185             return;
186         }
187     } else {
188         SHARING_LOGE("msg convert fail!");
189     }
190 }
191 
HandleEvent(SharingEvent &event)192 int32_t ScreenCaptureConsumer::HandleEvent(SharingEvent &event)
193 {
194     SHARING_LOGD("trace.");
195     RETURN_INVALID_IF_NULL(event.eventMsg);
196 
197     SHARING_LOGD("eventType: %{public}s = %{public}d, capture consumerId: %{public}u.",
198                  std::string(magic_enum::enum_name(event.eventMsg->type)).c_str(), event.eventMsg->type, GetId());
199     switch (event.eventMsg->type) {
200         case EventType::EVENT_SCREEN_CAPTURE_INIT:
201             HandleProsumerInitState(event);
202             break;
203         case EventType::EVENT_WFD_NOTIFY_RTSP_PLAYED:
204             HandleProsumerPlay(event);
205             break;
206         default:
207             SHARING_LOGI("none process case.");
208             break;
209     }
210 
211     return 0;
212 }
213 
UpdateOperation(ProsumerStatusMsg::Ptr &statusMsg)214 void ScreenCaptureConsumer::UpdateOperation(ProsumerStatusMsg::Ptr &statusMsg)
215 {
216     SHARING_LOGD("trace.");
217     RETURN_IF_NULL(statusMsg);
218 
219     SHARING_LOGD("status: %{public}s consumerId: %{public}u.",
220                  std::string(magic_enum::enum_name(static_cast<ProsumerOptRunningStatus>(statusMsg->status))).c_str(),
221                  GetId());
222     switch (statusMsg->status) {
223         case ProsumerOptRunningStatus::PROSUMER_INIT:
224             statusMsg->status = PROSUMER_NOTIFY_INIT_SUCCESS;
225             break;
226         case ProsumerOptRunningStatus::PROSUMER_START: {
227             paused_ = false;
228             statusMsg->status = PROSUMER_NOTIFY_START_SUCCESS;
229             break;
230         }
231         case ProsumerOptRunningStatus::PROSUMER_PAUSE: {
232             paused_ = true;
233             statusMsg->status = PROSUMER_NOTIFY_PAUSE_SUCCESS;
234             break;
235         }
236         case ProsumerOptRunningStatus::PROSUMER_RESUME: {
237             paused_ = false;
238             statusMsg->status = PROSUMER_NOTIFY_RESUME_SUCCESS;
239             break;
240         }
241         case ProsumerOptRunningStatus::PROSUMER_STOP: {
242             if (StopCapture()) {
243                 statusMsg->status = PROSUMER_NOTIFY_STOP_SUCCESS;
244             } else {
245                 statusMsg->status = PROSUMER_NOTIFY_ERROR;
246                 statusMsg->errorCode = ERR_PROSUMER_STOP;
247             }
248             break;
249         }
250         case ProsumerOptRunningStatus::PROSUMER_DESTROY:
251             Release();
252             SHARING_LOGD("prosumer release success, out.");
253             statusMsg->status = PROSUMER_NOTIFY_DESTROY_SUCCESS;
254             break;
255         default:
256             SHARING_LOGI("none process case.");
257             break;
258     }
259 
260     Notify(statusMsg);
261 }
262 
Release()263 int32_t ScreenCaptureConsumer::Release()
264 {
265     SHARING_LOGD("trace.");
266     StopCapture();
267 
268     SHARING_LOGD("consumerId: %{public}u, capture consumer release=out.", GetId());
269     return 0;
270 }
271 
IsPaused()272 bool ScreenCaptureConsumer::IsPaused()
273 {
274     SHARING_LOGD("trace.");
275     return paused_;
276 }
277 
ReleaseScreenBuffer()278 int32_t ScreenCaptureConsumer::ReleaseScreenBuffer()
279 {
280     SHARING_LOGD("trace.");
281     if (videoSourceScreen_ == nullptr) {
282         return ERR_GENERAL_ERROR;
283     }
284     return videoSourceScreen_->ReleaseScreenBuffer();
285 }
286 
InitCapture(uint64_t screenId)287 bool ScreenCaptureConsumer::InitCapture(uint64_t screenId)
288 {
289     SHARING_LOGD("capture consumerId: %{public}u, viddeo codecId: %{public}d, audio codecId: %{public}d.", GetId(),
290                  videoTrack_.codecId, audioTrack_.codecId);
291     std::lock_guard<std::mutex> lock(mutex_);
292     if (audioTrack_.codecId != CODEC_NONE) {
293         if (!InitAudioEncoder()) {
294             SHARING_LOGE("capture consumerId: %{public}u, init audio encoder failed.", GetId());
295             return false;
296         }
297 
298         if (!InitAudioCapture()) {
299             SHARING_LOGE("capture consumerId: %{public}u, init audio capture failed.", GetId());
300             return false;
301         }
302     }
303 
304     if ((videoTrack_.codecId != CODEC_NONE) && !InitVideoCapture(screenId)) {
305         SHARING_LOGE("capture consumerId: %{public}u, init video capture failed.", GetId());
306         return false;
307     }
308 
309     SHARING_LOGD("capture consumerId: %{public}u, success.", GetId());
310     isInit_ = true;
311     return true;
312 }
313 
InitVideoCapture(uint64_t screenId)314 bool ScreenCaptureConsumer::InitVideoCapture(uint64_t screenId)
315 {
316     SHARING_LOGD("consumerId: %{public}u.", GetId());
317     VideoSourceConfigure config;
318     config.srcScreenId_ = screenId;
319 
320     videoSourceEncoder_ = std::make_shared<VideoSourceEncoder>(shared_from_this());
321     if (!videoSourceEncoder_->InitEncoder(config)) {
322         SHARING_LOGE("InitEncoder failed! %{public}u.", GetId());
323         OnInitVideoCaptureError();
324         return false;
325     }
326     videoSourceScreen_ = std::make_shared<VideoSourceScreen>(videoSourceEncoder_->GetEncoderSurface());
327     videoSourceScreen_->InitScreenSource(config);
328     return true;
329 }
330 
InitAudioCapture()331 bool ScreenCaptureConsumer::InitAudioCapture()
332 {
333     SHARING_LOGD("trace.");
334 
335     if (audioEncoder_ == nullptr) {
336         SHARING_LOGE("InitAudioCapture failed, audioEncoder_ is null! %{public}u.", GetId());
337         return false;
338     }
339     audioSourceCapturer_ = std::make_shared<AudioSourceCapturer>(audioEncoder_);
340     return audioSourceCapturer_->InitAudioCapture();
341 }
342 
InitAudioEncoder()343 bool ScreenCaptureConsumer::InitAudioEncoder()
344 {
345     SHARING_LOGD("trace.");
346     audioEncoder_ = CodecFactory::CreateAudioEncoder(audioTrack_.codecId);
347     if (!audioEncoder_) {
348         SHARING_LOGE("create audio encoder failed.");
349         return false;
350     }
351 
352     audioEncoder_->Init(AUDIO_CHANNEL_STEREO, AUDIO_SAMPLE_BIT_S16LE, AUDIO_SAMPLE_RATE_48000);
353     audioEncoderReceiver_ = std::make_shared<AudioEncoderReceiver>(shared_from_this());
354     audioEncoder_->AddAudioDestination(audioEncoderReceiver_);
355     return true;
356 }
357 
StartCapture()358 bool ScreenCaptureConsumer::StartCapture()
359 {
360     SHARING_LOGD("trace.");
361     std::lock_guard<std::mutex> lock(mutex_);
362     if (isRunning_ || !isInit_) {
363         return false;
364     }
365 
366     isRunning_ = true;
367     if ((videoTrack_.codecId != CODEC_NONE) && !StartVideoCapture()) {
368         SHARING_LOGE("consumerId: %{public}u, start video capture failed.", GetId());
369         isRunning_ = false;
370         return false;
371     }
372 
373     if ((audioTrack_.codecId != CODEC_NONE) && !StartAudioCapture()) {
374         SHARING_LOGE("consumerId: %{public}u, start audio capture failed, only viedo works!", GetId());
375     }
376 
377     SHARING_LOGD("consumerId: %{public}u, success.", GetId());
378     return true;
379 }
380 
StartAudioCapture()381 bool ScreenCaptureConsumer::StartAudioCapture()
382 {
383     SHARING_LOGD("trace.");
384 
385     if (audioSourceCapturer_ == nullptr) {
386         SHARING_LOGE("start capture fail - capturer is null, consumerId: %{public}u.", GetId());
387         return false;
388     }
389     return audioSourceCapturer_->StartAudioCapture();
390 }
391 
StartVideoCapture()392 bool ScreenCaptureConsumer::StartVideoCapture()
393 {
394     SHARING_LOGD("trace.");
395     if (videoSourceScreen_ != nullptr) {
396         videoSourceScreen_->StartScreenSourceCapture();
397     }
398     if (videoSourceEncoder_ != nullptr) {
399         videoSourceEncoder_->StartEncoder();
400     }
401     return true;
402 }
403 
StopCapture()404 bool ScreenCaptureConsumer::StopCapture()
405 {
406     SHARING_LOGD("stop capture consumer, consumerId: %{public}u.", GetId());
407     std::lock_guard<std::mutex> lock(mutex_);
408     if (!isRunning_ || !isInit_) {
409         SHARING_LOGD("capture consumer already stoped, consumerId: %{public}u.", GetId());
410         return true;
411     }
412 
413     isRunning_ = false;
414     StopAudioCapture();
415     StopVideoCapture();
416 
417     return true;
418 }
419 
StopAudioCapture()420 bool ScreenCaptureConsumer::StopAudioCapture()
421 {
422     SHARING_LOGD("trace.");
423 
424     if (audioSourceCapturer_ == nullptr) {
425         SHARING_LOGE("stop capture fail - capturer is null, consumerId: %{public}u.", GetId());
426         return false;
427     }
428     return audioSourceCapturer_->StopAudioCapture();
429 }
430 
StopVideoCapture()431 bool ScreenCaptureConsumer::StopVideoCapture()
432 {
433     SHARING_LOGD("trace.");
434     if (videoTrack_.codecId != CODEC_NONE) {
435         if (videoSourceEncoder_ != nullptr) {
436             videoSourceEncoder_->StopEncoder();
437         }
438         if (videoSourceScreen_ != nullptr) {
439             videoSourceScreen_->StopScreenSourceCapture();
440         }
441     }
442     return true;
443 }
444 
OnInitVideoCaptureError()445 void ScreenCaptureConsumer::OnInitVideoCaptureError()
446 {
447     SHARING_LOGD("trace.");
448     auto statusMsg = std::make_shared<ProsumerStatusMsg>();
449     statusMsg->errorCode = ERR_PROSUMER_VIDEO_CAPTURE;
450     if (audioTrack_.codecId == CODEC_NONE) {
451         statusMsg->status = PROSUMER_NOTIFY_ERROR;
452     }
453 
454     Notify(statusMsg);
455 }
456 
457 REGISTER_CLASS_REFLECTOR(ScreenCaptureConsumer);
458 } // namespace Sharing
459 } // namespace OHOS
460