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: 732945338
(cherry picked from commit 814d368d9f54f47f6f0372c7905846104c549322)
This commit is contained in:
Googler 2025-03-03 09:08:22 -08:00 committed by tonihei
parent 53a3144f96
commit 25dc53551d
2 changed files with 156 additions and 33 deletions

View File

@ -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;
@ -321,6 +326,8 @@ public final class HttpEngineDataSource extends BaseDataSource implements HttpDa
// The size of read buffer passed to cronet UrlRequest.read().
private static final int READ_BUFFER_SIZE_BYTES = 32 * 1024;
private static final String TAG = "HttpEngineDataSource";
private final HttpEngine httpEngine;
private final Executor executor;
private final int requestPriority;
@ -709,7 +716,7 @@ public final class HttpEngineDataSource extends BaseDataSource implements HttpDa
@UnstableApi
@VisibleForTesting
@Nullable
UrlRequest.Callback getCurrentUrlRequestCallback() {
UrlRequestCallback getCurrentUrlRequestCallback() {
return currentUrlRequestWrapper == null
? null
: currentUrlRequestWrapper.getUrlRequestCallback();
@ -932,14 +939,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);
@ -957,6 +956,55 @@ public final class HttpEngineDataSource extends BaseDataSource implements HttpDa
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);
}
}
}
@VisibleForTesting
/* private */ static String getCookieHeader(String url) {
return getCookieHeader(url, ImmutableMap.of());
}
// 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();
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);
if (cookiesList != null) {
for (String cookie : cookiesList) {
cookies.append(cookie).append("; ");
}
}
}
return cookies.toString().stripTrailing();
}
/**
* A wrapper class that manages a {@link UrlRequest} and the {@link UrlRequestCallback} associated
* with that request.
@ -984,7 +1032,7 @@ public final class HttpEngineDataSource extends BaseDataSource implements HttpDa
urlRequest.cancel();
}
public UrlRequest.Callback getUrlRequestCallback() {
public UrlRequestCallback getUrlRequestCallback() {
return urlRequestCallback;
}
@ -1004,14 +1052,18 @@ public final class HttpEngineDataSource extends BaseDataSource implements HttpDa
}
}
private final class UrlRequestCallback implements UrlRequest.Callback {
final class UrlRequestCallback implements UrlRequest.Callback {
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 +1092,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 +1114,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) {
request.followRedirect();
return;
}
@Nullable
String cookieHeadersValue =
parseCookies(info.getHeaders().getAsMap().get(HttpHeaders.SET_COOKIE));
if (!shouldKeepPost && TextUtils.isEmpty(cookieHeadersValue)) {
request.followRedirect();
return;
if (!shouldKeepPost) {
// No cookies, or we're not handling them - so just follow the redirect.
if (!handleSetCookieRequests || TextUtils.isEmpty(cookieHeaders)) {
request.followRedirect();
return;
}
}
request.cancel();
@ -1075,13 +1137,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);
@ -1101,6 +1165,7 @@ public final class HttpEngineDataSource extends BaseDataSource implements HttpDa
if (isClosed) {
return;
}
storeCookiesFromHeaders(info);
responseInfo = info;
operation.open();
}

View File

@ -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.get(0)).isEqualTo(TEST_CONTENT_TYPE);
assertThat(testedContentTypes).containsExactly(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.getCookieHeader(TEST_URL)).isEmpty();
assertThat(CookieHandler.getDefault()).isNull();
}
@Test
public void getCookieHeader_emptyCookieHandler() {
CookieHandler.setDefault(new CookieManager());
assertThat(HttpEngineDataSource.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.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.getCookieHeader(TEST_URL))
.isEqualTo(TEST_REQUEST_COOKIE + "; " + TEST_REQUEST_COOKIE_2 + ";");
}
// Helper methods.
private void mockStatusResponse() {