Proper refresh of credentials

3 files changed, 121 insertions, 8 deletions

diff --git a/jdisc-cloud-aws/pom.xml b/jdisc-cloud-aws/pom.xml
index 3dfbf1e9154..af5c5c7cc8f 100644
--- a/jdisc-cloud-aws/pom.xml
+++ b/jdisc-cloud-aws/pom.xml
@@ -47,6 +47,23 @@
             <groupId>com.amazonaws</groupId>
             <artifactId>aws-java-sdk-ssm</artifactId>
         </dependency>
+        <dependency>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter-api</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter-engine</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>com.yahoo.vespa</groupId>
+            <artifactId>testutil</artifactId>
+            <version>${project.version}</version>
+            <scope>test</scope>
+        </dependency>
+
     </dependencies>
 
     <build>
diff --git a/jdisc-cloud-aws/src/main/java/com/yahoo/jdisc/cloud/aws/VespaAwsCredentialsProvider.java b/jdisc-cloud-aws/src/main/java/com/yahoo/jdisc/cloud/aws/VespaAwsCredentialsProvider.java
index fc9c03a824a..4bca505620b 100644
--- a/jdisc-cloud-aws/src/main/java/com/yahoo/jdisc/cloud/aws/VespaAwsCredentialsProvider.java
+++ b/jdisc-cloud-aws/src/main/java/com/yahoo/jdisc/cloud/aws/VespaAwsCredentialsProvider.java
@@ -14,45 +14,87 @@ import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.time.Clock;
+import java.time.Duration;
+import java.time.Instant;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 
 public class VespaAwsCredentialsProvider implements AWSCredentialsProvider {
 
+    private static final Logger logger = Logger.getLogger(VespaAwsCredentialsProvider.class.getName());
     private static final String DEFAULT_CREDENTIALS_PATH = "/opt/vespa/var/vespa/aws/credentials.json";
-
-    private final AtomicReference<AWSCredentials> credentials = new AtomicReference<>();
+    private static final Duration REFRESH_INTERVAL = Duration.ofMinutes(30);
+    private final AtomicReference<Credentials> credentials = new AtomicReference<>();
     private final Path credentialsPath;
-
+    private final Clock clock;
 
     public VespaAwsCredentialsProvider() {
-        this.credentialsPath = Path.of(DEFAULT_CREDENTIALS_PATH);
+        this(Path.of(DEFAULT_CREDENTIALS_PATH), Clock.systemUTC());
+    }
+
+    VespaAwsCredentialsProvider(Path credentialsPath, Clock clock) {
+        this.credentialsPath = credentialsPath;
+        this.clock = clock;
         refresh();
     }
 
     @Override
     public AWSCredentials getCredentials() {
-        return credentials.get();
+        Credentials sessionCredentials = credentials.get();
+        if (Duration.between(clock.instant(), sessionCredentials.expiry).abs().compareTo(REFRESH_INTERVAL)<0) {
+            refresh();
+            sessionCredentials = credentials.get();
+        }
+        return sessionCredentials;
     }
 
     @Override
     public void refresh() {
         try {
+            logger.log(Level.FINE, "Refreshing credentials from disk");
             credentials.set(readCredentials());
         } catch (Exception e) {
             throw new RuntimeException("Unable to get credentials. Please ensure cluster is configured as exclusive. See: https://cloud.vespa.ai/en/reference/services#nodes");
         }
     }
 
-    private AWSSessionCredentials readCredentials() {
+    private Credentials readCredentials() {
         try {
             Slime slime = SlimeUtils.jsonToSlime(Files.readAllBytes(credentialsPath));
             Cursor cursor = slime.get();
             String accessKey = cursor.field("awsAccessKey").asString();
             String secretKey = cursor.field("awsSecretKey").asString();
-            String sessionToken = cursor.field("sessionToken").asString();
-            return new BasicSessionCredentials(accessKey, secretKey, sessionToken);
+            String sessionToken = cursor.field("sessionTToken").asString();
+            Instant defaultExpiry = Instant.now().plus(Duration.ofHours(1));
+            Instant expiry;
+            try {
+                expiry = SlimeUtils.optionalString(cursor.field("expiry")).map(Instant::parse).orElse(defaultExpiry);
+            } catch (Exception e) {
+                expiry = defaultExpiry;
+                logger.warning("Unable to read expiry from credentials");
+            }
+            return new Credentials(accessKey, secretKey, sessionToken, expiry);
         } catch (IOException e) {
             throw new UncheckedIOException(e);
         }
     }
+
+    record Credentials (String awsAccessKey, String awsSecretKey, String sessionToken, Instant expiry) implements AWSSessionCredentials {
+        @Override
+        public String getSessionToken() {
+            return sessionToken;
+        }
+
+        @Override
+        public String getAWSAccessKeyId() {
+            return awsAccessKey;
+        }
+
+        @Override
+        public String getAWSSecretKey() {
+            return awsSecretKey;
+        }
+    }
 }
diff --git a/jdisc-cloud-aws/src/test/java/com/yahoo/jdisc/cloud/aws/VespaAwsCredentialsProviderTest.java b/jdisc-cloud-aws/src/test/java/com/yahoo/jdisc/cloud/aws/VespaAwsCredentialsProviderTest.java
new file mode 100644
index 00000000000..2725d28651e
--- /dev/null
+++ b/jdisc-cloud-aws/src/test/java/com/yahoo/jdisc/cloud/aws/VespaAwsCredentialsProviderTest.java
@@ -0,0 +1,54 @@
+package com.yahoo.jdisc.cloud.aws;
+
+import com.amazonaws.auth.AWSCredentials;
+import com.yahoo.test.ManualClock;
+import com.yahoo.vespa.test.file.TestFileSystem;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.file.FileSystem;
+import java.nio.file.FileSystems;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.time.Duration;
+import java.time.Instant;
+
+public class VespaAwsCredentialsProviderTest {
+
+    @Test
+    void refreshes_credentials() throws IOException {
+        Path credentialsPath = TestFileSystem.create().getPath("/credentials.json");
+        ManualClock clock = new ManualClock(Instant.now());
+
+        Instant originalExpiry = clock.instant().plus(Duration.ofHours(12));
+        writeCredentials(credentialsPath, originalExpiry);
+        VespaAwsCredentialsProvider credentialsProvider = new VespaAwsCredentialsProvider(credentialsPath, clock);
+
+        AWSCredentials credentials = credentialsProvider.getCredentials();
+        Assertions.assertEquals(originalExpiry.toString(), credentials.getAWSAccessKeyId());
+
+        Instant updatedExpiry = clock.instant().plus(Duration.ofHours(24));
+        writeCredentials(credentialsPath, updatedExpiry);
+        // File updated, but old credentials still valid
+        credentials = credentialsProvider.getCredentials();
+        Assertions.assertEquals(originalExpiry.toString(), credentials.getAWSAccessKeyId());
+
+        // Credentials refreshes when it is < 30 minutes left until expiry
+        clock.advance(Duration.ofHours(11).plus(Duration.ofMinutes(31)));
+        credentials = credentialsProvider.getCredentials();
+        Assertions.assertEquals(updatedExpiry.toString(), credentials.getAWSAccessKeyId());
+
+    }
+
+    private void writeCredentials(Path path, Instant expiry) throws IOException {
+        String content = """
+                {
+                   "awsAccessKey": "%1$s",
+                   "awsSecretKey": "%1$s",
+                   "sessionToken": "%1$s",
+                   "expiry": "%1$s"
+                 }""".formatted(expiry.toString());
+        Files.writeString(path, content);
+    }
+}