fixes #1248 update the TokenHandler to support multiple OAuth 2.0 pro…

egress-router/src/main/java/com/networknt/router/middleware/TokenHandler.java

@@ -16,46 +16,55 @@
 
 package com.networknt.router.middleware;
 
+import com.networknt.client.ClientConfig;
 import com.networknt.client.oauth.Jwt;
 import com.networknt.client.oauth.OauthHelper;
+import com.networknt.client.oauth.TokenKeyRequest;
 import com.networknt.config.Config;
 import com.networknt.handler.Handler;
 import com.networknt.handler.MiddlewareHandler;
 import com.networknt.httpstring.HttpStringConstants;
 import com.networknt.monad.Result;
+import com.networknt.status.Status;
+import com.networknt.utility.ConcurrentHashSet;
 import com.networknt.utility.ModuleRegistry;
 import io.undertow.Handlers;
 import io.undertow.server.HttpHandler;
 import io.undertow.server.HttpServerExchange;
+import io.undertow.util.HeaderValues;
 import io.undertow.util.Headers;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * This is a middleware handler that is responsible for getting a JWT access token from
  * OAuth 2.0 provider for the particular router client. The only token that is supported in
  * this handler is client credentials token as there is no user information available here.
  *
- * The client_id will be retrieved from client.yml and client_secret will be retrieved from
- * secret.yml
+ * The client_id and client_secret will be retrieved from client.yml and client_secret should
+ * be encrypted or set as an environment variable. In Kubernetes cluster, you can create a
+ * sealed secret for it.
  *
- * This handler will also responsible for checking if the cached token is about to expired
+ * This handler will also responsible for checking if the cached token is about expired
  * or not. In which case, it will renew the token in another thread. When request comes and
  * the cached token is already expired, then it will block the request and go to the OAuth
  * provider to get a new token and then resume the request to the next handler in the chain.
  *
  * The logic is very similar with client module in light-4j but this is implemented in a
- * handler instead.
+ * handler instead. Multiple OAuth 2.0 providers are supported and the token cache strategy
+ * can be defined based on your OAuth 2.0 providers.
  *
  * This light-router is designed for standalone or client that is not implemented in Java
  * Otherwise, you should use client module instead of this one. In the future, we might
  * add Authorization Code grant type support by providing an endpoint in the light-router
  * to accept Authorization Code redirect and then get the token from OAuth 2.0 provider.
  *
  * There is no specific configuration file for this handler just to enable or disable it. If
- * you want to bypass this handler, you can comment it out from service.yml middleware
+ * you want to bypass this handler, you can comment it out from handler.yml middleware
  * handler section or change the token.yml to disable it.
  *
  * Once the token is retrieved from OAuth 2.0 provider, it will be placed in the header as
@@ -65,27 +74,76 @@
 public class TokenHandler implements MiddlewareHandler {
     public static final String CONFIG_NAME = "token";
     public static final String ENABLED = "enabled";
+    private static final String HANDLER_DEPENDENCY_ERROR = "ERR10074";
 
     public static Map<String, Object> config = Config.getInstance().getJsonMapConfigNoCache(CONFIG_NAME);
     static Logger logger = LoggerFactory.getLogger(TokenHandler.class);
     protected volatile HttpHandler next;
-    // Cached jwt token for this handler on behalf of a client.
-    private final Jwt cachedJwt = new Jwt();
+    // Cached jwt token for this handler on behalf of a client by serviceId as the key
+    private final Map<String, Jwt> cache = new ConcurrentHashMap();
     public TokenHandler() { }
 
     @Override
     public void handleRequest(final HttpServerExchange exchange) throws Exception {
-        // check if there is a bear token in the authorization header in the request. If this
+        // This handler must be put after the prefix or dict handler so that the serviceId is
+        // readily available in the header resolved by the path or the endpoint from the request.
+        HeaderValues headerValues = exchange.getRequestHeaders().get(HttpStringConstants.SERVICE_ID);
+        String serviceId = null;
+        if(headerValues != null) serviceId = headerValues.getFirst();
+        if(serviceId == null) {
+            // this handler should be before the router and after the handler to resolve the serviceId from path
+            // or endpoint like the PathPrefixServiceHandler or ServiceDictHandler.
+            logger.error("The serviceId cannot be resolved. Do you have PathPrefixServiceHandler or ServiceDictHandler before this handler?");
+            setExchangeStatus(exchange, HANDLER_DEPENDENCY_ERROR, "TokenHandler", "PathPrefixServiceHandler");
+            return;
+        }
+        ClientConfig clientConfig = ClientConfig.get();
+        Map<String, Object> tokenConfig = clientConfig.getTokenConfig();
+        Map<String, Object> ccConfig = (Map<String, Object>)tokenConfig.get(ClientConfig.CLIENT_CREDENTIALS);
+
+        Jwt cachedJwt = cache.get(serviceId);
+        // get a new token if cachedJwt is null or the jwt is about expired.
+        if(cachedJwt == null || cachedJwt.getExpire() - (Long)tokenConfig.get(ClientConfig.TOKEN_RENEW_BEFORE_EXPIRED) < System.currentTimeMillis()) {
+            Jwt.Key key = new Jwt.Key(serviceId);
+            cachedJwt = new Jwt(key); // create a new instance if the cache is empty for the serviceId.
+
+            if(clientConfig.isMultipleAuthServers()) {
+                // get the right client credentials configuration based on the serviceId
+                Map<String, Object> serviceIdAuthServers = (Map<String, Object>)ccConfig.get(ClientConfig.SERVICE_ID_AUTH_SERVERS);
+                if(serviceIdAuthServers == null) {
+                    throw new RuntimeException("serviceIdAuthServers property is missing in the token client credentials configuration");
+                }
+                Map<String, Object> authServerConfig = (Map<String, Object>)serviceIdAuthServers.get(serviceId);
+                // overwrite some elements in the auth server config if it is not defined.
+                if(authServerConfig.get(ClientConfig.PROXY_HOST) == null) authServerConfig.put(ClientConfig.PROXY_HOST, tokenConfig.get(ClientConfig.PROXY_HOST));
+                if(authServerConfig.get(ClientConfig.PROXY_PORT) == null) authServerConfig.put(ClientConfig.PROXY_PORT, tokenConfig.get(ClientConfig.PROXY_PORT));
+                if(authServerConfig.get(ClientConfig.TOKEN_RENEW_BEFORE_EXPIRED) == null) authServerConfig.put(ClientConfig.TOKEN_RENEW_BEFORE_EXPIRED, tokenConfig.get(ClientConfig.TOKEN_RENEW_BEFORE_EXPIRED));
+                if(authServerConfig.get(ClientConfig.EXPIRED_REFRESH_RETRY_DELAY) == null) authServerConfig.put(ClientConfig.EXPIRED_REFRESH_RETRY_DELAY, tokenConfig.get(ClientConfig.EXPIRED_REFRESH_RETRY_DELAY));
+                if(authServerConfig.get(ClientConfig.EARLY_REFRESH_RETRY_DELAY) == null) authServerConfig.put(ClientConfig.EARLY_REFRESH_RETRY_DELAY, tokenConfig.get(ClientConfig.EARLY_REFRESH_RETRY_DELAY));
+                cachedJwt.setCcConfig(authServerConfig);
+            } else {
+                // only one client credentials configuration, populate some common elements to the ccConfig from tokenConfig.
+                ccConfig.put(ClientConfig.PROXY_HOST, tokenConfig.get(ClientConfig.PROXY_HOST));
+                ccConfig.put(ClientConfig.PROXY_PORT, tokenConfig.get(ClientConfig.PROXY_PORT));
+                ccConfig.put(ClientConfig.TOKEN_RENEW_BEFORE_EXPIRED, tokenConfig.get(ClientConfig.TOKEN_RENEW_BEFORE_EXPIRED));
+                ccConfig.put(ClientConfig.EXPIRED_REFRESH_RETRY_DELAY, tokenConfig.get(ClientConfig.EXPIRED_REFRESH_RETRY_DELAY));
+                ccConfig.put(ClientConfig.EARLY_REFRESH_RETRY_DELAY, tokenConfig.get(ClientConfig.EARLY_REFRESH_RETRY_DELAY));
+                cachedJwt.setCcConfig(ccConfig);
+            }
+            Result result = OauthHelper.populateCCToken(cachedJwt);
+            if(result.isFailure()) {
+                logger.error("Cannot populate or renew jwt for client credential grant type: " + result.getError().toString());
+                setExchangeStatus(exchange, result.getError());
+                return;
+            }
+            // put the cachedJwt to the cache.
+            cache.put(serviceId, cachedJwt);
+        }
+        // check if there is a bear token in the authorization header in the request. If there
         // is one, then this must be the subject token that is linked to the original user.
         // We will keep this token in the Authorization header but create a new token with
         // client credentials grant type with scopes for the particular client. (Can we just
         // assume that the subject token has the scope already?)
-        Result result = OauthHelper.populateCCToken(cachedJwt);
-        if(result.isFailure()) {
-            logger.error("cannot populate or renew jwt for client credential grant type");
-            OauthHelper.sendStatusToResponse(exchange, result.getError());
-            return;
-        }
         String token = exchange.getRequestHeaders().getFirst(Headers.AUTHORIZATION);
         if(token == null) {
             exchange.getRequestHeaders().put(Headers.AUTHORIZATION, "Bearer " + cachedJwt.getJwt());

egress-router/src/main/resources/config/token.yml

@@ -1,2 +1,6 @@
-enabled: false
-scopeToken: true
+# This is the configuration file for the TokenHandler that is responsible for getting
+# a client credentials token in http-sidecar and light-gateway when calling others.
+# The configuration for one or multiple OAuth 2.0 providers is in the client.yml file.
+
+# indicate if the handler is enabled.
+enabled: ${token.enabled:false}

egress-router/src/test/java/com/networknt/router/middleware/TokenHandlerTest.java

@@ -0,0 +1,140 @@
+package com.networknt.router.middleware;
+
+import com.networknt.client.Http2Client;
+import com.networknt.exception.ClientException;
+import com.networknt.httpstring.HttpStringConstants;
+import io.undertow.Handlers;
+import io.undertow.Undertow;
+import io.undertow.UndertowOptions;
+import io.undertow.client.ClientConnection;
+import io.undertow.client.ClientRequest;
+import io.undertow.client.ClientResponse;
+import io.undertow.server.HttpHandler;
+import io.undertow.server.RoutingHandler;
+import io.undertow.util.Headers;
+import io.undertow.util.Methods;
+import org.junit.*;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.xnio.IoUtils;
+import org.xnio.OptionMap;
+
+import java.net.URI;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicReference;
+
+/**
+ * This is the test class to ensure that the TokenHandler is working with multiple
+ * OAuth 2.0 providers with proper cache based on the serviceId and scopes.
+ *
+ * @author Steve Hu
+ */
+public class TokenHandlerTest {
+    static final Logger logger = LoggerFactory.getLogger(TokenHandlerTest.class);
+    static Undertow server = null;
+
+    @BeforeClass
+    public static void setUp() throws Exception{
+        if(server == null) {
+            logger.info("starting server");
+            HttpHandler handler = getTestHandler();
+            TokenHandler tokenHandler = new TokenHandler();
+            tokenHandler.setNext(handler);
+            handler = tokenHandler;
+            server = Undertow.builder()
+                    .setServerOption(UndertowOptions.ENABLE_HTTP2, true)
+                    .addHttpListener(7080, "localhost")
+                    .setHandler(handler)
+                    .build();
+            server.start();
+        }
+    }
+
+    @AfterClass
+    public static void tearDown() throws Exception {
+        if(server != null) {
+            try {
+                Thread.sleep(100);
+            } catch (InterruptedException ignored) {
+
+            }
+            server.stop();
+            logger.info("The server is stopped.");
+        }
+    }
+
+    static RoutingHandler getTestHandler() {
+        return Handlers.routing()
+                .add(Methods.POST, "/", exchange -> {
+                    exchange.getResponseSender().send("POST OK");
+                })
+                .add(Methods.GET, "/", exchange -> {
+                    exchange.getResponseSender().send("GET OK");
+                });
+    }
+
+    @Test
+    @Ignore
+    public void testOneGetService1Request() throws Exception {
+        final Http2Client client = Http2Client.getInstance();
+        final CountDownLatch latch = new CountDownLatch(1);
+        final ClientConnection connection;
+        try {
+            connection = client.connect(new URI("http://localhost:7080"), Http2Client.WORKER, Http2Client.SSL, Http2Client.BUFFER_POOL, OptionMap.create(UndertowOptions.ENABLE_HTTP2, true)).get();
+        } catch (Exception e) {
+            throw new ClientException(e);
+        }
+        final AtomicReference<ClientResponse> reference = new AtomicReference<>();
+        try {
+            ClientRequest request = new ClientRequest().setPath("/service1").setMethod(Methods.GET);
+            request.getRequestHeaders().put(HttpStringConstants.SERVICE_ID, "service1");
+            request.getRequestHeaders().put(Headers.HOST, "localhost");
+            connection.sendRequest(request, client.createClientCallback(reference, latch));
+            latch.await();
+        } catch (Exception e) {
+            logger.error("Exception: ", e);
+            throw new ClientException(e);
+        } finally {
+            IoUtils.safeClose(connection);
+        }
+        int statusCode = reference.get().getResponseCode();
+        String body = reference.get().getAttachment(Http2Client.RESPONSE_BODY);
+        Assert.assertEquals(200, statusCode);
+        if(statusCode == 200) {
+            Assert.assertEquals("GET OK", body);
+        }
+    }
+
+    @Test
+    @Ignore
+    public void testOnePostRequest() throws Exception {
+        final Http2Client client = Http2Client.getInstance();
+        final CountDownLatch latch = new CountDownLatch(1);
+        final ClientConnection connection;
+        try {
+            connection = client.connect(new URI("http://localhost:7080"), Http2Client.WORKER, Http2Client.SSL, Http2Client.BUFFER_POOL, OptionMap.create(UndertowOptions.ENABLE_HTTP2, true)).get();
+        } catch (Exception e) {
+            throw new ClientException(e);
+        }
+        final AtomicReference<ClientResponse> reference = new AtomicReference<>();
+        String requestBody = "{\"clientId\":\"FSC_0030303343303x32AA2\",\"loggedInUserEmail\":\"steve.hu@networknt.com\"}";
+        try {
+            ClientRequest request = new ClientRequest().setPath("/services/apexrest/NNT_ConquestApplicationServices/getAccountDetailsForConquestToUpdatePlan").setMethod(Methods.POST);
+            request.getRequestHeaders().put(Headers.CONTENT_TYPE, "application/json");
+            request.getRequestHeaders().put(Headers.TRANSFER_ENCODING, "chunked");
+            request.getRequestHeaders().put(Headers.HOST, "localhost");
+            connection.sendRequest(request, client.createClientCallback(reference, latch, requestBody));
+            latch.await();
+        } catch (Exception e) {
+            logger.error("Exception: ", e);
+            throw new ClientException(e);
+        } finally {
+            IoUtils.safeClose(connection);
+        }
+        int statusCode = reference.get().getResponseCode();
+        String body = reference.get().getAttachment(Http2Client.RESPONSE_BODY);
+        System.out.println("statusCode = " + statusCode + " body = " + body);
+        Assert.assertEquals(200, statusCode);
+    }
+
+}

egress-router/src/test/resources/config/values.yml

@@ -70,3 +70,26 @@ router.headerRewriteRules:
   /v3/address:
     - oldK: path
       newK: route
+
+# client.yml
+# test configuration for the TokenHandler with multiple OAuth 2.0 providers.
+client.multipleAuthServers: true
+client.tokenCcServiceIdAuthServers:
+  service1:
+    server_url: https://localhost:7771
+    enableHttp2: true
+    uri: /oauth1/token
+    client_id: f7d42348-c647-4efb-a52d-4c5787421e72
+    client_secret: f6h1FTI8Q3-7UScPZDzfXA
+    scope:
+      - petstore.r
+      - petstore.w
+  service2:
+    server_url: https://localhost:7772
+    enableHttp2: true
+    uri: /oauth2/token
+    client_id: f7d42348-c647-4efb-a52d-4c5787421e73
+    client_secret: f6h1FTI8Q3-7UScPZDzfXA
+    scope:
+      - market.r
+      - market.w

status/src/main/resources/config/status.yml

@@ -492,6 +492,11 @@ ERR10073:
   code: ERR10073
   message: LIMIT_KEY_NOT_FOUND
   description: Could not resolve the rate limit key with key type %s and resolver %s.
+ERR10074:
+  statusCode: 400
+  code: ERR10074
+  message: HANDLER_DEPENDENCY_ERROR
+  description: Handler %s depends on handler %s in the chain.