fix: allow TES restart on port changes (#1691)

junfuchen99 · alter-mage · alter-mage · commit 95b8efa953f5 · 2025-04-02T12:12:05.000-07:00
* fix: allow TES restart on port changes

* chore: fix unit tests

---------

Co-authored-by: Siddhant Srivastava &lt;sidsriv@amazon.com&gt;
diff --git a/src/integrationtests/java/com/aws/greengrass/integrationtests/e2e/tes/TESTest.java b/src/integrationtests/java/com/aws/greengrass/integrationtests/e2e/tes/TESTest.java
@@ -20,7 +20,9 @@
 import com.aws.greengrass.logging.api.Logger;
 import com.aws.greengrass.logging.impl.LogManager;
 import com.aws.greengrass.tes.CredentialRequestHandler;
+import com.aws.greengrass.tes.HttpServerImpl;
 import com.aws.greengrass.tes.TokenExchangeService;
+import com.aws.greengrass.util.Coerce;
 import com.aws.greengrass.util.IamSdkClientFactory;
 import com.aws.greengrass.util.IotSdkClientFactory;
 import org.junit.jupiter.api.AfterAll;
@@ -43,16 +45,19 @@
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Path;
 import java.util.Collections;
+import java.util.Objects;
 import java.util.UUID;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 
 import static com.aws.greengrass.componentmanager.KernelConfigResolver.CONFIGURATION_CONFIG_KEY;
+import static com.aws.greengrass.deployment.DeviceConfiguration.IOT_ROLE_ALIAS_TOPIC;
 import static com.aws.greengrass.easysetup.DeviceProvisioningHelper.ThingInfo;
 import static com.aws.greengrass.integrationtests.e2e.BaseE2ETestCase.E2ETEST_ENV_STAGE;
 import static com.aws.greengrass.lifecyclemanager.GreengrassService.SERVICES_NAMESPACE_TOPIC;
 import static com.aws.greengrass.lifecyclemanager.GreengrassService.SETENV_CONFIG_NAMESPACE;
-import static com.aws.greengrass.deployment.DeviceConfiguration.IOT_ROLE_ALIAS_TOPIC;
+import static com.aws.greengrass.tes.TokenExchangeService.PORT_TOPIC;
 import static com.aws.greengrass.tes.TokenExchangeService.TES_URI_ENV_VARIABLE_NAME;
 import static com.aws.greengrass.tes.TokenExchangeService.TOKEN_EXCHANGE_SERVICE_TOPICS;
 import static com.aws.greengrass.testcommons.testutilities.ExceptionLogProtector.ignoreExceptionUltimateCauseOfType;
@@ -131,6 +136,43 @@ static void tearDown() throws URISyntaxException {
         }
     }
 
+    @Test
+    void GIVEN_iot_role_alias_WHEN_port_changes_THEN_valid_credentials_are_returned_from_new_port() throws Exception {
+        // Get current port and calculate new port
+        int currentPort = Coerce.toInt(
+                kernel.getConfig().lookupTopics(SERVICES_NAMESPACE_TOPIC, TOKEN_EXCHANGE_SERVICE_TOPICS)
+                        .lookup(CONFIGURATION_CONFIG_KEY, PORT_TOPIC));
+        int newPort = currentPort + 1;
+        String expectedUri = String.format("http://localhost:%d%s", newPort, HttpServerImpl.URL);
+
+        // Setup server restart detection
+        CountDownLatch serverRestarted = new CountDownLatch(1);
+        AtomicReference<String> urlString = new AtomicReference<>();
+        kernel.getConfig().find(SETENV_CONFIG_NAMESPACE, TES_URI_ENV_VARIABLE_NAME).subscribe((why, newv) -> {
+            urlString.set(Coerce.toString(newv));
+            if (urlString.get().equals(expectedUri)) {
+                serverRestarted.countDown();
+            }
+        });
+
+        // Change port and wait for server restart
+        kernel.getConfig().lookupTopics(SERVICES_NAMESPACE_TOPIC, TOKEN_EXCHANGE_SERVICE_TOPICS)
+                .lookup(CONFIGURATION_CONFIG_KEY, PORT_TOPIC).withValue(newPort);
+        assertTrue(serverRestarted.await(5, TimeUnit.SECONDS), "Server did not restart within 5 seconds");
+        assertEquals(expectedUri, urlString.get(), "New URL does not match expected URI");
+
+        // Get authentication token and make request
+        String token = Objects.requireNonNull(kernel.getConfig()
+                        .findTopics(SERVICES_NAMESPACE_TOPIC, AuthenticationHandler.AUTHENTICATION_TOKEN_LOOKUP_KEY)).iterator()
+                .next().getName();
+
+        String response = getResponseString(new URL(urlString.get()), token);
+
+        // Verify response format
+        assertThat(response, matchesPattern(
+                "\\{\"AccessKeyId\":\".+\",\"SecretAccessKey\":\".+\",\"Expiration\":\".+\",\"Token\":\".+\"\\}"));
+    }
+
     @Test
     void GIVEN_iot_role_alias_WHEN_tes_is_queried_THEN_valid_credentials_are_returned(ExtensionContext context)
             throws Exception {
diff --git a/src/main/java/com/aws/greengrass/tes/TokenExchangeService.java b/src/main/java/com/aws/greengrass/tes/TokenExchangeService.java
@@ -19,7 +19,7 @@
 import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
 
 import java.io.IOException;
-import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashSet;
 import javax.inject.Inject;
 
@@ -56,10 +56,19 @@ public TokenExchangeService(Topics topics,
                                 CredentialRequestHandler credentialRequestHandler,
                                 AuthorizationHandler authZHandler, DeviceConfiguration deviceConfiguration) {
         super(topics);
-        // Port change should not be allowed
-        topics.lookup(CONFIGURATION_CONFIG_KEY, PORT_TOPIC).dflt(DEFAULT_PORT)
-                .subscribe((why, newv) -> port = Coerce.toInt(newv));
-
+        config.lookup(CONFIGURATION_CONFIG_KEY, PORT_TOPIC).dflt(DEFAULT_PORT);
+        config.subscribe((why, node) -> {
+            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)
+                            .log("Restarting TES server due to port config change");
+                    requestRestart();
+                }
+            }
+        });
         deviceConfiguration.getIotRoleAlias().subscribe((why, newv) -> {
             iotRoleAlias = Coerce.toString(newv);
         });
@@ -72,7 +81,8 @@ public TokenExchangeService(Topics topics,
     public void postInject() {
         super.postInject();
         try {
-            authZHandler.registerComponent(this.getName(), new HashSet<>(Arrays.asList(AUTHZ_TES_OPERATION)));
+            authZHandler.registerComponent(this.getName(),
+                    new HashSet<>(Collections.singletonList(AUTHZ_TES_OPERATION)));
         } catch (AuthorizationException e) {
             serviceErrored(e);
         }
diff --git a/src/test/java/com/aws/greengrass/tes/TokenExchangeServiceTest.java b/src/test/java/com/aws/greengrass/tes/TokenExchangeServiceTest.java
@@ -8,6 +8,7 @@
 import com.amazon.aws.iot.greengrass.component.common.ComponentType;
 import com.aws.greengrass.authorization.AuthorizationHandler;
 import com.aws.greengrass.authorization.exceptions.AuthorizationException;
+import com.aws.greengrass.config.ChildChanged;
 import com.aws.greengrass.config.Configuration;
 import com.aws.greengrass.config.Subscriber;
 import com.aws.greengrass.config.Topic;
@@ -131,14 +132,7 @@ void cleanup() throws Exception {
     @ParameterizedTest
     @ValueSource(ints = {0, 3000})
     void GIVEN_token_exchange_service_WHEN_started_THEN_correct_env_set(int port) throws Exception {
-        Topic portTopic = mock(Topic.class);
-        when(portTopic.dflt(anyInt())).thenReturn(portTopic);
-        when(portTopic.subscribe(any())).thenAnswer((a) -> {
-            ((Subscriber) a.getArgument(0)).published(WhatHappened.initialized, portTopic);
-            return null;
-        });
-        when(portTopic.getOnce()).thenReturn(port);
-
+        Topic portTopic = Topic.of(new Context(), PORT_TOPIC, port);
         Topic roleTopic = mock(Topic.class);
         when(roleTopic.subscribe(any())).thenAnswer((a) -> {
             ((Subscriber) a.getArgument(0)).published(WhatHappened.initialized, roleTopic);
@@ -157,6 +151,11 @@ void GIVEN_token_exchange_service_WHEN_started_THEN_correct_env_set(int port) th
         when(configuration.lookup(SERVICES_NAMESPACE_TOPIC, DEFAULT_NUCLEUS_COMPONENT_NAME, CONFIGURATION_CONFIG_KEY,
                 IOT_ROLE_ALIAS_TOPIC)).thenReturn(roleTopic);
 
+        when(config.subscribe(any())).thenAnswer((a) -> {
+            ((ChildChanged) a.getArgument(0)).childChanged(WhatHappened.initialized, portTopic);
+            return null;
+        });
+
         TokenExchangeService tes = new TokenExchangeService(config,
                 mockCredentialHandler,
                 mockAuthZHandler, deviceConfigurationWithRoleAlias(MOCK_ROLE_ALIAS));
@@ -180,7 +179,6 @@ void GIVEN_token_exchange_service_WHEN_started_THEN_correct_env_set(int port) th
         verify(mockAuthZHandler).registerComponent(stringArgumentCaptor.capture(), operationsCaptor.capture());
         assertEquals(TOKEN_EXCHANGE_SERVICE_TOPICS, stringArgumentCaptor.getValue());
         assertTrue(operationsCaptor.getValue().contains(TokenExchangeService.AUTHZ_TES_OPERATION));
-
     }
 
     @ParameterizedTest
@@ -202,13 +200,7 @@ void GIVEN_token_exchange_service_WHEN_started_with_empty_role_alias_THEN_server
         // set mock for port topic
         Topic portTopic = mock(Topic.class);
         when(portTopic.dflt(anyInt())).thenReturn(portTopic);
-        when(portTopic.subscribe(any())).thenAnswer((a) -> {
-            ((Subscriber) a.getArgument(0)).published(WhatHappened.initialized, portTopic);
-            return null;
-        });
-        when(portTopic.getOnce()).thenReturn(8080);
-
-
+        
         when(config.lookup(CONFIGURATION_CONFIG_KEY, PORT_TOPIC)).thenReturn(portTopic);
         when(configuration.lookup(SERVICES_NAMESPACE_TOPIC, DEFAULT_NUCLEUS_COMPONENT_NAME, CONFIGURATION_CONFIG_KEY,
                 IOT_ROLE_ALIAS_TOPIC)).thenReturn(roleTopic);
@@ -240,11 +232,6 @@ void GIVEN_token_exchange_service_WHEN_auth_errors_THEN_server_errors_out(Extens
         // set mock for port topic
         Topic portTopic = mock(Topic.class);
         when(portTopic.dflt(anyInt())).thenReturn(portTopic);
-        when(portTopic.subscribe(any())).thenAnswer((a) -> {
-            ((Subscriber) a.getArgument(0)).published(WhatHappened.initialized, portTopic);
-            return null;
-        });
-        when(portTopic.getOnce()).thenReturn(8080);
         when(config.lookup(CONFIGURATION_CONFIG_KEY, PORT_TOPIC)).thenReturn(portTopic);
         when(configuration.lookup(SERVICES_NAMESPACE_TOPIC, DEFAULT_NUCLEUS_COMPONENT_NAME, CONFIGURATION_CONFIG_KEY,
                 IOT_ROLE_ALIAS_TOPIC)).thenReturn(roleTopic);
