support authorization header in deterministic routing

Signed-off-by: Pablo Carle <pablo.carle@broadcom.com>
zowe · Feb 20, 2025 · f4ff995 · f4ff995
1 parent 540658e
commit f4ff995
Show file tree

Hide file tree

Showing 2 changed files with 103 additions and 7 deletions.
diff --git a/...-service/src/main/java/org/zowe/apiml/gateway/loadbalancer/DeterministicLoadBalancer.java b/...-service/src/main/java/org/zowe/apiml/gateway/loadbalancer/DeterministicLoadBalancer.java
@@ -25,6 +25,7 @@
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpStatus;
 import org.springframework.web.server.ResponseStatusException;
+import org.zowe.apiml.constants.ApimlConstants;
 import org.zowe.apiml.gateway.caching.LoadBalancerCache;
 import org.zowe.apiml.gateway.caching.LoadBalancerCache.LoadBalancerCacheRecord;
 import reactor.core.publisher.Flux;
@@ -114,12 +115,32 @@ private boolean isTooOld(LocalDateTime cachedDate) {
 
     private Mono<String> getSub(Object requestContext) {
         if (requestContext instanceof RequestDataContext ctx) {
-            var token = Optional.ofNullable(ctx.getClientRequest().getCookies().get("apimlAuthenticationToken")).map(list -> list.get(0)).orElse("");
+            var token = Optional.ofNullable(getTokenFromCookie(ctx))
+                                .orElseGet(() -> getTokenFromHeader(ctx));
             return Mono.just(extractSubFromToken(token));
         }
         return Mono.just("");
     }
 
+    private String getTokenFromCookie(RequestDataContext ctx) {
+        var tokens = ctx.getClientRequest().getCookies().get("apimlAuthenticationToken");
+        return tokens == null || tokens.isEmpty() ? null : tokens.get(0);
+    }
+
+    private String getTokenFromHeader(RequestDataContext ctx) {
+        var authHeader = ctx.getClientRequest().getHeaders().get(HttpHeaders.AUTHORIZATION);
+        var token = authHeader == null || authHeader.isEmpty() ? null : authHeader.get(0);
+        if (token != null && token.startsWith(ApimlConstants.BEARER_AUTHENTICATION_PREFIX)) {
+            token = token.replaceFirst(ApimlConstants.BEARER_AUTHENTICATION_PREFIX, "").trim();
+            if (token.isEmpty()) {
+                return null;
+            }
+
+            return token;
+        }
+        return null;
+    }
+
     /**
      * Filters the list of service instances to include only those with the specified instance ID.
      * Optional operation, it verifies if the conditions are met to actually filter the list of instances.
@@ -273,7 +294,7 @@ private Claims getJwtClaims(String jwt) {
     }
 
     private String extractSubFromToken(String token) {
-        if (!token.isEmpty()) {
+        if (token != null && !token.isEmpty()) {
             Claims claims = getJwtClaims(token);
             if (claims != null) {
                 return claims.getSubject();

diff --git a/...vice/src/test/java/org/zowe/apiml/gateway/loadbalancer/DeterministicLoadBalancerTest.java b/...vice/src/test/java/org/zowe/apiml/gateway/loadbalancer/DeterministicLoadBalancerTest.java
@@ -121,7 +121,7 @@ void whenNoToken_thenUseDefaultList() {
         }
 
         @Nested
-        class GivenTokenExists {
+        class GivenTokenExistsInCookie {
 
             @Test
             void whenTokenIsInvalid_thenUseDefaultList() {
@@ -278,10 +278,85 @@ void whenNoPreferece_thenCreateOne() {
 
         }
 
-        /**
-         *
-         *
-         */
+        @Nested
+        class GivenTokenExistsInHeader {
+
+            @Test
+            void whenTokenIsInvalid_thenUseDefaultList() {
+                var requestData = mock(RequestData.class);
+                var context = new RequestDataContext(requestData);
+
+                var headers = new HttpHeaders();
+                headers.add(HttpHeaders.AUTHORIZATION, "Basic aaaaaaaaaa");
+                when(requestData.getHeaders()).thenReturn(headers);
+
+                when(requestData.getCookies()).thenReturn(new LinkedMultiValueMap<>());
+                when(request.getContext()).thenReturn(context);
+
+                StepVerifier.create(loadBalancer.get(request))
+                    .assertNext(chosenInstances -> {
+                        assertNotNull(chosenInstances);
+                        assertEquals(2, chosenInstances.size());
+                    })
+                    .expectComplete()
+                    .verify();
+            }
+
+            @Nested
+            class GivenTokenIsValid {
+
+                @BeforeEach
+                void setUp() {
+                    var requestData = mock(RequestData.class);
+                    var context = new RequestDataContext(requestData);
+
+                    MultiValueMap<String, String> cookie = new LinkedMultiValueMap<>();
+
+                    var headers = new HttpHeaders();
+                    headers.add(HttpHeaders.AUTHORIZATION, "Bearer " + VALID_TOKEN);
+                    when(requestData.getHeaders()).thenReturn(headers);
+                    when(requestData.getCookies()).thenReturn(cookie);
+                    when(request.getContext()).thenReturn(context);
+                    when(clock.now()).thenReturn(Date.from(Instant.ofEpochSecond(1721552753)));
+                }
+
+                @Nested
+                class GivenServiceUsesSticky {
+
+                    @BeforeEach
+                    void setUp() {
+                        Map<String, String> metadata = new HashMap<>();
+                        metadata.put("apiml.lb.type", "authentication");
+                        when(instance1.getMetadata()).thenReturn(metadata);
+                    }
+
+                    @Nested
+                    class GivenCacheHasPreference {
+
+                        @Test
+                        void whenInstanceExists_thenUpdateList() {
+                            when(lbCache.retrieve("USER", "service")).thenReturn(Mono.just(new LoadBalancerCacheRecord("instance1")));
+                            when(lbCache.store(eq("USER"), eq("service"), argThat(cacheRecord -> cacheRecord.getInstanceId().equals("instance1"))))
+                                .thenReturn(Mono.empty());
+
+                            StepVerifier.create(loadBalancer.get(request))
+                                .assertNext(chosenInstances -> {
+                                    assertNotNull(chosenInstances);
+                                    assertEquals(1, chosenInstances.size());
+                                    assertEquals("instance1", chosenInstances.get(0).getInstanceId());
+                                })
+                                .expectComplete()
+                                .verify();
+                        }
+
+                    }
+
+                }
+
+            }
+
+        }
+
         @Nested
         class GivenStickyBalancerIgnored {