playwright/browser_patches/firefox/juggler/screencast/HeadlessWindowCapturer.cpp

/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#include "HeadlessWindowCapturer.h"

#include <bit>

#include "api/video/i420_buffer.h"
#include "HeadlessWidget.h"
#include "libyuv.h"
#include "mozilla/EndianUtils.h"
#include "mozilla/gfx/DataSurfaceHelpers.h"
#include "rtc_base/ref_counted_object.h"
#include "rtc_base/time_utils.h"
#include "api/scoped_refptr.h"

using namespace mozilla::widget;
using namespace webrtc;

namespace mozilla {

webrtc::scoped_refptr<webrtc::VideoCaptureModuleEx> HeadlessWindowCapturer::Create(HeadlessWidget* headlessWindow) {
  return webrtc::scoped_refptr<webrtc::VideoCaptureModuleEx>(
    new webrtc::RefCountedObject<HeadlessWindowCapturer>(headlessWindow)
  );
}

HeadlessWindowCapturer::HeadlessWindowCapturer(mozilla::widget::HeadlessWidget* window)
    : mWindow(window) {
}
HeadlessWindowCapturer::~HeadlessWindowCapturer() {
  StopCapture();
}


void HeadlessWindowCapturer::RegisterCaptureDataCallback(webrtc::VideoSinkInterface<webrtc::VideoFrame>* dataCallback) {
  webrtc::CritScope lock2(&_callBackCs);
  _dataCallBacks.insert(dataCallback);
}

void HeadlessWindowCapturer::RegisterCaptureDataCallback(webrtc::RawVideoSinkInterface* dataCallback) {
}

void HeadlessWindowCapturer::DeRegisterCaptureDataCallback() {
  webrtc::CritScope lock2(&_callBackCs);
  _dataCallBacks.clear();
}

void HeadlessWindowCapturer::RegisterRawFrameCallback(webrtc::RawFrameCallback* rawFrameCallback) {
  webrtc::CritScope lock2(&_callBackCs);
  _rawFrameCallbacks.insert(rawFrameCallback);
}

void HeadlessWindowCapturer::DeRegisterRawFrameCallback(webrtc::RawFrameCallback* rawFrameCallback) {
  webrtc::CritScope lock2(&_callBackCs);
  auto it = _rawFrameCallbacks.find(rawFrameCallback);
  if (it != _rawFrameCallbacks.end()) {
    _rawFrameCallbacks.erase(it);
  }
}

void HeadlessWindowCapturer::NotifyFrameCaptured(const webrtc::VideoFrame& frame) {
  webrtc::CritScope lock2(&_callBackCs);
  for (auto dataCallBack : _dataCallBacks)
    dataCallBack->OnFrame(frame);
}

int32_t HeadlessWindowCapturer::StartCapture(const webrtc::VideoCaptureCapability& capability) {
  mWindow->SetSnapshotListener([this] (RefPtr<gfx::DataSourceSurface>&& dataSurface){
    if (!NS_IsInCompositorThread()) {
      fprintf(stderr, "SnapshotListener is called not on the Compositor thread!\n");
      return;
    }

    if (dataSurface->GetFormat() != gfx::SurfaceFormat::B8G8R8A8) {
      fprintf(stderr, "Unexpected snapshot surface format: %hhd\n", dataSurface->GetFormat());
      return;
    }

    webrtc::VideoCaptureCapability frameInfo;
    frameInfo.width = dataSurface->GetSize().width;
    frameInfo.height = dataSurface->GetSize().height;
    if constexpr (std::endian::native == std::endian::little) {
      frameInfo.videoType = VideoType::kARGB;
    } else {
      frameInfo.videoType = VideoType::kBGRA;
    }

    {
      webrtc::CritScope lock2(&_callBackCs);
      for (auto rawFrameCallback : _rawFrameCallbacks) {
        rawFrameCallback->OnRawFrame(dataSurface->GetData(), dataSurface->Stride(), frameInfo);
      }
      if (!_dataCallBacks.size())
        return;
    }

    int width = dataSurface->GetSize().width;
    int height = dataSurface->GetSize().height;
    webrtc::scoped_refptr<I420Buffer> buffer = I420Buffer::Create(width, height);

    gfx::DataSourceSurface::ScopedMap map(dataSurface.get(), gfx::DataSourceSurface::MapType::READ);
    if (!map.IsMapped()) {
      fprintf(stderr, "Failed to map snapshot bytes!\n");
      return;
    }

    const int conversionResult = std::endian::native == std::endian::little ? libyuv::ARGBToI420(
      map.GetData(), map.GetStride(),
      buffer->MutableDataY(), buffer->StrideY(),
      buffer->MutableDataU(), buffer->StrideU(),
      buffer->MutableDataV(), buffer->StrideV(),
      width, height
    ) : libyuv::BGRAToI420(
      map.GetData(), map.GetStride(),
      buffer->MutableDataY(), buffer->StrideY(),
      buffer->MutableDataU(), buffer->StrideU(),
      buffer->MutableDataV(), buffer->StrideV(),
      width, height
    );
    if (conversionResult != 0) {
      fprintf(stderr, "Failed to convert capture frame to I420: %d\n", conversionResult);
      return;
    }

    VideoFrame captureFrame(buffer, 0, webrtc::TimeMillis(), kVideoRotation_0);
    NotifyFrameCaptured(captureFrame);
  });
  return 0;
}

int32_t HeadlessWindowCapturer::StopCapture() {
  if (!CaptureStarted())
    return 0;
  mWindow->SetSnapshotListener(nullptr);
  return 0;
}

bool HeadlessWindowCapturer::CaptureStarted() {
  return true;
}

}  // namespace mozilla