Cookie-based validation in CronetDataSource

Using cookie validation from streamer, streamer can enforce that only clients who have the cookie are able to stream the video.

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=171999924

extensions/cronet/src/androidTest/java/com/google/android/exoplayer2/ext/cronet/CronetDataSourceTest.java

@@ -18,6 +18,7 @@ package com.google.android.exoplayer2.ext.cronet;
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.Matchers.any;
@@ -45,6 +46,7 @@ import com.google.android.exoplayer2.upstream.TransferListener;
 import com.google.android.exoplayer2.util.Clock;
 import com.google.android.exoplayer2.util.Predicate;
 import java.io.IOException;
+import java.net.CookieManager;
 import java.net.SocketTimeoutException;
 import java.net.UnknownHostException;
 import java.nio.ByteBuffer;
@@ -102,12 +104,14 @@ public final class CronetDataSourceTest {
   @Mock private CronetEngine mockCronetEngine;
 
   private CronetDataSource dataSourceUnderTest;
+  private boolean redirectCalled;
 
   @Before
   public void setUp() throws Exception {
     System.setProperty("dexmaker.dexcache",
         InstrumentationRegistry.getTargetContext().getCacheDir().getPath());
     initMocks(this);
+    CookieManager cookieManager = new CookieManager();
     dataSourceUnderTest = spy(
         new CronetDataSource(
             mockCronetEngine,
@@ -118,7 +122,8 @@ public final class CronetDataSourceTest {
             TEST_READ_TIMEOUT_MS,
             true, // resetTimeoutOnRedirects
             mockClock,
-            null));
+            null,
+            cookieManager));
     when(mockContentTypePredicate.evaluate(anyString())).thenReturn(true);
     when(mockCronetEngine.newUrlRequestBuilder(
             anyString(), any(UrlRequest.Callback.class), any(Executor.class)))
@@ -138,10 +143,14 @@ public final class CronetDataSourceTest {
   }
 
   private UrlResponseInfo createUrlResponseInfo(int statusCode) {
+    return createUrlResponseInfoWithUrl(TEST_URL, statusCode);
+  }
+
+  private UrlResponseInfo createUrlResponseInfoWithUrl(String url, int statusCode) {
     ArrayList<Map.Entry<String, String>> responseHeaderList = new ArrayList<>();
     responseHeaderList.addAll(testResponseHeader.entrySet());
     return new UrlResponseInfoImpl(
-        Collections.singletonList(TEST_URL),
+        Collections.singletonList(url),
         statusCode,
         null, // httpStatusText
         responseHeaderList,
@@ -150,11 +159,11 @@ public final class CronetDataSourceTest {
         null); // proxyServer
   }
 
-  @Test(expected = IllegalStateException.class)
+  @Test
   public void testOpeningTwiceThrows() throws HttpDataSourceException {
     mockResponseStartSuccess();
     dataSourceUnderTest.open(testDataSpec);
-    dataSourceUnderTest.open(testDataSpec);
+    assertThrows(IllegalStateException.class, () -> dataSourceUnderTest.open(testDataSpec));
   }
 
   @Test
@@ -650,6 +659,27 @@ public final class CronetDataSourceTest {
   }
 
   @Test
+  public void testRedirectParseAndAttachCookie() throws HttpDataSourceException {
+    mockSingleRedirectSuccess();
+
+    testResponseHeader.put("Set-Cookie", "testcookie=testcookie; Path=/video");
+    dataSourceUnderTest.open(testDataSpec);
+    verify(mockUrlRequestBuilder).addHeader(eq("Cookie"), any(String.class));
+    verify(mockUrlRequest, never()).followRedirect();
+    verify(mockUrlRequest, times(2)).start();
+  }
+
+  @Test
+  public void testRedirectNoSetCookieFollowsRedirect() throws HttpDataSourceException {
+    mockSingleRedirectSuccess();
+    mockFollowRedirectSuccess();
+
+    dataSourceUnderTest.open(testDataSpec);
+    verify(mockUrlRequestBuilder, never()).addHeader(eq("Cookie"), any(String.class));
+    verify(mockUrlRequest).followRedirect();
+  }
+
+  @Test
   public void testExceptionFromTransferListener() throws HttpDataSourceException {
     mockResponseStartSuccess();
 
@@ -731,6 +761,38 @@ public final class CronetDataSourceTest {
     }).when(mockUrlRequest).start();
   }
 
+  private void mockSingleRedirectSuccess() {
+    doAnswer(new Answer<Object>() {
+      @Override
+      public Object answer(InvocationOnMock invocation) throws Throwable {
+        if (!redirectCalled) {
+          redirectCalled = true;
+          dataSourceUnderTest.onRedirectReceived(
+              mockUrlRequest,
+              createUrlResponseInfoWithUrl("http://example.com/video", 300),
+              "http://example.com/video/redirect");
+        } else {
+          dataSourceUnderTest.onResponseStarted(
+              mockUrlRequest,
+              testUrlResponseInfo);
+        }
+        return null;
+      }
+    }).when(mockUrlRequest).start();
+  }
+
+  private void mockFollowRedirectSuccess() {
+    doAnswer(new Answer<Object>() {
+      @Override
+      public Object answer(InvocationOnMock invocation) throws Throwable {
+          dataSourceUnderTest.onResponseStarted(
+              mockUrlRequest,
+              testUrlResponseInfo);
+        return null;
+      }
+    }).when(mockUrlRequest).followRedirect();
+  }
+
   private void mockResponseStartFailure() {
     doAnswer(new Answer<Object>() {
       @Override

extensions/cronet/src/main/java/com/google/android/exoplayer2/ext/cronet/CronetDataSource.java

@@ -29,9 +29,13 @@ import com.google.android.exoplayer2.util.Assertions;
 import com.google.android.exoplayer2.util.Clock;
 import com.google.android.exoplayer2.util.Predicate;
 import java.io.IOException;
+import java.net.CookieManager;
 import java.net.SocketTimeoutException;
+import java.net.URI;
+import java.net.URISyntaxException;
 import java.net.UnknownHostException;
 import java.nio.ByteBuffer;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
@@ -93,6 +97,7 @@ public class CronetDataSource extends UrlRequest.Callback implements HttpDataSou
       Pattern.compile("^bytes (\\d+)-(\\d+)/(\\d+)$");
   // The size of read buffer passed to cronet UrlRequest.read().
   private static final int READ_BUFFER_SIZE_BYTES = 32 * 1024;
+  private static final String SET_COOKIE = "Set-Cookie";
 
   private final CronetEngine cronetEngine;
   private final Executor executor;
@@ -105,6 +110,7 @@ public class CronetDataSource extends UrlRequest.Callback implements HttpDataSou
   private final RequestProperties requestProperties;
   private final ConditionVariable operation;
   private final Clock clock;
+  private final CookieManager cookieManager;
 
   // Accessed by the calling thread only.
   private boolean opened;
@@ -144,7 +150,7 @@ public class CronetDataSource extends UrlRequest.Callback implements HttpDataSou
   public CronetDataSource(CronetEngine cronetEngine, Executor executor,
       Predicate<String> contentTypePredicate, TransferListener<? super CronetDataSource> listener) {
     this(cronetEngine, executor, contentTypePredicate, listener, DEFAULT_CONNECT_TIMEOUT_MILLIS,
-        DEFAULT_READ_TIMEOUT_MILLIS, false, null);
+        DEFAULT_READ_TIMEOUT_MILLIS, false, null, /* cookieManager= */ null);
   }
 
   /**
@@ -168,13 +174,42 @@ public class CronetDataSource extends UrlRequest.Callback implements HttpDataSou
       int connectTimeoutMs, int readTimeoutMs, boolean resetTimeoutOnRedirects,
       RequestProperties defaultRequestProperties) {
     this(cronetEngine, executor, contentTypePredicate, listener, connectTimeoutMs,
-        readTimeoutMs, resetTimeoutOnRedirects, Clock.DEFAULT, defaultRequestProperties);
+        readTimeoutMs, resetTimeoutOnRedirects, Clock.DEFAULT, defaultRequestProperties,
+        /* cookieManager= */ null);
+  }
+
+
+  /**
+   * @param cronetEngine A CronetEngine.
+   * @param executor The {@link java.util.concurrent.Executor} that will handle responses.
+   *     This may be a direct executor (i.e. executes tasks on the calling thread) in order
+   *     to avoid a thread hop from Cronet's internal network thread to the response handling
+   *     thread. However, to avoid slowing down overall network performance, care must be taken
+   *     to make sure response handling is a fast operation when using a direct executor.
+   * @param contentTypePredicate An optional {@link Predicate}. If a content type is rejected by the
+   *     predicate then an {@link InvalidContentTypeException} is thrown from
+   *     {@link #open(DataSpec)}.
+   * @param listener An optional listener.
+   * @param connectTimeoutMs The connection timeout, in milliseconds.
+   * @param readTimeoutMs The read timeout, in milliseconds.
+   * @param resetTimeoutOnRedirects Whether the connect timeout is reset when a redirect occurs.
+   * @param defaultRequestProperties The default request properties to be used.
+   * @param cookieManager An optional {@link CookieManager} to be used to handle "Set-Cookie"
+   *     requests. If this is null, then "Set-Cookie" requests will be ignored.
+   */
+  public CronetDataSource(CronetEngine cronetEngine, Executor executor,
+      Predicate<String> contentTypePredicate, TransferListener<? super CronetDataSource> listener,
+      int connectTimeoutMs, int readTimeoutMs, boolean resetTimeoutOnRedirects,
+      RequestProperties defaultRequestProperties, CookieManager cookieManager) {
+    this(cronetEngine, executor, contentTypePredicate, listener, connectTimeoutMs,
+        readTimeoutMs, resetTimeoutOnRedirects, Clock.DEFAULT, defaultRequestProperties,
+        cookieManager);
   }
 
   /* package */ CronetDataSource(CronetEngine cronetEngine, Executor executor,
       Predicate<String> contentTypePredicate, TransferListener<? super CronetDataSource> listener,
       int connectTimeoutMs, int readTimeoutMs, boolean resetTimeoutOnRedirects, Clock clock,
-      RequestProperties defaultRequestProperties) {
+      RequestProperties defaultRequestProperties, CookieManager cookieManager) {
     this.cronetEngine = Assertions.checkNotNull(cronetEngine);
     this.executor = Assertions.checkNotNull(executor);
     this.contentTypePredicate = contentTypePredicate;
@@ -184,6 +219,7 @@ public class CronetDataSource extends UrlRequest.Callback implements HttpDataSou
     this.resetTimeoutOnRedirects = resetTimeoutOnRedirects;
     this.clock = Assertions.checkNotNull(clock);
     this.defaultRequestProperties = defaultRequestProperties;
+    this.cookieManager = cookieManager;
     requestProperties = new RequestProperties();
     operation = new ConditionVariable();
   }
@@ -223,7 +259,12 @@ public class CronetDataSource extends UrlRequest.Callback implements HttpDataSou
     operation.close();
     resetConnectTimeout();
     currentDataSpec = dataSpec;
-    currentUrlRequest = buildRequest(dataSpec);
+
+    try {
+      currentUrlRequest = buildRequest(dataSpec);
+    } catch (IOException e) {
+      throw new OpenException(e, dataSpec, Status.IDLE);
+    }
     currentUrlRequest.start();
     boolean requestStarted = blockUntilConnectTimeout();
 
@@ -379,7 +420,24 @@ public class CronetDataSource extends UrlRequest.Callback implements HttpDataSou
     if (resetTimeoutOnRedirects) {
       resetConnectTimeout();
     }
-    request.followRedirect();
+
+    Map<String, List<String>> headers = info.getAllHeaders();
+    if (cookieManager == null || isEmpty(headers.get(SET_COOKIE))) {
+      request.followRedirect();
+    } else {
+      currentUrlRequest.cancel();
+      UrlRequest.Builder requestBuilder =
+          cronetEngine.newUrlRequestBuilder(newLocationUrl, this, executor).allowDirectExecutor();
+      try {
+        parseCookies(info);
+        attachCookies(newLocationUrl, requestBuilder);
+        currentUrlRequest = requestBuilder.build();
+        currentUrlRequest.start();
+      } catch (IOException e) {
+        exception = e;
+        operation.open();
+      }
+    }
   }
 
   @Override
@@ -387,7 +445,12 @@ public class CronetDataSource extends UrlRequest.Callback implements HttpDataSou
     if (request != currentUrlRequest) {
       return;
     }
-    responseInfo = info;
+    try {
+      parseCookies(info);
+      responseInfo = info;
+    } catch (IOException e) {
+      exception = e;
+    }
     operation.open();
   }
 
@@ -427,7 +490,7 @@ public class CronetDataSource extends UrlRequest.Callback implements HttpDataSou
 
   // Internal methods.
 
-  private UrlRequest buildRequest(DataSpec dataSpec) throws OpenException {
+  private UrlRequest buildRequest(DataSpec dataSpec) throws IOException {
     UrlRequest.Builder requestBuilder = cronetEngine.newUrlRequestBuilder(
         dataSpec.uri.toString(), this, executor).allowDirectExecutor();
     // Set the headers.
@@ -474,6 +537,7 @@ public class CronetDataSource extends UrlRequest.Callback implements HttpDataSou
             executor);
       }
     }
+    attachCookies(dataSpec.uri.toString(), requestBuilder);
     return requestBuilder.build();
   }
 
@@ -491,6 +555,37 @@ public class CronetDataSource extends UrlRequest.Callback implements HttpDataSou
     currentConnectTimeoutMs = clock.elapsedRealtime() + connectTimeoutMs;
   }
 
+  private void parseCookies(UrlResponseInfo info) throws IOException {
+    if (cookieManager == null) {
+      return;
+    }
+    try {
+      cookieManager.put(new URI(info.getUrl()), info.getAllHeaders());
+    } catch (URISyntaxException e) {
+      throw new IOException("Failed to parse cookies", e);
+    }
+  }
+
+  private void attachCookies(String url, UrlRequest.Builder requestBuilder) throws IOException {
+    if (cookieManager == null) {
+      return;
+    }
+    try {
+      for (Entry<String, List<String>> headers :
+          cookieManager.get(new URI(url), Collections.emptyMap()).entrySet()) {
+        StringBuilder cookies = new StringBuilder();
+        for (String cookie : headers.getValue()) {
+          cookies.append(cookie).append("; ");
+        }
+        if (cookies.length() > 0) {
+          requestBuilder.addHeader(headers.getKey(), cookies.substring(0, cookies.length() - 1));
+        }
+      }
+    } catch (URISyntaxException e) {
+      throw new IOException("Failed to attach cookies", e);
+    }
+  }
+
   private static boolean getIsCompressed(UrlResponseInfo info) {
     for (Map.Entry<String, String> entry : info.getAllHeadersAsList()) {
       if (entry.getKey().equalsIgnoreCase("Content-Encoding")) {