Refresh token

pom.xml

@@ -66,6 +66,11 @@
       <artifactId>bcprov-jdk15on</artifactId>
       <version>1.70</version> 
     </dependency>
+    <dependency>
+      <groupId>org.springframework</groupId>
+      <artifactId>spring-webmvc</artifactId>
+      <version>5.3.23</version>
+    </dependency>
   </dependencies>
 
   <build>

src/main/java/city/makeour/moc/FetchCognitoToken.java

@@ -50,6 +50,25 @@ public void setAuthParameters(String username, String password) {
         this.password = password;
     }
 
+    public Token refleshToken(String refreshToken) {
+        InitiateAuthRequest initiateAuthRequest = InitiateAuthRequest.builder()
+                .authFlow(AuthFlowType.REFRESH_TOKEN_AUTH)
+                .clientId(this.cognitoClientId)
+                .authParameters(Map.of(
+                        "REFRESH_TOKEN", refreshToken))
+                .build();
+        InitiateAuthResponse initiateAuthResponse = cognitoClient.initiateAuth(initiateAuthRequest);
+        AuthenticationResultType result = initiateAuthResponse.authenticationResult();
+
+        Token token = new Token(
+                result.idToken(),
+                result.refreshToken(),
+                result.accessToken(),
+                result.expiresIn());
+
+        return token;
+    }
+
     public Token fetchTokenWithSrpAuth() throws InvalidKeyException, NoSuchAlgorithmException {
         InitiateAuthRequest initiateAuthRequest = InitiateAuthRequest.builder()
                 .authFlow(AuthFlowType.USER_SRP_AUTH)
@@ -87,10 +106,8 @@ public Token fetchTokenWithSrpAuth() throws InvalidKeyException, NoSuchAlgorithm
         RespondToAuthChallengeResponse authChallengeResponse = cognitoClient.respondToAuthChallenge(respondRequest);
         AuthenticationResultType authResult = authChallengeResponse.authenticationResult();
 
-        return new Token(authResult.idToken(), authResult.refreshToken());
-    }
+        Integer expiresIn = authResult.expiresIn();
 
-    public Token fetchToken() throws InvalidKeyException, NoSuchAlgorithmException {
-        return this.fetchTokenWithSrpAuth();
+        return new Token(authResult.idToken(), authResult.refreshToken(), authResult.accessToken(), expiresIn);
     }
 }

src/main/java/city/makeour/moc/MocClient.java

@@ -35,15 +35,48 @@ public void setMocAuthInfo(String cognitoUserPoolId, String cognitoClientId) {
         this.tokenFetcher = new FetchCognitoToken(cognitoUserPoolId, cognitoClientId);
     }
 
+    public void auth(String username, String password) throws InvalidKeyException, NoSuchAlgorithmException {
+        if (this.tokenFetcher == null) {
+            throw new IllegalStateException("MocClient is not initialized with Cognito auth info.");
+        }
+
+        this.tokenFetcher.setAuthParameters(username, password);
+
+        if (this.refreshTokenStorage.hasRefreshToken()) {
+            RefreshTokenInterface refreshToken = this.refreshTokenStorage.getRefreshToken();
+            if (!refreshToken.isExpired()) {
+                this.refreshToken();
+            }
+            return;
+        }
+
+        this.login(username, password);
+    }
+
     public void login(String username, String password) throws InvalidKeyException, NoSuchAlgorithmException {
         if (this.tokenFetcher == null) {
             throw new IllegalStateException("MocClient is not initialized with Cognito auth info.");
         }
 
         this.tokenFetcher.setAuthParameters(username, password);
-        Token token = this.tokenFetcher.fetchToken();
+        Token token = this.tokenFetcher.fetchTokenWithSrpAuth();
+
+        this.setToken(token.getIdToken());
+        this.refreshTokenStorage.setRefreshToken(token);
+    }
+
+    public void refreshToken() throws InvalidKeyException, NoSuchAlgorithmException {
+        if (this.tokenFetcher == null) {
+            throw new IllegalStateException("MocClient is not initialized with Cognito auth info.");
+        }
+
+        if (!this.refreshTokenStorage.hasRefreshToken()) {
+            throw new IllegalStateException("No refresh token available.");
+        }
+
+        RefreshTokenInterface refreshToken = this.refreshTokenStorage.getRefreshToken();
+        Token token = this.tokenFetcher.refleshToken(refreshToken.getRefreshToken());
         this.setToken(token.getIdToken());
-        this.refreshTokenStorage.setRefreshToken(token.getRefreshToken());
     }
 
     public void setToken(String token) {

src/main/java/city/makeour/moc/RefreshTokenInterface.java

@@ -0,0 +1,16 @@
+package city.makeour.moc;
+
+import java.time.ZonedDateTime;
+
+public interface RefreshTokenInterface {
+
+    public String getRefreshToken();
+
+    public ZonedDateTime getExpirationDate();
+
+    public boolean isExpired();
+
+    public boolean isExpired(ZonedDateTime now);
+
+    public boolean hasRefreshToken();
+}

src/main/java/city/makeour/moc/RefreshTokenStorage.java

@@ -2,17 +2,17 @@
 
 public class RefreshTokenStorage implements RefreshTokenStorageInterface {
 
-    private String refreshToken;
+    private RefreshTokenInterface refreshToken;
 
     public RefreshTokenStorage() {
         this.refreshToken = null;
     }
 
-    public String getRefreshToken() {
-        return refreshToken;
+    public RefreshTokenInterface getRefreshToken() {
+        return this.refreshToken;
     }
 
-    public void setRefreshToken(String refreshToken) {
+    public void setRefreshToken(RefreshTokenInterface refreshToken) {
         this.refreshToken = refreshToken;
     }
 

src/main/java/city/makeour/moc/RefreshTokenStorageInterface.java

@@ -2,9 +2,9 @@
 
 public interface RefreshTokenStorageInterface {
 
-    public String getRefreshToken();
+    public RefreshTokenInterface getRefreshToken();
 
-    public void setRefreshToken(String refreshToken);
+    public void setRefreshToken(RefreshTokenInterface refreshToken);
 
     public boolean hasRefreshToken();
 }

src/main/java/city/makeour/moc/Token.java

@@ -1,12 +1,32 @@
 package city.makeour.moc;
 
-public class Token {
+import java.time.ZonedDateTime;
+
+public class Token implements RefreshTokenInterface {
     private String idToken;
     private String refreshToken;
+    private String accessToken;
+    private ZonedDateTime expirationDate;
+
+    public Token(String idToken, String refreshToken, String accessToken, Integer expiresIn) {
+        this(idToken, refreshToken, accessToken, expiresIn, ZonedDateTime.now());
+    }
 
-    public Token(String idToken, String refreshToken) {
+    public Token(String idToken, String refreshToken, String accessToken, Integer expiresIn, ZonedDateTime now) {
+        if (expiresIn <= 0) {
+            throw new IllegalArgumentException("Expiration date must be greater than 0");
+        }
         this.idToken = idToken;
         this.refreshToken = refreshToken;
+        this.accessToken = accessToken;
+
+        if (expiresIn != null) {
+            if (now == null) {
+                now = ZonedDateTime.now();
+            }
+            ZonedDateTime date = now.plusSeconds(expiresIn);
+            this.expirationDate = date;
+        }
     }
 
     public String getIdToken() {
@@ -16,4 +36,24 @@ public String getIdToken() {
     public String getRefreshToken() {
         return refreshToken;
     }
+
+    public String getAccessToken() {
+        return accessToken;
+    }
+
+    public ZonedDateTime getExpirationDate() {
+        return expirationDate;
+    }
+
+    public boolean isExpired() {
+        return this.isExpired(ZonedDateTime.now());
+    }
+
+    public boolean isExpired(ZonedDateTime now) {
+        return now.isAfter(expirationDate);
+    }
+
+    public boolean hasRefreshToken() {
+        return refreshToken != null && !refreshToken.isEmpty();
+    }
 }

src/main/java/city/makeour/moc/TokenFetcherInterface.java

@@ -6,5 +6,7 @@
 public interface TokenFetcherInterface {
     void setAuthParameters(String username, String password);
 
-    Token fetchToken() throws InvalidKeyException, NoSuchAlgorithmException;
+    Token fetchTokenWithSrpAuth() throws InvalidKeyException, NoSuchAlgorithmException;
+
+    Token refleshToken(String refreshToken) throws InvalidKeyException, NoSuchAlgorithmException;
 }

src/main/java/city/makeour/moc/examples/Auth.java

@@ -0,0 +1,38 @@
+package city.makeour.moc.examples;
+
+import java.util.UUID;
+
+import org.springframework.web.client.RestClientResponseException;
+
+import city.makeour.moc.MocClient;
+import city.makeour.ngsi.v2.model.CreateEntityRequest;
+
+public class Auth {
+
+    public static void main(String[] args) {
+        // Create a new MocClient instance
+        MocClient mocClient = new MocClient("https://orion.sandbox.makeour.city");
+
+        // Set the Cognito auth info
+        mocClient.setMocAuthInfo("your_cognito_user_pool_id", "your_cognito_client_id");
+
+        // Authenticate with username and password
+        try {
+            mocClient.auth("your_username", "your_password");
+        } catch (Exception e) {
+            e.printStackTrace();
+        }
+
+        String uuid = UUID.randomUUID().toString();
+
+        CreateEntityRequest entity = new CreateEntityRequest();
+        entity.setType("TestEntity");
+        entity.setId("urn:ngsi-ld:TestEntity:" + uuid);
+        try {
+            mocClient.entities().createEntity("application/json", entity, "keyValues");
+        } catch (RestClientResponseException e) {
+            System.err.println("Error creating entity: " + e.getMessage());
+        }
+
+    }
+}

src/test/java/city/makeour/moc/FetchCognitoTokenTest.java

@@ -1,6 +1,7 @@
 package city.makeour.moc;
 
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
 
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
@@ -28,5 +29,13 @@ public void testFetchTokenWithSrpAuth() throws Exception {
         Token token = fetchCognitoToken.fetchTokenWithSrpAuth();
         assertNotNull(token.getIdToken());
         assertNotNull(token.getRefreshToken());
+
+        // Refresh tokenを使ってトークンを更新する トークンのローテーションはしない
+        Token refreshedToken = fetchCognitoToken.refleshToken(token.getRefreshToken());
+        assertNotNull(refreshedToken.getIdToken());
+        assertNull(refreshedToken.getRefreshToken());
+
+        Token refreshedToken2 = fetchCognitoToken.refleshToken(token.getRefreshToken());
+        assertNotNull(refreshedToken2.getIdToken());
     }
 }

src/test/java/city/makeour/moc/MocClientTest.java

@@ -84,6 +84,33 @@ void testSetMocAuthInfo() throws GeneralSecurityException, NoSuchAlgorithmExcept
         entity.setId("urn:ngsi-ld:TestEntity:" + uuid);
 
         client.entities().createEntity("application/json", entity, "keyValues");
+    }
 
+    @Test
+    @DisplayName("トークンリフレッシュなど行えるかのテスト")
+    @EnabledIfEnvironmentVariables({
+            @EnabledIfEnvironmentVariable(named = "TEST_COGNITO_USER_POOL_ID", matches = ".*"),
+            @EnabledIfEnvironmentVariable(named = "TEST_COGNITO_CLIENT_ID", matches = ".*"),
+            @EnabledIfEnvironmentVariable(named = "TEST_COGNITO_USERNAME", matches = ".*"),
+            @EnabledIfEnvironmentVariable(named = "TEST_COGNITO_PASSWORD", matches = ".*")
+    })
+    void testAuth() throws GeneralSecurityException, NoSuchAlgorithmException {
+        String cognitoUserPoolId = System.getenv("TEST_COGNITO_USER_POOL_ID");
+        String cognitoClientId = System.getenv("TEST_COGNITO_CLIENT_ID");
+        String username = System.getenv("TEST_COGNITO_USERNAME");
+        String password = System.getenv("TEST_COGNITO_PASSWORD");
+
+        MocClient client = new MocClient();
+        client.setMocAuthInfo(cognitoUserPoolId, cognitoClientId);
+        client.auth(username, password);
+        client.auth(username, password); // トークンリフレッシュを行う
+
+        String uuid = UUID.randomUUID().toString();
+
+        CreateEntityRequest entity = new CreateEntityRequest();
+        entity.setType("TestEntity");
+        entity.setId("urn:ngsi-ld:TestEntity:" + uuid);
+
+        client.entities().createEntity("application/json", entity, "keyValues");
     }
 }