Fix a race condition in HttpEngineDataSource

Pass the temporary CookieHandler as a parameter instead of setting it as a temporary process-default. This avoids a rare race condition, where the player is sending a request with preserveCookies option and runs the `CookieHandler.getDefault()` before a different thread sets a default cookie handler. Over 5 or so lines of code the new default might be reverted to null - this is now fixed.

PiperOrigin-RevId: 738052152
(cherry picked from commit d0d76f2)

Author: Googler

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

@@ -956,32 +956,38 @@ private static int copyByteBuffer(ByteBuffer src, ByteBuffer dst) {
     return remaining;
   }
 
-  /**
-   * 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 .
-   */
-  @VisibleForTesting
-  /* private */ 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);
-      }
+  // Stores the cookie headers from the response in the default {@link CookieHandler}.
+  private static void storeCookiesFromHeaders(UrlResponseInfo info) {
+    storeCookiesFromHeaders(info, CookieHandler.getDefault());
+  }
+
+  // Stores the cookie headers from the response in the provided {@link CookieHandler}.
+  private static void storeCookiesFromHeaders(
+      UrlResponseInfo info, @Nullable CookieHandler cookieHandler) {
+    if (cookieHandler == null) {
+      return;
+    }
+
+    try {
+      cookieHandler.put(new URI(info.getUrl()), info.getHeaders().getAsMap());
+    } catch (Exception e) {
+      Log.w(TAG, "Failed to store cookies in CookieHandler", e);
     }
   }
 
   @VisibleForTesting
   /* private */ static String getCookieHeader(String url) {
-    return getCookieHeader(url, ImmutableMap.of());
+    return getCookieHeader(url, ImmutableMap.of(), CookieHandler.getDefault());
   }
 
-  // getCookieHeader maps Set-Cookie2 (RFC 2965) to Cookie just like CookieManager does.
   @VisibleForTesting
-  /* private */ static String getCookieHeader(String url, Map<String, List<String>> headers) {
-    CookieHandler cookieHandler = CookieHandler.getDefault();
+  /* private */ static String getCookieHeader(String url, @Nullable CookieHandler cookieHandler) {
+    return getCookieHeader(url, ImmutableMap.of(), cookieHandler);
+  }
+
+  // getCookieHeader maps Set-Cookie2 (RFC 2965) to Cookie just like CookieManager does.
+  private static String getCookieHeader(
+      String url, Map<String, List<String>> headers, @Nullable CookieHandler cookieHandler) {
     if (cookieHandler == null) {
       return "";
     }
@@ -1059,11 +1065,6 @@ 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) {
@@ -1092,20 +1093,17 @@ public synchronized void onRedirectReceived(
         resetConnectTimeout();
       }
 
-      boolean hasDefaultCookieHandler = CookieHandler.getDefault() != null;
-      if (!hasDefaultCookieHandler && handleSetCookieRequests) {
+      CookieHandler cookieHandler = CookieHandler.getDefault();
+
+      if (cookieHandler == null && 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);
+        cookieHandler = new CookieManager();
       }
 
-      storeCookiesFromHeaders(info);
-      String cookieHeaders = getCookieHeader(info.getUrl(), info.getHeaders().getAsMap());
-
-      if (!hasDefaultCookieHandler && handleSetCookieRequests) {
-        resetDefaultCookieHandler();
-      }
+      storeCookiesFromHeaders(info, cookieHandler);
+      String cookieHeaders =
+          getCookieHeader(info.getUrl(), info.getHeaders().getAsMap(), cookieHandler);
 
       boolean shouldKeepPost =
           keepPostFor302Redirects

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

@@ -1486,6 +1486,18 @@ public void getCookieHeader_cookieHandler() throws Exception {
         .isEqualTo(TEST_REQUEST_COOKIE + "; " + TEST_REQUEST_COOKIE_2 + ";");
   }
 
+  @Test
+  public void getCookieHeader_cookieHandlerCustomHandler() 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)));
+
+    assertThat(HttpEngineDataSource.getCookieHeader(TEST_URL, cm))
+        .isEqualTo(TEST_REQUEST_COOKIE + "; " + TEST_REQUEST_COOKIE_2 + ";");
+  }
+
   @Test
   public void getCookieHeader_cookieHandlerCookie2() throws Exception {
     CookieManager cm = new CookieManager();