Add CookieHandler support to HttpEngineDataSource.

The original behaviour of the `handleSetCookieRequests` is preserved.
A bug is addressed, where it woulld set `Set-Cookie` header in the client request (as opposed to the expected `Cookies`).

PiperOrigin-RevId: 730857996

libraries/datasource/src/main/java/androidx/media3/datasource/HttpEngineDataSource.java

@@ -36,16 +36,21 @@ import androidx.media3.common.PlaybackException;
 import androidx.media3.common.util.Assertions;
 import androidx.media3.common.util.Clock;
 import androidx.media3.common.util.ConditionVariable;
+import androidx.media3.common.util.Log;
 import androidx.media3.common.util.UnstableApi;
 import androidx.media3.common.util.Util;
 import com.google.common.base.Ascii;
 import com.google.common.base.Predicate;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.net.HttpHeaders;
 import com.google.common.primitives.Longs;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import java.io.IOException;
 import java.io.InterruptedIOException;
+import java.net.CookieHandler;
+import java.net.CookieManager;
 import java.net.SocketTimeoutException;
+import java.net.URI;
 import java.net.UnknownHostException;
 import java.nio.ByteBuffer;
 import java.util.Arrays;
@@ -932,14 +937,6 @@ public final class HttpEngineDataSource extends BaseDataSource implements HttpDa
     return false;
   }
 
-  @Nullable
-  private static String parseCookies(@Nullable List<String> setCookieHeaders) {
-    if (setCookieHeaders == null || setCookieHeaders.isEmpty()) {
-      return null;
-    }
-    return TextUtils.join(";", setCookieHeaders);
-  }
-
   @Nullable
   private static String getFirstHeader(Map<String, List<String>> allHeaders, String headerName) {
     @Nullable List<String> headers = allHeaders.get(headerName);
@@ -961,7 +958,8 @@ public final class HttpEngineDataSource extends BaseDataSource implements HttpDa
    * A wrapper class that manages a {@link UrlRequest} and the {@link UrlRequestCallback} associated
    * with that request.
    */
-  private static final class UrlRequestWrapper {
+  @VisibleForTesting
+  /* private */ static final class UrlRequestWrapper {
 
     private final UrlRequest urlRequest;
     private final UrlRequestCallback urlRequestCallback;
@@ -1004,14 +1002,21 @@ public final class HttpEngineDataSource extends BaseDataSource implements HttpDa
     }
   }
 
-  private final class UrlRequestCallback implements UrlRequest.Callback {
+  @VisibleForTesting
+  /* private */ final class UrlRequestCallback implements UrlRequest.Callback {
 
+    private static final String TAG = "HttpEngineDataSource";
     private volatile boolean isClosed = false;
 
     public void close() {
       this.isClosed = true;
     }
 
+    @SuppressWarnings("argument.type.incompatible")
+    private void resetDefaultCookieHandler() {
+      CookieHandler.setDefault(null);
+    }
+
     @Override
     public synchronized void onRedirectReceived(
         UrlRequest request, UrlResponseInfo info, String newLocationUrl) {
@@ -1040,6 +1045,21 @@ public final class HttpEngineDataSource extends BaseDataSource implements HttpDa
         resetConnectTimeout();
       }
 
+      boolean hasDefaultCookieHandler = CookieHandler.getDefault() != null;
+      if (!hasDefaultCookieHandler && handleSetCookieRequests) {
+        // a temporary CookieManager is created for the duration of this request - this guarantees
+        // redirects preserve the cookies correctly.
+        CookieManager cm = new CookieManager();
+        CookieHandler.setDefault(cm);
+      }
+
+      storeCookiesFromHeaders(info);
+      String cookieHeaders = getCookieHeader(info.getUrl(), info.getHeaders().getAsMap());
+
+      if (!hasDefaultCookieHandler && handleSetCookieRequests) {
+        resetDefaultCookieHandler();
+      }
+
       boolean shouldKeepPost =
           keepPostFor302Redirects
               && dataSpec.httpMethod == DataSpec.HTTP_METHOD_POST
@@ -1047,17 +1067,12 @@ public final class HttpEngineDataSource extends BaseDataSource implements HttpDa
 
       // request.followRedirect() transforms a POST request into a GET request, so if we want to
       // keep it as a POST we need to fall through to the manual redirect logic below.
-      if (!shouldKeepPost && !handleSetCookieRequests) {
+      if (!shouldKeepPost) {
+        // No cookies, or we're not handling them - so just follow the redirect.
+        if (!handleSetCookieRequests || TextUtils.isEmpty(cookieHeaders)) {
           request.followRedirect();
           return;
         }
-
-      @Nullable
-      String cookieHeadersValue =
-          parseCookies(info.getHeaders().getAsMap().get(HttpHeaders.SET_COOKIE));
-      if (!shouldKeepPost && TextUtils.isEmpty(cookieHeadersValue)) {
-        request.followRedirect();
-        return;
       }
 
       request.cancel();
@@ -1075,13 +1090,15 @@ public final class HttpEngineDataSource extends BaseDataSource implements HttpDa
       } else {
         redirectUrlDataSpec = dataSpec.withUri(Uri.parse(newLocationUrl));
       }
-      if (!TextUtils.isEmpty(cookieHeadersValue)) {
+
+      if (!TextUtils.isEmpty(cookieHeaders)) {
         Map<String, String> requestHeaders = new HashMap<>();
         requestHeaders.putAll(dataSpec.httpRequestHeaders);
-        requestHeaders.put(HttpHeaders.COOKIE, cookieHeadersValue);
+        requestHeaders.put(HttpHeaders.COOKIE, cookieHeaders);
         redirectUrlDataSpec =
             redirectUrlDataSpec.buildUpon().setHttpRequestHeaders(requestHeaders).build();
       }
+
       UrlRequestWrapper redirectUrlRequestWrapper;
       try {
         redirectUrlRequestWrapper = buildRequestWrapper(redirectUrlDataSpec);
@@ -1096,11 +1113,56 @@ public final class HttpEngineDataSource extends BaseDataSource implements HttpDa
       currentUrlRequestWrapper.start();
     }
 
+    /**
+     * Stores the cookie headers from the response in the default {@link CookieHandler}.
+     *
+     * <p>This is a no-op if the {@link CookieHandler} is not set .
+     */
+    static void storeCookiesFromHeaders(UrlResponseInfo info) {
+      CookieHandler cookieHandler = CookieHandler.getDefault();
+      if (cookieHandler != null) {
+        try {
+          cookieHandler.put(new URI(info.getUrl()), info.getHeaders().getAsMap());
+        } catch (Exception e) {
+          Log.w(TAG, "Failed to store cookies in CookieHandler", e);
+        }
+      }
+    }
+
+    static String getCookieHeader(String url) {
+      return getCookieHeader(url, ImmutableMap.of());
+    }
+
+    // getCookieHeader maps Set-Cookie2 (RFC 2965) to Cookie just like CookieManager does.
+    static String getCookieHeader(String url, Map<String, List<String>> headers) {
+      CookieHandler cookieHandler = CookieHandler.getDefault();
+      if (cookieHandler == null) {
+        return "";
+      }
+
+      Map<String, List<String>> cookieHeaders = ImmutableMap.of();
+      try {
+        cookieHeaders = cookieHandler.get(new URI(url), headers);
+      } catch (Exception e) {
+        Log.w(TAG, "Failed to read cookies from CookieHandler", e);
+      }
+
+      StringBuilder cookies = new StringBuilder();
+      if (cookieHeaders.containsKey(HttpHeaders.COOKIE)) {
+        List<String> cookiesList = cookieHeaders.get(HttpHeaders.COOKIE);
+        for (String cookie : cookiesList) {
+          cookies.append(cookie).append("; ");
+        }
+      }
+      return cookies.toString().stripTrailing();
+    }
+
     @Override
     public synchronized void onResponseStarted(UrlRequest request, UrlResponseInfo info) {
       if (isClosed) {
         return;
       }
+      storeCookiesFromHeaders(info);
       responseInfo = info;
       operation.open();
     }

libraries/datasource/src/test/java/androidx/media3/datasource/HttpEngineDataSourceTest.java

@@ -45,9 +45,14 @@ import androidx.media3.common.util.Util;
 import androidx.media3.datasource.HttpDataSource.HttpDataSourceException;
 import androidx.media3.datasource.HttpDataSource.InvalidResponseCodeException;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import java.io.IOException;
 import java.io.InterruptedIOException;
+import java.net.CookieHandler;
+import java.net.CookieManager;
 import java.net.SocketTimeoutException;
+import java.net.URI;
 import java.net.UnknownHostException;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
@@ -80,8 +85,14 @@ public final class HttpEngineDataSourceTest {
 
   private static final int TEST_CONNECT_TIMEOUT_MS = 100;
   private static final int TEST_READ_TIMEOUT_MS = 100;
-  private static final String TEST_URL = "http://google.com";
+  private static final String TEST_URL = "http://google.com/video/";
   private static final String TEST_CONTENT_TYPE = "test/test";
+  private static final String TEST_REQUEST_COOKIE = "foo=bar";
+  private static final String TEST_REQUEST_COOKIE_2 = "baz=qux";
+  private static final String TEST_RESPONSE_SET_COOKIE =
+      TEST_REQUEST_COOKIE + ";path=/video; expires 31-12-2099 23:59:59 GMT";
+  private static final String TEST_RESPONSE_SET_COOKIE_2 =
+      TEST_REQUEST_COOKIE_2 + ";path=/; expires 31-12-2099 23:59:59 GMT";
   private static final byte[] TEST_POST_BODY = Util.getUtf8Bytes("test post body");
   private static final long TEST_CONTENT_LENGTH = 16000L;
 
@@ -141,6 +152,8 @@ public final class HttpEngineDataSourceTest {
     // This value can be anything since the DataSpec is unset.
     testResponseHeader.put("Content-Length", Long.toString(TEST_CONTENT_LENGTH));
     testUrlResponseInfo = createUrlResponseInfo(/* statusCode= */ 200);
+
+    CookieHandler.setDefault(null);
   }
 
   @After
@@ -272,15 +285,15 @@ public final class HttpEngineDataSourceTest {
   @Test
   public void requestHeadersSet() throws HttpDataSourceException {
     Map<String, String> headersSet = new HashMap<>();
-    doAnswer(
+    when(mockUrlRequestBuilder.addHeader(
+            ArgumentMatchers.anyString(), ArgumentMatchers.anyString()))
+        .thenAnswer(
             (invocation) -> {
               String key = invocation.getArgument(0);
               String value = invocation.getArgument(1);
               headersSet.put(key, value);
               return null;
-            })
+            });
-        .when(mockUrlRequestBuilder)
-        .addHeader(ArgumentMatchers.anyString(), ArgumentMatchers.anyString());
 
     dataSourceUnderTest.setRequestProperty("defaultHeader2", "dataSourceOverridesDefault");
     dataSourceUnderTest.setRequestProperty("dataSourceHeader1", "dataSourceValue1");
@@ -447,8 +460,7 @@ public final class HttpEngineDataSourceTest {
       assertThat(e).isInstanceOf(HttpDataSource.InvalidContentTypeException.class);
       // Check for connection not automatically closed.
       verify(mockUrlRequest, never()).cancel();
-      assertThat(testedContentTypes).hasSize(1);
+      assertThat(testedContentTypes).containsExactly(TEST_CONTENT_TYPE);
-      assertThat(testedContentTypes.get(0)).isEqualTo(TEST_CONTENT_TYPE);
     }
   }
 
@@ -1277,7 +1289,7 @@ public final class HttpEngineDataSourceTest {
                 .createDataSource();
     mockSingleRedirectSuccess(/* responseCode= */ 302);
     dataSourceUnderTest.setRequestProperty("Content-Type", TEST_CONTENT_TYPE);
-    testResponseHeader.put("Set-Cookie", "testcookie=testcookie; Path=/video");
+    testResponseHeader.put("Set-Cookie", TEST_RESPONSE_SET_COOKIE);
 
     dataSourceUnderTest.open(testPostDataSpec);
 
@@ -1449,6 +1461,52 @@ public final class HttpEngineDataSourceTest {
     verify(mockUrlRequestBuilder).setDirectExecutorAllowed(true);
   }
 
+  @Test
+  public void getCookieHeader_noCookieHandler() {
+    assertThat(HttpEngineDataSource.UrlRequestCallback.getCookieHeader(TEST_URL)).isEmpty();
+    assertThat(CookieHandler.getDefault()).isNull();
+  }
+
+  @Test
+  public void getCookieHeader_emptyCookieHandler() {
+    CookieHandler.setDefault(new CookieManager());
+    assertThat(HttpEngineDataSource.UrlRequestCallback.getCookieHeader(TEST_URL)).isEmpty();
+  }
+
+  @Test
+  public void getCookieHeader_cookieHandler() throws Exception {
+    CookieManager cm = new CookieManager();
+    cm.put(
+        new URI(TEST_URL),
+        ImmutableMap.of(
+            "Set-Cookie", ImmutableList.of(TEST_RESPONSE_SET_COOKIE, TEST_RESPONSE_SET_COOKIE_2)));
+    CookieHandler.setDefault(cm);
+
+    assertThat(HttpEngineDataSource.UrlRequestCallback.getCookieHeader(TEST_URL))
+        .isEqualTo(TEST_REQUEST_COOKIE + "; " + TEST_REQUEST_COOKIE_2 + ";");
+  }
+
+  @Test
+  public void getCookieHeader_cookieHandlerCookie2() throws Exception {
+    CookieManager cm = new CookieManager();
+    cm.put(
+        new URI(TEST_URL),
+        ImmutableMap.of(
+            "Set-Cookie2", ImmutableList.of(TEST_RESPONSE_SET_COOKIE, TEST_RESPONSE_SET_COOKIE_2)));
+    CookieHandler.setDefault(cm);
+
+    // This asserts the surprising behavior of CookieManager - Set-Cookie2 is translated to Cookie,
+    // not Cookie2.
+    assertThat(cm.get(new URI(TEST_URL), ImmutableMap.of("", ImmutableList.of()))).isNotEmpty();
+    assertThat(cm.get(new URI(TEST_URL), ImmutableMap.of("", ImmutableList.of())).get("Cookie"))
+        .containsExactly(TEST_REQUEST_COOKIE, TEST_REQUEST_COOKIE_2);
+    assertThat(cm.get(new URI(TEST_URL), ImmutableMap.of("", ImmutableList.of())))
+        .doesNotContainKey("Cookie2");
+
+    assertThat(HttpEngineDataSource.UrlRequestCallback.getCookieHeader(TEST_URL))
+        .isEqualTo(TEST_REQUEST_COOKIE + "; " + TEST_REQUEST_COOKIE_2 + ";");
+  }
+
   // Helper methods.
 
   private void mockStatusResponse() {