Localized cookie handling

Bug: 177562073
Test: manual - VoWiFi enetitlement success
Test: unit test pass
Change-Id: Ie544e17026feaf6a0df6377061310fa2fe9be59f

5 files changed, 46 insertions, 31 deletions

diff --git a/java/com/android/libraries/entitlement/eapaka/EapAkaApi.java b/java/com/android/libraries/entitlement/eapaka/EapAkaApi.java
index 049ae11..4ad0332 100644
--- a/java/com/android/libraries/entitlement/eapaka/EapAkaApi.java
+++ b/java/com/android/libraries/entitlement/eapaka/EapAkaApi.java
@@ -41,9 +41,6 @@ import com.google.common.net.HttpHeaders;
 import org.json.JSONException;
 import org.json.JSONObject;
 
-import java.net.CookieHandler;
-import java.net.CookieManager;
-
 public class EapAkaApi {
     private static final String TAG = "ServiceEntitlement";
 
@@ -138,9 +135,6 @@ public class EapAkaApi {
     public String queryEntitlementStatus(
             ImmutableList<String> appIds, String serverUrl, ServiceEntitlementRequest request)
             throws ServiceEntitlementException {
-        // TODO(b/177562073): localize cookie management instead of VM global CookieHandler
-        CookieHandler.setDefault(new CookieManager());
-
         HttpRequest httpRequest =
                 HttpRequest.builder()
                         .setUrl(entitlementStatusUrl(appIds, serverUrl, request))
@@ -166,7 +160,7 @@ public class EapAkaApi {
             }
             return challengeResponse(
                     new EapAkaResponse(responseData).getEapAkaChallengeResponse(mContext,
-                            mSimSubscriptionId), serverUrl);
+                            mSimSubscriptionId), serverUrl, response.cookie());
         } else {
             // Result of fast AuthN
             Log.d(TAG, "fast AuthN");
@@ -174,7 +168,7 @@ public class EapAkaApi {
         }
     }
 
-    private String challengeResponse(String akaChallengeResponse, String serverUrl)
+    private String challengeResponse(String akaChallengeResponse, String serverUrl, String cookie)
             throws ServiceEntitlementException {
         Log.d(TAG, "challengeResponse");
         JSONObject postData = new JSONObject();
@@ -191,6 +185,7 @@ public class EapAkaApi {
                         .setPostData(postData)
                         .addRequestProperty(HttpHeaders.ACCEPT, ACCEPT_CONTENT_TYPE_JSON_AND_XML)
                         .addRequestProperty(HttpHeaders.CONTENT_TYPE, REQUEST_CONTENT_TYPE_JSON)
+                        .addRequestProperty(HttpHeaders.COOKIE, cookie)
                         .build();
         return mHttpClient.request(request).body();
     }
diff --git a/java/com/android/libraries/entitlement/http/HttpClient.java b/java/com/android/libraries/entitlement/http/HttpClient.java
index 4a3bdaf..1144c5a 100644
--- a/java/com/android/libraries/entitlement/http/HttpClient.java
+++ b/java/com/android/libraries/entitlement/http/HttpClient.java
@@ -130,6 +130,7 @@ public class HttpClient {
             throw new ServiceEntitlementException(
                     ERROR_HTTP_STATUS_NOT_SUCCESS, "Read response code failed!", e);
         }
+        responseBuilder.setCookie(nullToEmpty(getCookie(connection)));
         try {
             String responseBody = readResponse(connection);
             logPii("HttpClient.response body: " + responseBody);
@@ -163,4 +164,8 @@ public class HttpClient {
         }
         return ContentType.UNKNOWN;
     }
+
+    private static String getCookie(URLConnection connection) {
+        return connection.getHeaderField(HttpHeaders.SET_COOKIE);
+    }
 }
diff --git a/java/com/android/libraries/entitlement/http/HttpRequest.java b/java/com/android/libraries/entitlement/http/HttpRequest.java
index 733355f..4a56a60 100644
--- a/java/com/android/libraries/entitlement/http/HttpRequest.java
+++ b/java/com/android/libraries/entitlement/http/HttpRequest.java
@@ -45,7 +45,6 @@ public abstract class HttpRequest {
      */
     @AutoValue.Builder
     public abstract static class Builder {
-
         private final Map<String, String> values = new ArrayMap<>();
         private final Map<String, String> properties = new ArrayMap<>();
 
diff --git a/java/com/android/libraries/entitlement/http/HttpResponse.java b/java/com/android/libraries/entitlement/http/HttpResponse.java
index b6ca35f..e331b99 100644
--- a/java/com/android/libraries/entitlement/http/HttpResponse.java
+++ b/java/com/android/libraries/entitlement/http/HttpResponse.java
@@ -37,11 +37,15 @@ public abstract class HttpResponse {
     public abstract String responseMessage();
 
     /**
+     * Content of the "Set-Cookie" response header.
+     */
+    public abstract String cookie();
+
+    /**
      * Builder of {@link HttpResponse}.
      */
     @AutoValue.Builder
     public abstract static class Builder {
-
         public abstract HttpResponse build();
 
         public abstract Builder setContentType(int contentType);
@@ -51,6 +55,11 @@ public abstract class HttpResponse {
         public abstract Builder setResponseCode(int responseCode);
 
         public abstract Builder setResponseMessage(String responseMessage);
+
+        /**
+         * Content of the "Set-Cookie" response header.
+         */
+        public abstract Builder setCookie(String cookie);
     }
 
     public static Builder builder() {
@@ -58,7 +67,8 @@ public abstract class HttpResponse {
                 .setContentType(ContentType.UNKNOWN)
                 .setBody("")
                 .setResponseCode(0)
-                .setResponseMessage("");
+                .setResponseMessage("")
+                .setCookie("");
     }
 
     @Override
@@ -73,7 +83,9 @@ public abstract class HttpResponse {
                 .append(responseCode())
                 .append(" responseMessage=")
                 .append(responseMessage())
-                .append("}")
+                .append(" cookie=(")
+                .append(cookie().length())
+                .append(" characters)}")
                 .toString();
     }
 }
diff --git a/tests/src/com/android/libraries/entitlement/eapaka/EapAkaApiTest.java b/tests/src/com/android/libraries/entitlement/eapaka/EapAkaApiTest.java
index 2609c69..4d4f49f 100644
--- a/tests/src/com/android/libraries/entitlement/eapaka/EapAkaApiTest.java
+++ b/tests/src/com/android/libraries/entitlement/eapaka/EapAkaApiTest.java
@@ -20,6 +20,7 @@ import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 import static org.testng.Assert.fail;
@@ -34,13 +35,11 @@ import com.android.libraries.entitlement.ServiceEntitlement;
 import com.android.libraries.entitlement.ServiceEntitlementException;
 import com.android.libraries.entitlement.ServiceEntitlementRequest;
 import com.android.libraries.entitlement.http.HttpClient;
-import com.android.libraries.entitlement.http.HttpConstants;
 import com.android.libraries.entitlement.http.HttpConstants.ContentType;
 import com.android.libraries.entitlement.http.HttpRequest;
 import com.android.libraries.entitlement.http.HttpResponse;
 
 import com.google.common.collect.ImmutableList;
-import com.google.common.net.HttpHeaders;
 
 import org.json.JSONException;
 import org.junit.Before;
@@ -61,6 +60,9 @@ public class EapAkaApiTest {
                     + "AQIAfBcBAAABBQAAXOZSkCjxysgE4"
                     + "3GWqHJvgQIFAABrikWGrekAALNU4TxmCDPoCwUAAJT0nqXeAYlqzT0UGXINENWBBQAA7z3fhImk"
                     + "q+vcCKWIZBdvuYIJAAAPRUFp7KWFo+Thr78Qj9hEkB2zA0i6KakODsufBC+BJQ==\"}";
+    // com.google.common.net.HttpHeaders.COOKIE
+    private static final String HTTP_HEADER_COOKIE = "Cookie";
+    private static final String COOKIE_VALUE = "COOKIE=abcdefg";
     private static final String GSM_SECURITY_CONTEXT_REQUEST =
             "EFzmUpAo8crIBONxlqhyb4EQa4pFhq3pAACzVOE8Zggz6A==";
     private static final String GSM_SECURITY_CONTEXT_RESPONSE =
@@ -111,22 +113,24 @@ public class EapAkaApiTest {
     }
 
     @Test
-    public void queryEntitlementStatus_hasAuthenticationToken_fastAuthN() throws Exception {
+    public void queryEntitlementStatus_hasAuthenticationToken() throws Exception {
         HttpResponse response =
                 HttpResponse.builder().setContentType(ContentType.XML).setBody(RESPONSE_XML)
                         .build();
         when(mMockHttpClient.request(any())).thenReturn(response);
-
         ServiceEntitlementRequest request =
                 ServiceEntitlementRequest.builder().setAuthenticationToken(TOKEN).build();
 
-        assertThat(mEapAkaApi.queryEntitlementStatus(
-                ImmutableList.of(ServiceEntitlement.APP_VOWIFI), TEST_URL, request))
-                .isEqualTo(RESPONSE_XML);
+        String respopnse =
+                mEapAkaApi.queryEntitlementStatus(
+                        ImmutableList.of(ServiceEntitlement.APP_VOWIFI), TEST_URL, request);
+
+        assertThat(respopnse).isEqualTo(RESPONSE_XML);
+        verify(mMockHttpClient).request(mHttpRequestCaptor.capture());
     }
 
     @Test
-    public void queryEntitlementStatus_noAuthenticationToken_initialAuthN() throws Exception {
+    public void queryEntitlementStatus_noAuthenticationToken() throws Exception {
         when(mMockTelephonyManagerForSubId.getSubscriberId()).thenReturn(IMSI);
         when(mMockTelephonyManagerForSubId.getSimOperator()).thenReturn(MCCMNC);
         when(mMockTelephonyManagerForSubId.getIccAuthentication(
@@ -134,24 +138,26 @@ public class EapAkaApiTest {
                 TelephonyManager.AUTHTYPE_EAP_AKA,
                 GSM_SECURITY_CONTEXT_REQUEST))
                 .thenReturn(GSM_SECURITY_CONTEXT_RESPONSE);
-
         HttpResponse eapChallengeResponse =
                 HttpResponse
                         .builder().setContentType(ContentType.JSON).setBody(EAP_AKA_CHALLENGE)
-                        .build();
+                        .setCookie(COOKIE_VALUE).build();
         HttpResponse xmlResponse =
                 HttpResponse.builder().setContentType(ContentType.XML).setBody(RESPONSE_XML)
                         .build();
         when(mMockHttpClient.request(any()))
                 .thenReturn(eapChallengeResponse).thenReturn(xmlResponse);
-
         ServiceEntitlementRequest request = ServiceEntitlementRequest.builder().build();
 
-        assertThat(
-                mEapAkaApi.queryEntitlementStatus(ImmutableList.of(ServiceEntitlement.APP_VOWIFI),
-                        TEST_URL,
-                        request))
-                .isEqualTo(RESPONSE_XML);
+        String respopnse =
+                mEapAkaApi.queryEntitlementStatus(
+                        ImmutableList.of(ServiceEntitlement.APP_VOWIFI), TEST_URL, request);
+
+        assertThat(respopnse).isEqualTo(RESPONSE_XML);
+        // Verify that the 2nd request has cookie set by the 1st response
+        verify(mMockHttpClient, times(2)).request(mHttpRequestCaptor.capture());
+        assertThat(mHttpRequestCaptor.getAllValues().get(1).requestProperties())
+                .containsEntry(HTTP_HEADER_COOKIE, COOKIE_VALUE);
     }
 
     @Test
@@ -183,7 +189,6 @@ public class EapAkaApiTest {
                 HttpResponse.builder().setContentType(ContentType.XML).setBody(RESPONSE_XML)
                         .build();
         when(mMockHttpClient.request(any())).thenReturn(xmlResponse);
-
         ServiceEntitlementRequest request = ServiceEntitlementRequest.builder().build();
 
         try {
@@ -204,7 +209,6 @@ public class EapAkaApiTest {
         HttpResponse eapChallengeResponse =
                 HttpResponse.builder().setContentType(ContentType.JSON).build();
         when(mMockHttpClient.request(any())).thenReturn(eapChallengeResponse);
-
         ServiceEntitlementRequest request = ServiceEntitlementRequest.builder().build();
 
         try {
@@ -219,4 +223,4 @@ public class EapAkaApiTest {
             assertThat(exception.getCause()).isInstanceOf(JSONException.class);
         }
     }
-}
\ No newline at end of file
+}