Cache JWT tokens instead of generating per request

Author: wendigo

core/trino-main/src/main/java/io/trino/server/InternalAuthenticationManager.java

@@ -30,8 +30,12 @@
 
 import javax.crypto.SecretKey;
 
+import java.time.Instant;
 import java.time.ZonedDateTime;
 import java.util.Date;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
+import java.util.function.Supplier;
 
 import static io.airlift.http.client.Request.Builder.fromRequest;
 import static io.jsonwebtoken.security.Keys.hmacShaKeyFor;
@@ -41,18 +45,23 @@
 import static jakarta.ws.rs.core.MediaType.TEXT_PLAIN_TYPE;
 import static jakarta.ws.rs.core.Response.Status.UNAUTHORIZED;
 import static java.nio.charset.StandardCharsets.UTF_8;
+import static java.time.temporal.ChronoUnit.MINUTES;
 import static java.util.Objects.requireNonNull;
 
 public class InternalAuthenticationManager
         implements HttpRequestFilter
 {
     private static final Logger log = Logger.get(InternalAuthenticationManager.class);
+    private static final Supplier<Instant> DEFAULT_EXPIRATION_SUPPLIER = () -> ZonedDateTime.now().plusMinutes(6).toInstant();
+    // Leave a 5 minute buffer to allow for clock skew and GC pauses
+    private static final Function<Instant, Instant> TOKEN_REUSE_THRESHOLD = instant -> instant.minus(5, MINUTES);
 
     private static final String TRINO_INTERNAL_BEARER = "X-Trino-Internal-Bearer";
 
     private final SecretKey hmac;
     private final String nodeId;
     private final JwtParser jwtParser;
+    private final AtomicReference<InternalToken> currentToken;
 
     @Inject
     public InternalAuthenticationManager(InternalCommunicationConfig internalCommunicationConfig, SecurityConfig securityConfig, NodeInfo nodeInfo)
@@ -84,6 +93,7 @@ public InternalAuthenticationManager(String sharedSecret, String nodeId)
         this.hmac = hmacShaKeyFor(Hashing.sha256().hashString(sharedSecret, UTF_8).asBytes());
         this.nodeId = nodeId;
         this.jwtParser = newJwtParserBuilder().verifyWith(hmac).build();
+        this.currentToken = new AtomicReference<>(createJwt());
     }
 
     public static boolean isInternalRequest(ContainerRequestContext request)
@@ -118,17 +128,34 @@ public void handleInternalRequest(ContainerRequestContext request)
     public Request filterRequest(Request request)
     {
         return fromRequest(request)
-                .addHeader(TRINO_INTERNAL_BEARER, generateJwt())
+                .addHeader(TRINO_INTERNAL_BEARER, getOrGenerateJwt())
                 .build();
     }
 
-    private String generateJwt()
+    private String getOrGenerateJwt()
     {
-        return newJwtBuilder()
+        InternalToken token = currentToken.get();
+        if (token.isExpired()) {
+            InternalToken newToken = createJwt();
+            if (currentToken.compareAndSet(token, newToken)) {
+                token = newToken;
+            }
+            else {
+                // Another thread already generated a new token
+                token = currentToken.get();
+            }
+        }
+        return token.token();
+    }
+
+    private InternalToken createJwt()
+    {
+        Instant expiration = DEFAULT_EXPIRATION_SUPPLIER.get();
+        return new InternalToken(expiration, newJwtBuilder()
                 .signWith(hmac)
                 .subject(nodeId)
-                .expiration(Date.from(ZonedDateTime.now().plusMinutes(5).toInstant()))
-                .compact();
+                .expiration(Date.from(expiration))
+                .compact());
     }
 
     private String parseJwt(String jwt)
@@ -138,4 +165,18 @@ private String parseJwt(String jwt)
                 .getPayload()
                 .getSubject();
     }
+
+    private record InternalToken(Instant expiration, String token)
+    {
+        public InternalToken
+        {
+            expiration = TOKEN_REUSE_THRESHOLD.apply(requireNonNull(expiration, "expiration is null"));
+            requireNonNull(token, "token is null");
+        }
+
+        public boolean isExpired()
+        {
+            return Instant.now().isAfter(expiration);
+        }
+    }
 }

core/trino-main/src/main/java/io/trino/server/InternalCommunicationModule.java

@@ -14,6 +14,7 @@
 package io.trino.server;
 
 import com.google.inject.Binder;
+import com.google.inject.Scopes;
 import com.google.inject.multibindings.Multibinder;
 import io.airlift.configuration.AbstractConfigurationAwareModule;
 import io.airlift.discovery.client.ForDiscoveryClient;
@@ -52,7 +53,7 @@ protected void setup(Binder binder)
         }
         discoveryFilterBinder.addBinding().to(InternalAuthenticationManager.class);
         configBinder(binder).bindConfigDefaults(HttpClientConfig.class, ForDiscoveryClient.class, config -> configureClient(config, internalCommunicationConfig));
-        binder.bind(InternalAuthenticationManager.class);
+        binder.bind(InternalAuthenticationManager.class).in(Scopes.SINGLETON);
     }
 
     private static class DiscoveryEncodeAddressAsHostname