feat: add ability to configure tes error cache expire time (#1751)

rauthsuj · aws-kevinrickard · commit d44a286cf20f · 2025-12-16T17:16:47.000-08:00
added the ability to configure TES error cache expire time, TES wouldn't restart when error cache is changed and if the time is less than 10 secs or greater than 12 hrs it will fallback to default time.
diff --git a/src/main/java/com/aws/greengrass/tes/CredentialRequestHandler.java b/src/main/java/com/aws/greengrass/tes/CredentialRequestHandler.java
@@ -68,10 +68,14 @@ public class CredentialRequestHandler implements HttpHandler {
     public static final String AUTH_HEADER = "Authorization";
     public static final String IOT_CREDENTIALS_HTTP_VERB = "GET";
     public static final String SUPPORTED_REQUEST_VERB = "GET";
-    public static final int TIME_BEFORE_CACHE_EXPIRE_IN_MIN = 5;
-    public static final int CLOUD_4XX_ERROR_CACHE_IN_MIN = 2;
-    public static final int CLOUD_5XX_ERROR_CACHE_IN_MIN = 1;
-    public static final int UNKNOWN_ERROR_CACHE_IN_MIN = 5;
+    public static final int DEFAULT_TIME_BEFORE_CACHE_EXPIRE_IN_SEC = 300;
+    public static final int DEFAULT_CLOUD_4XX_ERROR_CACHE_IN_SEC = 120;
+    public static final int DEFAULT_CLOUD_5XX_ERROR_CACHE_IN_SEC = 60;
+    public static final int DEFAULT_UNKNOWN_ERROR_CACHE_IN_SEC = 300;
+
+    private volatile int cloud4xxErrorCacheInSec = DEFAULT_CLOUD_4XX_ERROR_CACHE_IN_SEC;
+    private volatile int cloud5xxErrorCacheInSec = DEFAULT_CLOUD_5XX_ERROR_CACHE_IN_SEC;
+    private volatile int unknownErrorCacheInSec = DEFAULT_UNKNOWN_ERROR_CACHE_IN_SEC;
 
     private String iotCredentialsPath;
 
@@ -142,6 +146,20 @@ void setIotCredentialsPath(String iotRoleAlias) {
         this.iotCredentialsPath = "/role-aliases/" + iotRoleAlias + "/credentials";
     }
 
+    /**
+     * Configure error cache settings for error responses.
+     *
+     * @param cloud4xxErrorCache error cache duration in seconds for 4xx errors.
+     * @param cloud5xxErrorCache error cache duration in seconds for 5xx errors.
+     * @param unknownErrorCache error cache duration in seconds for unknown errors.
+     */
+    public void configureCacheSettings(int cloud4xxErrorCache, int cloud5xxErrorCache, int unknownErrorCache) {
+        this.cloud4xxErrorCacheInSec = cloud4xxErrorCache;
+        this.cloud5xxErrorCacheInSec = cloud5xxErrorCache;
+        this.unknownErrorCacheInSec = unknownErrorCache;
+    }
+
+
     @Override
     @SuppressWarnings("PMD.AvoidCatchingThrowable")
     public void handle(final HttpExchange exchange) throws IOException {
@@ -281,14 +299,14 @@ private byte[] getCredentialsBypassCache() {
                         LOGGER.atError().kv(IOT_CRED_PATH_KEY, iotCredentialsPath)
                                 .log("Unable to cache expired credentials which expired at {}", expiry);
                     } else {
-                        newExpiry = expiry.minus(Duration.ofMinutes(TIME_BEFORE_CACHE_EXPIRE_IN_MIN));
+                        newExpiry = expiry.minus(Duration.ofSeconds(DEFAULT_TIME_BEFORE_CACHE_EXPIRE_IN_SEC));
                         tesCache.get(iotCredentialsPath).responseCode = HttpURLConnection.HTTP_OK;
 
                         if (newExpiry.isBefore(Instant.now(clock))) {
                             LOGGER.atWarn().kv(IOT_CRED_PATH_KEY, iotCredentialsPath)
                                     .log("Can't cache credentials as new credentials {} will "
-                                                    + "expire in less than {} minutes", expiry,
-                                            TIME_BEFORE_CACHE_EXPIRE_IN_MIN);
+                                                    + "expire in less than {} seconds", expiry,
+                                            DEFAULT_TIME_BEFORE_CACHE_EXPIRE_IN_SEC);
                         } else {
                             LOGGER.atInfo().kv(IOT_CRED_PATH_KEY, iotCredentialsPath)
                                     .log("Received IAM credentials that will be cached until {}", newExpiry);
@@ -318,7 +336,7 @@ private byte[] getCredentialsBypassCache() {
             String responseString = "Failed to get connection";
             response = responseString.getBytes(StandardCharsets.UTF_8);
             // Use unknown error cache policy for SSL/TLS connection errors to prevent excessive retries
-            newExpiry = Instant.now(clock).plus(Duration.ofMinutes(UNKNOWN_ERROR_CACHE_IN_MIN));
+            newExpiry = Instant.now(clock).plus(Duration.ofSeconds(unknownErrorCacheInSec));
             tesCache.get(iotCredentialsPath).responseCode = HttpURLConnection.HTTP_INTERNAL_ERROR;
             tesCache.get(iotCredentialsPath).expiry = newExpiry;
             tesCache.get(iotCredentialsPath).credentials = response;
@@ -421,16 +439,16 @@ private String parseExpiryFromResponse(final String credentials) throws AWSIotEx
     }
 
     private Instant getExpiryPolicyForErr(int statusCode) {
-        int expiryTime = UNKNOWN_ERROR_CACHE_IN_MIN; // In case of unrecognized cloud errors, back off
+        int expiryTime = unknownErrorCacheInSec; // In case of unrecognized cloud errors, back off
         // Add caching Time-To-Live (TTL) for TES cloud errors
         if (statusCode >= 400 && statusCode < 500) {
             // 4xx retries are only meaningful unless a user action has been adopted, TTL should be longer
-            expiryTime = CLOUD_4XX_ERROR_CACHE_IN_MIN;
+            expiryTime = cloud4xxErrorCacheInSec;
         } else if (statusCode >= 500 && statusCode < 600) {
             // 5xx could be a temporary cloud unavailability, TTL should be shorter
-            expiryTime = CLOUD_5XX_ERROR_CACHE_IN_MIN;
+            expiryTime = cloud5xxErrorCacheInSec;
         }
-        return Instant.now(clock).plus(Duration.ofMinutes(expiryTime));
+        return Instant.now(clock).plus(Duration.ofSeconds(expiryTime));
     }
 
     /**
diff --git a/src/main/java/com/aws/greengrass/tes/TokenExchangeService.java b/src/main/java/com/aws/greengrass/tes/TokenExchangeService.java
@@ -9,6 +9,7 @@
 import com.aws.greengrass.authorization.exceptions.AuthorizationException;
 import com.aws.greengrass.config.Topic;
 import com.aws.greengrass.config.Topics;
+import com.aws.greengrass.config.WhatHappened;
 import com.aws.greengrass.dependency.ImplementsService;
 import com.aws.greengrass.dependency.State;
 import com.aws.greengrass.deployment.DeviceConfiguration;
@@ -41,6 +42,15 @@ public class TokenExchangeService extends GreengrassService implements AwsCreden
     private String iotRoleAlias;
     private HttpServerImpl server;
 
+    public static final String CLOUD_4XX_ERROR_CACHE_TOPIC = "error4xxCredentialRetryInSec";
+    public static final String CLOUD_5XX_ERROR_CACHE_TOPIC = "error5xxCredentialRetryInSec";
+    public static final String UNKNOWN_ERROR_CACHE_TOPIC = "errorUnknownCredentialRetryInSec";
+    private static final int MINIMUM_ERROR_CACHE_IN_SEC = 10;
+    private static final int MAXIMUM_ERROR_CACHE_IN_SEC = 42_900;
+    private int cloud4xxErrorCache;
+    private int cloud5xxErrorCache;
+    private int unknownErrorCache;
+
     private final AuthorizationHandler authZHandler;
     private final CredentialRequestHandler credentialRequestHandler;
 
@@ -57,24 +67,77 @@ public TokenExchangeService(Topics topics,
                                 AuthorizationHandler authZHandler, DeviceConfiguration deviceConfiguration) {
         super(topics);
         port = Coerce.toInt(config.lookup(CONFIGURATION_CONFIG_KEY, PORT_TOPIC).dflt(DEFAULT_PORT));
+        deviceConfiguration.getIotRoleAlias().subscribe((why, newv) -> {
+            iotRoleAlias = Coerce.toString(newv);
+        });
+
+        this.authZHandler = authZHandler;
+        this.credentialRequestHandler = credentialRequestHandler;
+
+        cloud4xxErrorCache = validateErrorCacheConfig(Coerce.toInt(config.findOrDefault(
+                CredentialRequestHandler.DEFAULT_CLOUD_4XX_ERROR_CACHE_IN_SEC, CONFIGURATION_CONFIG_KEY,
+                CLOUD_4XX_ERROR_CACHE_TOPIC)), CLOUD_4XX_ERROR_CACHE_TOPIC,
+                CredentialRequestHandler.DEFAULT_CLOUD_4XX_ERROR_CACHE_IN_SEC);
+        cloud5xxErrorCache = validateErrorCacheConfig(Coerce.toInt(config.findOrDefault(
+                CredentialRequestHandler.DEFAULT_CLOUD_5XX_ERROR_CACHE_IN_SEC, CONFIGURATION_CONFIG_KEY,
+                CLOUD_5XX_ERROR_CACHE_TOPIC)), CLOUD_5XX_ERROR_CACHE_TOPIC,
+                CredentialRequestHandler.DEFAULT_CLOUD_5XX_ERROR_CACHE_IN_SEC);
+        unknownErrorCache = validateErrorCacheConfig(Coerce.toInt(config.findOrDefault(
+                CredentialRequestHandler.DEFAULT_UNKNOWN_ERROR_CACHE_IN_SEC, CONFIGURATION_CONFIG_KEY,
+                UNKNOWN_ERROR_CACHE_TOPIC)), UNKNOWN_ERROR_CACHE_TOPIC,
+                CredentialRequestHandler.DEFAULT_UNKNOWN_ERROR_CACHE_IN_SEC);
+
+        credentialRequestHandler.configureCacheSettings(cloud4xxErrorCache, cloud5xxErrorCache, unknownErrorCache);
+
         config.subscribe((why, node) -> {
+            logger.atDebug("tes-config-change").kv("node", node).kv("what", why).log();
+            if (why.equals(WhatHappened.timestampUpdated)) {
+                return;
+            }
+            if (node != null && (node.childOf(CLOUD_4XX_ERROR_CACHE_TOPIC)
+                    || node.childOf(CLOUD_5XX_ERROR_CACHE_TOPIC)
+                    || node.childOf(UNKNOWN_ERROR_CACHE_TOPIC))) {
+
+                int newCloud4xxErrorCache = validateErrorCacheConfig(Coerce.toInt(config.findOrDefault(
+                        CredentialRequestHandler.DEFAULT_CLOUD_4XX_ERROR_CACHE_IN_SEC, CONFIGURATION_CONFIG_KEY,
+                        CLOUD_4XX_ERROR_CACHE_TOPIC)), CLOUD_4XX_ERROR_CACHE_TOPIC,
+                        CredentialRequestHandler.DEFAULT_CLOUD_4XX_ERROR_CACHE_IN_SEC);
+                int newCloud5xxErrorCache = validateErrorCacheConfig(Coerce.toInt(config.findOrDefault(
+                        CredentialRequestHandler.DEFAULT_CLOUD_5XX_ERROR_CACHE_IN_SEC, CONFIGURATION_CONFIG_KEY,
+                        CLOUD_5XX_ERROR_CACHE_TOPIC)), CLOUD_5XX_ERROR_CACHE_TOPIC,
+                        CredentialRequestHandler.DEFAULT_CLOUD_5XX_ERROR_CACHE_IN_SEC);
+                int newUnknownErrorCache = validateErrorCacheConfig(Coerce.toInt(config.findOrDefault(
+                        CredentialRequestHandler.DEFAULT_UNKNOWN_ERROR_CACHE_IN_SEC, CONFIGURATION_CONFIG_KEY,
+                        UNKNOWN_ERROR_CACHE_TOPIC)), UNKNOWN_ERROR_CACHE_TOPIC,
+                        CredentialRequestHandler.DEFAULT_UNKNOWN_ERROR_CACHE_IN_SEC);
+
+                if (cloud4xxErrorCache != newCloud4xxErrorCache
+                        || cloud5xxErrorCache != newCloud5xxErrorCache
+                        || unknownErrorCache != newUnknownErrorCache) {
+
+                    cloud4xxErrorCache = newCloud4xxErrorCache;
+                    cloud5xxErrorCache = newCloud5xxErrorCache;
+                    unknownErrorCache = newUnknownErrorCache;
+
+                    credentialRequestHandler.configureCacheSettings(
+                            newCloud4xxErrorCache, newCloud5xxErrorCache, newUnknownErrorCache);
+
+                    logger.atInfo("tes-error-cache-config-change")
+                            .kv("node", node).kv("why", why)
+                            .log("TES error cache configuration updated");
+                }
+            }
             if (node != null && node.childOf(PORT_TOPIC)) {
-                logger.atDebug("tes-config-change").kv("node", node).kv("why", why).log();
                 port = Coerce.toInt(node);
                 Topic activePortTopic = config.lookup(CONFIGURATION_CONFIG_KEY, ACTIVE_PORT_TOPIC);
+
                 if (port != Coerce.toInt(activePortTopic)) {
-                    logger.atInfo("tes-config-change").kv(PORT_TOPIC, port).kv("node", node).kv("why", why)
+                    logger.atInfo("tes-port-config-change").kv(PORT_TOPIC, port).kv("node", node).kv("why", why)
                             .log("Restarting TES server due to port config change");
                     requestRestart();
                 }
             }
         });
-        deviceConfiguration.getIotRoleAlias().subscribe((why, newv) -> {
-            iotRoleAlias = Coerce.toString(newv);
-        });
-
-        this.authZHandler = authZHandler;
-        this.credentialRequestHandler = credentialRequestHandler;
     }
 
     @Override
@@ -130,6 +193,16 @@ private void validateConfig() {
         }
     }
 
+    private int validateErrorCacheConfig(int newCacheValue, String topic, int defaultCacheValue) {
+        if (newCacheValue < MINIMUM_ERROR_CACHE_IN_SEC || newCacheValue > MAXIMUM_ERROR_CACHE_IN_SEC) {
+            logger.atError()
+                    .log("Error cache value must be between {} and {} seconds, setting {} to default value {}",
+                    MINIMUM_ERROR_CACHE_IN_SEC, MAXIMUM_ERROR_CACHE_IN_SEC, topic, defaultCacheValue);
+            return defaultCacheValue;
+        }
+        return newCacheValue;
+    }
+
     @Override
     public AwsCredentials resolveCredentials() {
         return credentialRequestHandler.getAwsCredentials();
diff --git a/src/test/java/com/aws/greengrass/tes/CredentialRequestHandlerTest.java b/src/test/java/com/aws/greengrass/tes/CredentialRequestHandlerTest.java
@@ -47,10 +47,10 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 
-import static com.aws.greengrass.tes.CredentialRequestHandler.CLOUD_4XX_ERROR_CACHE_IN_MIN;
-import static com.aws.greengrass.tes.CredentialRequestHandler.CLOUD_5XX_ERROR_CACHE_IN_MIN;
-import static com.aws.greengrass.tes.CredentialRequestHandler.TIME_BEFORE_CACHE_EXPIRE_IN_MIN;
-import static com.aws.greengrass.tes.CredentialRequestHandler.UNKNOWN_ERROR_CACHE_IN_MIN;
+import static com.aws.greengrass.tes.CredentialRequestHandler.DEFAULT_CLOUD_4XX_ERROR_CACHE_IN_SEC;
+import static com.aws.greengrass.tes.CredentialRequestHandler.DEFAULT_CLOUD_5XX_ERROR_CACHE_IN_SEC;
+import static com.aws.greengrass.tes.CredentialRequestHandler.DEFAULT_UNKNOWN_ERROR_CACHE_IN_SEC;
+import static com.aws.greengrass.tes.CredentialRequestHandler.DEFAULT_TIME_BEFORE_CACHE_EXPIRE_IN_SEC;
 import static com.aws.greengrass.testcommons.testutilities.ExceptionLogProtector.ignoreExceptionOfType;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.containsString;
@@ -320,15 +320,15 @@ void GIVEN_credential_handler_WHEN_called_handle_THEN_caches_creds() throws Exce
         verify(mockStream, times(1)).write(expectedResponse);
 
         // Expiry time in recent future won't give error but there wil be no caching
-        expirationTime = Instant.now().plus(Duration.ofMinutes(TIME_BEFORE_CACHE_EXPIRE_IN_MIN - 1));
+        expirationTime = Instant.now().plus(Duration.ofSeconds(DEFAULT_TIME_BEFORE_CACHE_EXPIRE_IN_SEC - 60));
         responseStr = String.format(RESPONSE_STR, expirationTime.toString());
         mockResponse = new IotCloudResponse(responseStr.getBytes(StandardCharsets.UTF_8), 200);
         when(mockCloudHelper.sendHttpRequest(any(), any(), any(), any(), any())).thenReturn(mockResponse);
         handler.handle(mockExchange);
         verify(mockCloudHelper, times(2)).sendHttpRequest(any(), any(), any(), any(), any());
 
         // Expiry time in future will result in credentials being cached
-        expirationTime = Instant.now().plus(Duration.ofMinutes(TIME_BEFORE_CACHE_EXPIRE_IN_MIN + 1));
+        expirationTime = Instant.now().plus(Duration.ofSeconds(DEFAULT_TIME_BEFORE_CACHE_EXPIRE_IN_SEC + 60));
         responseStr = String.format(RESPONSE_STR, expirationTime.toString());
         mockResponse = new IotCloudResponse(responseStr.getBytes(StandardCharsets.UTF_8), 200);
         when(mockCloudHelper.sendHttpRequest(any(), any(), any(), any(), any())).thenReturn(mockResponse);
@@ -401,7 +401,7 @@ void GIVEN_4xx_response_code_WHEN_called_handle_THEN_expire_in_2_minutes() throw
                 String.format("TES responded with status code: %d. Caching response. ", expectedStatus).getBytes();
         // expire in 2 minutes
         handler.getAwsCredentials();
-        Instant expirationTime = Instant.now().plus(Duration.ofMinutes(CLOUD_4XX_ERROR_CACHE_IN_MIN));
+        Instant expirationTime = Instant.now().plus(Duration.ofSeconds(DEFAULT_CLOUD_4XX_ERROR_CACHE_IN_SEC));
         Clock mockClock = Clock.fixed(expirationTime, ZoneId.of("UTC"));
         handler.setClock(mockClock);
         handler.getAwsCredentials();
@@ -425,7 +425,7 @@ void GIVEN_5xx_response_code_WHEN_called_handle_THEN_expire_in_1_minute() throws
                 String.format("TES responded with status code: %d. Caching response. ", expectedStatus).getBytes();
         // expire in 1 minute
         handler.getAwsCredentials();
-        Instant expirationTime = Instant.now().plus(Duration.ofMinutes(CLOUD_5XX_ERROR_CACHE_IN_MIN));
+        Instant expirationTime = Instant.now().plus(Duration.ofSeconds(DEFAULT_CLOUD_5XX_ERROR_CACHE_IN_SEC));
         Clock mockClock = Clock.fixed(expirationTime, ZoneId.of("UTC"));
         handler.setClock(mockClock);
         handler.getAwsCredentials();
@@ -449,7 +449,7 @@ void GIVEN_unknown_error_response_code_WHEN_called_handle_THEN_expire_in_5_minut
                 String.format("TES responded with status code: %d. Caching response. ", expectedStatus).getBytes();
         // expire in 5 minutes
         handler.getAwsCredentials();
-        Instant expirationTime = Instant.now().plus(Duration.ofMinutes(UNKNOWN_ERROR_CACHE_IN_MIN));
+        Instant expirationTime = Instant.now().plus(Duration.ofSeconds(DEFAULT_UNKNOWN_ERROR_CACHE_IN_SEC));
         Clock mockClock = Clock.fixed(expirationTime, ZoneId.of("UTC"));
         handler.setClock(mockClock);
         handler.getAwsCredentials();
diff --git a/src/test/java/com/aws/greengrass/tes/TokenExchangeServiceTest.java b/src/test/java/com/aws/greengrass/tes/TokenExchangeServiceTest.java
