From 5066db5409406b52817ed20e95361be0a7e89bf1 Mon Sep 17 00:00:00 2001
From: Andrew Azores <aazores@redhat.com>
Date: Mon, 22 Jul 2024 15:46:58 -0400
Subject: [PATCH] fix(token): retry request if token may have potentially
 expired between reading from fs and handling by remote authenticator

---
 .../io/cryostat/agent/AuthorizationType.java  | 16 ++++++--
 .../java/io/cryostat/agent/ConfigModule.java  |  2 +-
 .../java/io/cryostat/agent/MainModule.java    | 37 ++++++++++++++++++-
 3 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/src/main/java/io/cryostat/agent/AuthorizationType.java b/src/main/java/io/cryostat/agent/AuthorizationType.java
index 5dc2cb18..891458da 100644
--- 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
index 31b70525..428efc3e 100644
--- a/src/main/java/io/cryostat/agent/ConfigModule.java
+++ b/src/main/java/io/cryostat/agent/ConfigModule.java
@@ -245,7 +245,7 @@ public static int provideCryostatAgentWebclientResponseTimeoutMs(Config config)
 
     @Provides
     @Singleton
-    @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
index 7793953c..2be47050 100644
--- a/src/main/java/io/cryostat/agent/MainModule.java
+++ b/src/main/java/io/cryostat/agent/MainModule.java
@@ -65,11 +65,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;
 
@@ -242,6 +245,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,
@@ -259,7 +263,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);