Hold a reference to AndroidVideoTrackSource while calling onFrameCaptured.

This makes it safe to deliver frames to the sink from VideoProcessor
even after setSink has been called with null reference without danger
of use after free.

Bug: b/148063550
Change-Id: Ib78f75ac49fc6117f744c55da1a4e671bbdcdf22
Reviewed-on: https://webrtc-review.googlesource.com/c/src/+/168160
Reviewed-by: Paulina Hensman <phensman@webrtc.org>
Commit-Queue: Sami Kalliomäki <sakal@webrtc.org>
Cr-Commit-Position: refs/heads/master@{#30455}

sdk/android/BUILD.gn

@@ -281,8 +281,11 @@ if (is_android) {
       "api/org/webrtc/NativePeerConnectionFactory.java",
       "api/org/webrtc/NetEqFactoryFactory.java",
       "api/org/webrtc/NetworkControllerFactoryFactory.java",
-      "api/org/webrtc/NetworkMonitor.java",  # TODO(sakal): Break dependencies and move to base_java.
-      "api/org/webrtc/NetworkMonitorAutoDetect.java",  # TODO(sakal): Break dependencies and move to base_java.
+      "api/org/webrtc/NetworkMonitor.java",  # TODO(sakal): Break dependencies
+                                             # and move to base_java.
+      "api/org/webrtc/NetworkMonitorAutoDetect.java",  # TODO(sakal): Break
+                                                       # dependencies and move
+                                                       # to base_java.
       "api/org/webrtc/NetworkStatePredictorFactoryFactory.java",
       "api/org/webrtc/PeerConnection.java",
       "api/org/webrtc/PeerConnectionDependencies.java",
@@ -1537,6 +1540,7 @@ if (is_android) {
       "tests/src/org/webrtc/GlGenericDrawerTest.java",
       "tests/src/org/webrtc/HardwareVideoEncoderTest.java",
       "tests/src/org/webrtc/IceCandidateTest.java",
+      "tests/src/org/webrtc/RefCountDelegateTest.java",
       "tests/src/org/webrtc/ScalingSettingsTest.java",
     ]
 

sdk/android/api/org/webrtc/MediaSource.java

@@ -25,9 +25,11 @@ public class MediaSource {
     }
   }
 
+  private final RefCountDelegate refCountDelegate;
   private long nativeSource;
 
   public MediaSource(long nativeSource) {
+    refCountDelegate = new RefCountDelegate(() -> JniCommon.nativeReleaseRef(nativeSource));
     this.nativeSource = nativeSource;
   }
 
@@ -38,7 +40,7 @@ public class MediaSource {
 
   public void dispose() {
     checkMediaSourceExists();
-    JniCommon.nativeReleaseRef(nativeSource);
+    refCountDelegate.release();
     nativeSource = 0;
   }
 
@@ -48,6 +50,20 @@ public class MediaSource {
     return nativeSource;
   }
 
+  /**
+   * Runs code in {@code runnable} holding a reference to the media source. If the object has
+   * already been released, does nothing.
+   */
+  void runWithReference(Runnable runnable) {
+    if (refCountDelegate.safeRetain()) {
+      try {
+        runnable.run();
+      } finally {
+        refCountDelegate.release();
+      }
+    }
+  }
+
   private void checkMediaSourceExists() {
     if (nativeSource == 0) {
       throw new IllegalStateException("MediaSource has been disposed.");

sdk/android/api/org/webrtc/VideoProcessor.java

@@ -54,7 +54,7 @@ public interface VideoProcessor extends CapturerObserver {
 
   /**
    * Set the sink that receives the output from this processor. Null can be passed in to unregister
-   * a sink. After this call returns, no frames should be delivered to an unregistered sink.
+   * a sink.
    */
   void setSink(@Nullable VideoSink sink);
 

sdk/android/api/org/webrtc/VideoSource.java

@@ -135,7 +135,9 @@ public class VideoSource extends MediaSource {
       }
       videoProcessor = newVideoProcessor;
       if (newVideoProcessor != null) {
-        newVideoProcessor.setSink(nativeAndroidVideoTrackSource::onFrameCaptured);
+        newVideoProcessor.setSink(
+            (frame)
+                -> runWithReference(() -> nativeAndroidVideoTrackSource.onFrameCaptured(frame)));
         if (isCapturerRunning) {
           newVideoProcessor.onCapturerStarted(/* success= */ true);
         }

sdk/android/src/java/org/webrtc/RefCountDelegate.java

@@ -45,4 +45,19 @@ class RefCountDelegate implements RefCounted {
       releaseCallback.run();
     }
   }
+
+  /**
+   * Tries to retain the object. Can be used in scenarios where it is unknown if the object has
+   * already been released. Returns true if successful or false if the object was already released.
+   */
+  boolean safeRetain() {
+    int currentRefCount = refCount.get();
+    while (currentRefCount != 0) {
+      if (refCount.weakCompareAndSet(currentRefCount, currentRefCount + 1)) {
+        return true;
+      }
+      currentRefCount = refCount.get();
+    }
+    return false;
+  }
 }

sdk/android/tests/src/org/webrtc/RefCountDelegateTest.java

@@ -0,0 +1,83 @@
+/*
+ *  Copyright 2020 The WebRTC Project Authors. All rights reserved.
+ *
+ *  Use of this source code is governed by a BSD-style license
+ *  that can be found in the LICENSE file in the root of the source
+ *  tree. An additional intellectual property rights grant can be found
+ *  in the file PATENTS.  All contributing project authors may
+ *  be found in the AUTHORS file in the root of the source tree.
+ */
+
+package org.webrtc;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+
+import org.chromium.testing.local.LocalRobolectricTestRunner;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.annotation.Config;
+
+@RunWith(LocalRobolectricTestRunner.class)
+@Config(manifest = Config.NONE)
+public class RefCountDelegateTest {
+  @Mock Runnable mockReleaseCallback;
+  private RefCountDelegate refCountDelegate;
+
+  @Before
+  public void setUp() {
+    MockitoAnnotations.initMocks(this);
+
+    refCountDelegate = new RefCountDelegate(mockReleaseCallback);
+  }
+
+  @Test
+  public void testReleaseRunsReleaseCallback() {
+    refCountDelegate.release();
+    verify(mockReleaseCallback).run();
+  }
+
+  @Test
+  public void testRetainIncreasesRefCount() {
+    refCountDelegate.retain();
+
+    refCountDelegate.release();
+    verify(mockReleaseCallback, never()).run();
+
+    refCountDelegate.release();
+    verify(mockReleaseCallback).run();
+  }
+
+  @Test(expected = IllegalStateException.class)
+  public void testReleaseAfterFreeThrowsIllegalStateException() {
+    refCountDelegate.release();
+    refCountDelegate.release();
+  }
+
+  @Test(expected = IllegalStateException.class)
+  public void testRetainAfterFreeThrowsIllegalStateException() {
+    refCountDelegate.release();
+    refCountDelegate.retain();
+  }
+
+  @Test
+  public void testSafeRetainBeforeFreeReturnsTrueAndIncreasesRefCount() {
+    assertThat(refCountDelegate.safeRetain()).isTrue();
+
+    refCountDelegate.release();
+    verify(mockReleaseCallback, never()).run();
+
+    refCountDelegate.release();
+    verify(mockReleaseCallback).run();
+  }
+
+  @Test
+  public void testSafeRetainAfterFreeReturnsFalse() {
+    refCountDelegate.release();
+    assertThat(refCountDelegate.safeRetain()).isFalse();
+  }
+}