chore: zaas-client support multiple set-cookie headers

Signed-off-by: Richard Salac <richard.salac@broadcom.com>
zowe · Sep 11, 2024 · 42e1d56 · 42e1d56
1 parent a330907
commit 42e1d56
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 15 deletions.
diff --git a/zaas-client/src/main/java/org/zowe/apiml/zaasclient/service/internal/ZaasJwtService.java b/zaas-client/src/main/java/org/zowe/apiml/zaasclient/service/internal/ZaasJwtService.java
@@ -38,10 +38,9 @@
 
 import java.io.IOException;
 import java.util.Arrays;
-import java.util.Map;
 import java.util.Optional;
 import java.util.function.Predicate;
-import java.util.stream.Stream;
+import java.util.stream.Collectors;
 
 @Slf4j
 class ZaasJwtService implements TokenService {
@@ -93,9 +92,7 @@ private ClassicHttpRequest loginWithCredentials(String userId, char[] password,
 
     private SimpleHttpResponse processJwtTokenResponse(ClassicHttpResponse response) throws ParseException, IOException {
         var headers = Arrays.stream(response.getHeaders(HttpHeaders.SET_COOKIE))
-            .findFirst()
-            .map(header -> Map.of(HttpHeaders.SET_COOKIE, header))
-            .orElse(Map.of());
+            .collect(Collectors.groupingBy((__) -> HttpHeaders.SET_COOKIE));
         if (response.getEntity() != null) {
             return new SimpleHttpResponse(response.getCode(), EntityUtils.toString(response.getEntity()), headers);
         } else {
@@ -245,8 +242,8 @@ private String extractToken(SimpleHttpResponse response) throws ZaasClientExcept
         String token = "";
         int httpResponseCode = response.getCode();
         if (httpResponseCode == 204) {
-            var vals = response.getHeaders().get(HttpHeaders.SET_COOKIE).getValue().split(";");
-            var apimlAuthCookie = Stream.of(vals).filter(v -> v.startsWith(zassConfigProperties.getTokenPrefix())).map(v -> v.substring(v.indexOf("=") + 1)).findFirst();
+            var cookies = response.getHeaders().get(HttpHeaders.SET_COOKIE).stream().map(header -> header.getValue().split(";")).flatMap(Arrays::stream).toList();
+            var apimlAuthCookie = cookies.stream().filter(v -> v.startsWith(zassConfigProperties.getTokenPrefix())).map(v -> v.substring(v.indexOf("=") + 1)).findFirst();
             if (apimlAuthCookie.isPresent()) {
                 token = apimlAuthCookie.get();
             }

diff --git a/zaas-client/src/main/java/org/zowe/apiml/zaasclient/util/SimpleHttpResponse.java b/zaas-client/src/main/java/org/zowe/apiml/zaasclient/util/SimpleHttpResponse.java
@@ -19,6 +19,7 @@
 import org.apache.hc.core5.http.io.entity.EntityUtils;
 
 import java.io.IOException;
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -40,25 +41,25 @@ public class SimpleHttpResponse {
     private final int code;
     private final byte[] byteBody;
     private final String stringBody;
-    private final Map<String, Header> headers;
+    private final Map<String, List<Header>> headers;
 
     public SimpleHttpResponse(int code, byte[] byteBody) {
-        this(code, byteBody, null, null);
+        this(code, byteBody, null, Map.of());
     }
 
     public SimpleHttpResponse(int code, String stringBody) {
-        this(code, null, stringBody, null);
+        this(code, null, stringBody, Map.of());
     }
 
-    public SimpleHttpResponse(int code, String stringBody, Map<String, Header> headers) {
+    public SimpleHttpResponse(int code, String stringBody, Map<String, List<Header>> headers) {
         this(code, null, stringBody, headers);
     }
 
     public SimpleHttpResponse(int code) {
-        this(code, null, null, null);
+        this(code, null, null, Map.of());
     }
 
-    public SimpleHttpResponse(int code, Map<String, Header> headers) {
+    public SimpleHttpResponse(int code, Map<String, List<Header>> headers) {
         this(code, null, null, headers);
     }
 

diff --git a/...nt/src/test/java/org/zowe/apiml/zaasclient/service/internal/ZaasClientImplHttpsTests.java b/...nt/src/test/java/org/zowe/apiml/zaasclient/service/internal/ZaasClientImplHttpsTests.java
@@ -188,7 +188,7 @@ private void assertThatExceptionContainValidCode(ZaasClientException zce, ZaasCl
     }
 
     @Test
-    void testLoginWithCredentials_ValidUserName_ValidPassword() throws ZaasClientException, IOException {
+    void testLoginWithCredentials_ValidUserName_ValidPassword() throws ZaasClientException {
         prepareResponse(HttpStatus.SC_NO_CONTENT, true);
 
         String token = tokenService.login(VALID_USER, VALID_PASSWORD);
@@ -197,6 +197,21 @@ void testLoginWithCredentials_ValidUserName_ValidPassword() throws ZaasClientExc
         assertEquals("token", token, "Token Mismatch");
     }
 
+    @Test
+    void testLoginWithCredentials_ValidUserName_ValidPassword_multipleResponseHeaders() throws ZaasClientException {
+        var response = prepareResponse(HttpStatus.SC_NO_CONTENT, false);
+        var tokenCookieHeader = mock(Header.class);
+        Header[] headers = new Header[]{header, tokenCookieHeader};
+        when(response.getHeaders(HttpHeaders.SET_COOKIE)).thenReturn(headers);
+        when(header.getValue()).thenReturn("someCookie=cookieValue");
+        when(tokenCookieHeader.getValue()).thenReturn("apimlAuthenticationToken=token");
+
+        String token = tokenService.login(VALID_USER, VALID_PASSWORD);
+        assertNotNull(token, "null Token obtained");
+        assertNotEquals(EMPTY_STRING, token, "Empty Token obtained");
+        assertEquals("token", token, "Token Mismatch");
+    }
+
     private static Stream<Arguments> provideInvalidUsernamePassword() {
         return Stream.of(
             Arguments.of(HttpStatus.SC_UNAUTHORIZED, INVALID_USER, VALID_PASSWORD, ZaasClientErrorCodes.INVALID_AUTHENTICATION),
@@ -335,7 +350,7 @@ void testPassTicketWithToken_ValidToken_ValidPassTicket() throws Exception {
     }
 
     @Test
-    void givenValidToken_whenLogout_thenSuccess() throws ZaasClientException, IOException {
+    void givenValidToken_whenLogout_thenSuccess() throws ZaasClientException {
         prepareResponse(HttpStatus.SC_NO_CONTENT, true);
         String token = tokenService.login(getAuthHeader(VALID_USER, VALID_PASSWORD));
         assertDoesNotThrow(() -> tokenService.logout(token));