fix(token): retry request if token may have potentially expired betwe…

…en reading from fs and handling by remote authenticator
cryostatio · Aug 20, 2024 · 07bcc70 · 07bcc70
1 parent 50e64b3
commit 07bcc70
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 5 deletions.
diff --git a/src/main/java/io/cryostat/agent/AuthorizationType.java b/src/main/java/io/cryostat/agent/AuthorizationType.java
@@ -25,15 +25,17 @@
 import io.cryostat.agent.util.StringUtils;
 
 public enum AuthorizationType implements Function<String, String> {
-    NONE(v -> null),
-    BEARER(v -> String.format("Bearer %s", v)),
+    NONE(false, v -> null),
+    BEARER(false, v -> String.format("Bearer %s", v)),
     BASIC(
+            false,
             v ->
                     String.format(
                             "Basic %s",
                             Base64.getEncoder()
                                     .encodeToString(v.getBytes(StandardCharsets.UTF_8)))),
     KUBERNETES(
+            true,
             v -> {
                 try {
                     File file = new File(v);
@@ -45,6 +47,7 @@ public enum AuthorizationType implements Function<String, String> {
                 }
             }),
     AUTO(
+            true,
             v -> {
                 try {
                     String k8s = KUBERNETES.fn.apply(v);
@@ -58,12 +61,19 @@ public enum AuthorizationType implements Function<String, String> {
             }),
     ;
 
+    private final boolean dynamic;
     private final Function<String, String> fn;
 
-    private AuthorizationType(Function<String, String> fn) {
+    private AuthorizationType(boolean dynamic, Function<String, String> fn) {
+        this.dynamic = dynamic;
         this.fn = fn;
     }
 
+    public boolean isDynamic() {
+        // if the authorization value may change between invocations
+        return this.dynamic;
+    }
+
     @Override
     public String apply(String in) {
         return fn.apply(in);

diff --git a/src/main/java/io/cryostat/agent/ConfigModule.java b/src/main/java/io/cryostat/agent/ConfigModule.java
@@ -313,7 +313,7 @@ public static List<TruststoreConfig> provideCryostatAgentWecblientTlsTruststoreC
         return truststoreConfigs;
     }
 
-    @Named(CRYOSTAT_AGENT_WEBCLIENT_RESPONSE_TIMEOUT_MS)
+    @Named(CRYOSTAT_AGENT_WEBCLIENT_RESPONSE_RETRY_COUNT)
     public static int provideCryostatAgentWebclientResponseRetryCount(Config config) {
         return config.getValue(CRYOSTAT_AGENT_WEBCLIENT_RESPONSE_RETRY_COUNT, int.class);
     }

diff --git a/src/main/java/io/cryostat/agent/MainModule.java b/src/main/java/io/cryostat/agent/MainModule.java
@@ -66,11 +66,14 @@
 import dagger.Provides;
 import org.apache.commons.io.IOUtils;
 import org.apache.commons.lang3.exception.ExceptionUtils;
+import org.apache.http.HttpResponse;
 import org.apache.http.client.HttpClient;
 import org.apache.http.client.config.RequestConfig;
+import org.apache.http.client.protocol.HttpClientContext;
 import org.apache.http.impl.client.HttpClientBuilder;
 import org.apache.http.impl.client.HttpClients;
 import org.apache.http.impl.client.StandardHttpRequestRetryHandler;
+import org.apache.http.protocol.HttpContext;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -330,6 +333,7 @@ public static Optional<SSLContext> provideServerSslContext(
     @Provides
     @Singleton
     public static HttpClient provideHttpClient(
+            AuthorizationType authorizationType,
             @Named(HTTP_CLIENT_SSL_CTX) SSLContext sslContext,
             @Named(ConfigModule.CRYOSTAT_AGENT_WEBCLIENT_TLS_VERIFY_HOSTNAME)
                     boolean verifyHostname,
@@ -347,7 +351,38 @@ public static HttpClient provideHttpClient(
                                         .setSocketTimeout(responseTimeout)
                                         .setRedirectsEnabled(true)
                                         .build())
-                        .setRetryHandler(new StandardHttpRequestRetryHandler(retryCount, true));
+                        .setRetryHandler(
+                                new StandardHttpRequestRetryHandler(retryCount, true) {
+                                    @Override
+                                    public boolean retryRequest(
+                                            IOException exception,
+                                            int executionCount,
+                                            HttpContext context) {
+                                        // if the Authorization header we should send may change
+                                        // over time, ex. we read a Bearer token from a file, then
+                                        // it is possible that we get a 401 or 403 response because
+                                        // the token expired in between the time that we read it
+                                        // from our filesystem and when it was received by the
+                                        // authenticator. So, in this set of conditions, we should
+                                        // refresh our header value and try again right away
+                                        if (authorizationType.isDynamic()) {
+                                            HttpClientContext clientCtx =
+                                                    HttpClientContext.adapt(context);
+                                            if (clientCtx.isRequestSent()) {
+                                                HttpResponse resp = clientCtx.getResponse();
+                                                if (resp != null && resp.getStatusLine() != null) {
+                                                    int sc = resp.getStatusLine().getStatusCode();
+                                                    if (executionCount < 2
+                                                            && (sc == 401 || sc == 403)) {
+                                                        return true;
+                                                    }
+                                                }
+                                            }
+                                        }
+                                        return super.retryRequest(
+                                                exception, executionCount, context);
+                                    }
+                                });
 
         if (!verifyHostname) {
             builder = builder.setSSLHostnameVerifier((hostname, session) -> true);