browser(firefox): cross thread sync in screencast (#16320)

* nsIScreencastServiceClient is not thread safe refcounted so we make nsScreencastService::Session a thread safe refcounted object and keep it alive while there are inflight frames. Once such frames get handled on the main thread we check if the session has been stopped.
* Removed mCaptureCallbackCs in favor of atomic counter (mClient is not accessed only on the main thread).
* HeadlessWindowCapturer now holds RefPtr to the headless window object to avoid use after free when clearing it as a listener on the widget.
* ScreencastEncoder is not ref counted anymore.

Pretty-diff: 5f5042ff1e

browser_patches/firefox-beta/BUILD_NUMBER

@@ -1,2 +1,2 @@
-1342
-Changed: yurys@chromium.org Thu Aug  4 18:33:22 PDT 2022
+1343
+Changed: yurys@chromium.org Fri Aug  5 15:17:14 PDT 2022

browser_patches/firefox-beta/juggler/screencast/HeadlessWindowCapturer.h

@@ -55,7 +55,7 @@ class HeadlessWindowCapturer : public webrtc::VideoCaptureModuleEx {
  private:
   void NotifyFrameCaptured(const webrtc::VideoFrame& frame);
 
-  mozilla::widget::HeadlessWidget* mWindow = nullptr;
+  RefPtr<mozilla::widget::HeadlessWidget> mWindow;
   rtc::RecursiveCriticalSection _callBackCs;
   std::set<rtc::VideoSinkInterface<webrtc::VideoFrame>*> _dataCallBacks;
   std::set<webrtc::RawFrameCallback*> _rawFrameCallbacks;

browser_patches/firefox-beta/juggler/screencast/ScreencastEncoder.cpp

@@ -299,7 +299,7 @@ ScreencastEncoder::~ScreencastEncoder()
 {
 }
 
-RefPtr<ScreencastEncoder> ScreencastEncoder::create(nsCString& errorString, const nsCString& filePath, int width, int height, const gfx::IntMargin& margin)
+std::unique_ptr<ScreencastEncoder> ScreencastEncoder::create(nsCString& errorString, const nsCString& filePath, int width, int height, const gfx::IntMargin& margin)
 {
     vpx_codec_iface_t* codec_interface = vpx_codec_vp8_cx();
     if (!codec_interface) {
@@ -340,7 +340,7 @@ RefPtr<ScreencastEncoder> ScreencastEncoder::create(nsCString& errorString, cons
 
     std::unique_ptr<VPXCodec> vpxCodec(new VPXCodec(std::move(codec), cfg, file));
     // fprintf(stderr, "ScreencastEncoder initialized with: %s\n", vpx_codec_iface_name(codec_interface));
-    return new ScreencastEncoder(std::move(vpxCodec), margin);
+    return std::make_unique<ScreencastEncoder>(std::move(vpxCodec), margin);
 }
 
 void ScreencastEncoder::flushLastFrame()

browser_patches/firefox-beta/juggler/screencast/ScreencastEncoder.h

@@ -19,22 +19,20 @@ class VideoFrame;
 namespace mozilla {
 
 class ScreencastEncoder {
-    NS_INLINE_DECL_THREADSAFE_REFCOUNTING(ScreencastEncoder)
 public:
     static constexpr int fps = 25;
 
-    static RefPtr<ScreencastEncoder> create(nsCString& errorString, const nsCString& filePath, int width, int height, const gfx::IntMargin& margin);
+    static std::unique_ptr<ScreencastEncoder> create(nsCString& errorString, const nsCString& filePath, int width, int height, const gfx::IntMargin& margin);
 
     class VPXCodec;
     ScreencastEncoder(std::unique_ptr<VPXCodec>, const gfx::IntMargin& margin);
+    ~ScreencastEncoder();
 
     void encodeFrame(const webrtc::VideoFrame& videoFrame);
 
     void finish(std::function<void()>&& callback);
 
 private:
-    ~ScreencastEncoder();
-
     void flushLastFrame();
 
     std::unique_ptr<VPXCodec> m_vpxCodec;

browser_patches/firefox-beta/juggler/screencast/nsScreencastService.cpp

@@ -78,11 +78,10 @@ nsresult generateUid(nsString& uid) {
 
 class nsScreencastService::Session : public rtc::VideoSinkInterface<webrtc::VideoFrame>,
                                      public webrtc::RawFrameCallback {
- public:
   Session(
     nsIScreencastServiceClient* client,
     rtc::scoped_refptr<webrtc::VideoCaptureModuleEx>&& capturer,
-    RefPtr<ScreencastEncoder>&& encoder,
+    std::unique_ptr<ScreencastEncoder> encoder,
     int width, int height,
     int viewportWidth, int viewportHeight,
     gfx::IntMargin margin,
@@ -97,6 +96,20 @@ class nsScreencastService::Session : public rtc::VideoSinkInterface<webrtc::Vide
       , mViewportHeight(viewportHeight)
       , mMargin(margin) {
   }
+  ~Session() override = default;
+
+ public:
+  NS_INLINE_DECL_THREADSAFE_REFCOUNTING(Session)
+  static RefPtr<Session> Create(
+    nsIScreencastServiceClient* client,
+    rtc::scoped_refptr<webrtc::VideoCaptureModuleEx>&& capturer,
+    std::unique_ptr<ScreencastEncoder> encoder,
+    int width, int height,
+    int viewportWidth, int viewportHeight,
+    gfx::IntMargin margin,
+    uint32_t jpegQuality) {
+    return do_AddRef(new Session(client, std::move(capturer), std::move(encoder), width, height, viewportWidth, viewportHeight, margin, jpegQuality));
+  }
 
   bool Start() {
     webrtc::VideoCaptureCapability capability;
@@ -119,6 +132,11 @@ class nsScreencastService::Session : public rtc::VideoSinkInterface<webrtc::Vide
   }
 
   void Stop() {
+    if (mStopped) {
+      fprintf(stderr, "Screencast session has already been stopped\n");
+      return;
+    }
+    mStopped = true;
     if (mEncoder)
       mCaptureModule->DeRegisterCaptureDataCallback(this);
     else
@@ -128,23 +146,23 @@ class nsScreencastService::Session : public rtc::VideoSinkInterface<webrtc::Vide
       fprintf(stderr, "StopCapture error %d\n", error);
     }
     if (mEncoder) {
-      rtc::CritScope lock(&mCaptureCallbackCs);
-      mEncoder->finish([client = std::move(mClient)] {
+      mEncoder->finish([this, protect = RefPtr{this}] {
         NS_DispatchToMainThread(NS_NewRunnableFunction(
-            "NotifyScreencastStopped", [client = std::move(client)]() -> void {
-              client->ScreencastStopped();
+            "NotifyScreencastStopped", [this, protect = std::move(protect)]() -> void {
+              mClient->ScreencastStopped();
             }));
       });
     } else {
-      rtc::CritScope lock(&mCaptureCallbackCs);
       mClient->ScreencastStopped();
-      mClient = nullptr;
     }
   }
 
   void ScreencastFrameAck() {
-    rtc::CritScope lock(&mCaptureCallbackCs);
-    --mFramesInFlight;
+    if (mFramesInFlight.load() == 0) {
+      fprintf(stderr, "ScreencastFrameAck is called while there are no inflight frames\n");
+      return;
+    }
+    mFramesInFlight.fetch_sub(1);
   }
 
   // These callbacks end up running on the VideoCapture thread.
@@ -167,15 +185,8 @@ class nsScreencastService::Session : public rtc::VideoSinkInterface<webrtc::Vide
     if (mViewportHeight && pageHeight > mViewportHeight)
       pageHeight = mViewportHeight;
 
-    {
-      rtc::CritScope lock(&mCaptureCallbackCs);
-      if (mFramesInFlight >= kMaxFramesInFlight) {
-        return;
-      }
-      ++mFramesInFlight;
-      if (!mClient)
-        return;
-    }
+    if (mFramesInFlight.load() >= kMaxFramesInFlight)
+      return;
 
     int screenshotWidth = pageWidth;
     int screenshotHeight = pageHeight;
@@ -256,21 +267,23 @@ class nsScreencastService::Session : public rtc::VideoSinkInterface<webrtc::Vide
       return;
     }
 
-    nsIScreencastServiceClient* client = mClient.get();
+    mFramesInFlight.fetch_add(1);
     NS_DispatchToMainThread(NS_NewRunnableFunction(
-        "NotifyScreencastFrame", [client, base64, pageWidth, pageHeight]() -> void {
+        "NotifyScreencastFrame", [this, protect = RefPtr{this}, base64, pageWidth, pageHeight]() -> void {
+          if (mStopped)
+            return;
           NS_ConvertUTF8toUTF16 utf16(base64);
-          client->ScreencastFrame(utf16, pageWidth, pageHeight);
+          mClient->ScreencastFrame(utf16, pageWidth, pageHeight);
         }));
   }
 
  private:
   RefPtr<nsIScreencastServiceClient> mClient;
   rtc::scoped_refptr<webrtc::VideoCaptureModuleEx> mCaptureModule;
-  RefPtr<ScreencastEncoder> mEncoder;
+  std::unique_ptr<ScreencastEncoder> mEncoder;
   uint32_t mJpegQuality;
-  rtc::RecursiveCriticalSection mCaptureCallbackCs;
-  uint32_t mFramesInFlight = 0;
+  bool mStopped = false;
+  std::atomic<uint32_t> mFramesInFlight = 0;
   int mWidth;
   int mHeight;
   int mViewportWidth;
@@ -322,7 +335,7 @@ nsresult nsScreencastService::StartVideoRecording(nsIScreencastServiceClient* aC
   margin.top += offsetTop;
 
   nsCString error;
-  RefPtr<ScreencastEncoder> encoder;
+  std::unique_ptr<ScreencastEncoder> encoder;
   if (isVideo) {
     encoder = ScreencastEncoder::create(error, PromiseFlatCString(aVideoFileName), width, height, margin);
     if (!encoder) {
@@ -336,7 +349,7 @@ nsresult nsScreencastService::StartVideoRecording(nsIScreencastServiceClient* aC
   NS_ENSURE_SUCCESS(rv, rv);
   sessionId = uid;
 
-  auto session = std::make_unique<Session>(aClient, std::move(capturer), std::move(encoder), width, height, viewportWidth, viewportHeight, margin, isVideo ? 0 : quality);
+  auto session = Session::Create(aClient, std::move(capturer), std::move(encoder), width, height, viewportWidth, viewportHeight, margin, isVideo ? 0 : quality);
   if (!session->Start())
     return NS_ERROR_FAILURE;
   mIdToSession.emplace(sessionId, std::move(session));

browser_patches/firefox-beta/juggler/screencast/nsScreencastService.h

@@ -23,7 +23,7 @@ class nsScreencastService final : public nsIScreencastService {
   ~nsScreencastService();
 
   class Session;
-  std::map<nsString, std::unique_ptr<Session>> mIdToSession;
+  std::map<nsString, RefPtr<Session>> mIdToSession;
 };
 
 }  // namespace mozilla

browser_patches/firefox/BUILD_NUMBER

@@ -1,2 +1,2 @@
-1343
-Changed: yurys@chromium.org Thu Aug  4 18:29:49 PDT 2022
+1344
+Changed: yurys@chromium.org Fri Aug  5 15:15:08 PDT 2022

browser_patches/firefox/juggler/screencast/HeadlessWindowCapturer.h

@@ -55,7 +55,7 @@ class HeadlessWindowCapturer : public webrtc::VideoCaptureModuleEx {
  private:
   void NotifyFrameCaptured(const webrtc::VideoFrame& frame);
 
-  mozilla::widget::HeadlessWidget* mWindow = nullptr;
+  RefPtr<mozilla::widget::HeadlessWidget> mWindow;
   rtc::RecursiveCriticalSection _callBackCs;
   std::set<rtc::VideoSinkInterface<webrtc::VideoFrame>*> _dataCallBacks;
   std::set<webrtc::RawFrameCallback*> _rawFrameCallbacks;

browser_patches/firefox/juggler/screencast/ScreencastEncoder.cpp

@@ -299,7 +299,7 @@ ScreencastEncoder::~ScreencastEncoder()
 {
 }
 
-RefPtr<ScreencastEncoder> ScreencastEncoder::create(nsCString& errorString, const nsCString& filePath, int width, int height, const gfx::IntMargin& margin)
+std::unique_ptr<ScreencastEncoder> ScreencastEncoder::create(nsCString& errorString, const nsCString& filePath, int width, int height, const gfx::IntMargin& margin)
 {
     vpx_codec_iface_t* codec_interface = vpx_codec_vp8_cx();
     if (!codec_interface) {
@@ -340,7 +340,7 @@ RefPtr<ScreencastEncoder> ScreencastEncoder::create(nsCString& errorString, cons
 
     std::unique_ptr<VPXCodec> vpxCodec(new VPXCodec(std::move(codec), cfg, file));
     // fprintf(stderr, "ScreencastEncoder initialized with: %s\n", vpx_codec_iface_name(codec_interface));
-    return new ScreencastEncoder(std::move(vpxCodec), margin);
+    return std::make_unique<ScreencastEncoder>(std::move(vpxCodec), margin);
 }
 
 void ScreencastEncoder::flushLastFrame()

browser_patches/firefox/juggler/screencast/ScreencastEncoder.h

@@ -19,22 +19,20 @@ class VideoFrame;
 namespace mozilla {
 
 class ScreencastEncoder {
-    NS_INLINE_DECL_THREADSAFE_REFCOUNTING(ScreencastEncoder)
 public:
     static constexpr int fps = 25;
 
-    static RefPtr<ScreencastEncoder> create(nsCString& errorString, const nsCString& filePath, int width, int height, const gfx::IntMargin& margin);
+    static std::unique_ptr<ScreencastEncoder> create(nsCString& errorString, const nsCString& filePath, int width, int height, const gfx::IntMargin& margin);
 
     class VPXCodec;
     ScreencastEncoder(std::unique_ptr<VPXCodec>, const gfx::IntMargin& margin);
+    ~ScreencastEncoder();
 
     void encodeFrame(const webrtc::VideoFrame& videoFrame);
 
     void finish(std::function<void()>&& callback);
 
 private:
-    ~ScreencastEncoder();
-
     void flushLastFrame();
 
     std::unique_ptr<VPXCodec> m_vpxCodec;

browser_patches/firefox/juggler/screencast/nsScreencastService.cpp

@@ -78,11 +78,10 @@ nsresult generateUid(nsString& uid) {
 
 class nsScreencastService::Session : public rtc::VideoSinkInterface<webrtc::VideoFrame>,
                                      public webrtc::RawFrameCallback {
- public:
   Session(
     nsIScreencastServiceClient* client,
     rtc::scoped_refptr<webrtc::VideoCaptureModuleEx>&& capturer,
-    RefPtr<ScreencastEncoder>&& encoder,
+    std::unique_ptr<ScreencastEncoder> encoder,
     int width, int height,
     int viewportWidth, int viewportHeight,
     gfx::IntMargin margin,
@@ -97,6 +96,20 @@ class nsScreencastService::Session : public rtc::VideoSinkInterface<webrtc::Vide
       , mViewportHeight(viewportHeight)
       , mMargin(margin) {
   }
+  ~Session() override = default;
+
+ public:
+  NS_INLINE_DECL_THREADSAFE_REFCOUNTING(Session)
+  static RefPtr<Session> Create(
+    nsIScreencastServiceClient* client,
+    rtc::scoped_refptr<webrtc::VideoCaptureModuleEx>&& capturer,
+    std::unique_ptr<ScreencastEncoder> encoder,
+    int width, int height,
+    int viewportWidth, int viewportHeight,
+    gfx::IntMargin margin,
+    uint32_t jpegQuality) {
+    return do_AddRef(new Session(client, std::move(capturer), std::move(encoder), width, height, viewportWidth, viewportHeight, margin, jpegQuality));
+  }
 
   bool Start() {
     webrtc::VideoCaptureCapability capability;
@@ -119,6 +132,11 @@ class nsScreencastService::Session : public rtc::VideoSinkInterface<webrtc::Vide
   }
 
   void Stop() {
+    if (mStopped) {
+      fprintf(stderr, "Screencast session has already been stopped\n");
+      return;
+    }
+    mStopped = true;
     if (mEncoder)
       mCaptureModule->DeRegisterCaptureDataCallback(this);
     else
@@ -128,23 +146,23 @@ class nsScreencastService::Session : public rtc::VideoSinkInterface<webrtc::Vide
       fprintf(stderr, "StopCapture error %d\n", error);
     }
     if (mEncoder) {
-      rtc::CritScope lock(&mCaptureCallbackCs);
-      mEncoder->finish([client = std::move(mClient)] {
+      mEncoder->finish([this, protect = RefPtr{this}] {
         NS_DispatchToMainThread(NS_NewRunnableFunction(
-            "NotifyScreencastStopped", [client = std::move(client)]() -> void {
-              client->ScreencastStopped();
+            "NotifyScreencastStopped", [this, protect = std::move(protect)]() -> void {
+              mClient->ScreencastStopped();
             }));
       });
     } else {
-      rtc::CritScope lock(&mCaptureCallbackCs);
       mClient->ScreencastStopped();
-      mClient = nullptr;
     }
   }
 
   void ScreencastFrameAck() {
-    rtc::CritScope lock(&mCaptureCallbackCs);
-    --mFramesInFlight;
+    if (mFramesInFlight.load() == 0) {
+      fprintf(stderr, "ScreencastFrameAck is called while there are no inflight frames\n");
+      return;
+    }
+    mFramesInFlight.fetch_sub(1);
   }
 
   // These callbacks end up running on the VideoCapture thread.
@@ -167,15 +185,8 @@ class nsScreencastService::Session : public rtc::VideoSinkInterface<webrtc::Vide
     if (mViewportHeight && pageHeight > mViewportHeight)
       pageHeight = mViewportHeight;
 
-    {
-      rtc::CritScope lock(&mCaptureCallbackCs);
-      if (mFramesInFlight >= kMaxFramesInFlight) {
-        return;
-      }
-      ++mFramesInFlight;
-      if (!mClient)
-        return;
-    }
+    if (mFramesInFlight.load() >= kMaxFramesInFlight)
+      return;
 
     int screenshotWidth = pageWidth;
     int screenshotHeight = pageHeight;
@@ -256,21 +267,23 @@ class nsScreencastService::Session : public rtc::VideoSinkInterface<webrtc::Vide
       return;
     }
 
-    nsIScreencastServiceClient* client = mClient.get();
+    mFramesInFlight.fetch_add(1);
     NS_DispatchToMainThread(NS_NewRunnableFunction(
-        "NotifyScreencastFrame", [client, base64, pageWidth, pageHeight]() -> void {
+        "NotifyScreencastFrame", [this, protect = RefPtr{this}, base64, pageWidth, pageHeight]() -> void {
+          if (mStopped)
+            return;
           NS_ConvertUTF8toUTF16 utf16(base64);
-          client->ScreencastFrame(utf16, pageWidth, pageHeight);
+          mClient->ScreencastFrame(utf16, pageWidth, pageHeight);
         }));
   }
 
  private:
   RefPtr<nsIScreencastServiceClient> mClient;
   rtc::scoped_refptr<webrtc::VideoCaptureModuleEx> mCaptureModule;
-  RefPtr<ScreencastEncoder> mEncoder;
+  std::unique_ptr<ScreencastEncoder> mEncoder;
   uint32_t mJpegQuality;
-  rtc::RecursiveCriticalSection mCaptureCallbackCs;
-  uint32_t mFramesInFlight = 0;
+  bool mStopped = false;
+  std::atomic<uint32_t> mFramesInFlight = 0;
   int mWidth;
   int mHeight;
   int mViewportWidth;
@@ -322,7 +335,7 @@ nsresult nsScreencastService::StartVideoRecording(nsIScreencastServiceClient* aC
   margin.top += offsetTop;
 
   nsCString error;
-  RefPtr<ScreencastEncoder> encoder;
+  std::unique_ptr<ScreencastEncoder> encoder;
   if (isVideo) {
     encoder = ScreencastEncoder::create(error, PromiseFlatCString(aVideoFileName), width, height, margin);
     if (!encoder) {
@@ -336,7 +349,7 @@ nsresult nsScreencastService::StartVideoRecording(nsIScreencastServiceClient* aC
   NS_ENSURE_SUCCESS(rv, rv);
   sessionId = uid;
 
-  auto session = std::make_unique<Session>(aClient, std::move(capturer), std::move(encoder), width, height, viewportWidth, viewportHeight, margin, isVideo ? 0 : quality);
+  auto session = Session::Create(aClient, std::move(capturer), std::move(encoder), width, height, viewportWidth, viewportHeight, margin, isVideo ? 0 : quality);
   if (!session->Start())
     return NS_ERROR_FAILURE;
   mIdToSession.emplace(sessionId, std::move(session));

browser_patches/firefox/juggler/screencast/nsScreencastService.h

@@ -23,7 +23,7 @@ class nsScreencastService final : public nsIScreencastService {
   ~nsScreencastService();
 
   class Session;
-  std::map<nsString, std::unique_ptr<Session>> mIdToSession;
+  std::map<nsString, RefPtr<Session>> mIdToSession;
 };
 
 }  // namespace mozilla