fix: Fixed kid cache management. (#113)

pagopa · Jun 15, 2024 · 96afa8b · 96afa8b
1 parent a9d9459
commit 96afa8b
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 16 deletions.
diff --git a/src/main/java/it/pagopa/swclient/mil/auth/bean/KeyIdCache.java b/src/main/java/it/pagopa/swclient/mil/auth/bean/KeyIdCache.java
@@ -69,4 +69,13 @@ public boolean isValid(long remainingLife) {
 
 		return true;
 	}
+
+	/**
+	 * 
+	 */
+	public void clean() {
+		kid = null;
+		exp = 0;
+		storedAt = 0;
+	}
 }
diff --git a/src/main/java/it/pagopa/swclient/mil/auth/service/KeyManCapabilities.java b/src/main/java/it/pagopa/swclient/mil/auth/service/KeyManCapabilities.java
@@ -18,6 +18,7 @@
 import it.pagopa.swclient.mil.auth.util.UniGenerator;
 import it.pagopa.swclient.mil.azureservices.keyvault.keys.bean.JsonWebKeyType;
 import it.pagopa.swclient.mil.azureservices.keyvault.keys.bean.KeyAttributes;
+import it.pagopa.swclient.mil.azureservices.keyvault.keys.bean.KeyBundle;
 import it.pagopa.swclient.mil.azureservices.keyvault.keys.bean.KeyCreateParameters;
 import it.pagopa.swclient.mil.azureservices.keyvault.keys.service.AzureKeyVaultKeysExtReactiveService;
 import it.pagopa.swclient.mil.azureservices.keyvault.keys.service.AzureKeyVaultKeysReactiveService;
@@ -53,13 +54,14 @@ abstract class KeyManCapabilities {
 	 * 
 	 */
 	private KeyIdCache keyIdCache;
-	
+
 	/**
 	 * 
 	 */
 	KeyManCapabilities() {
+		keyIdCache = new KeyIdCache();
 	}
-	
+
 	/**
 	 * 
 	 * @param keysExtService
@@ -70,14 +72,25 @@ abstract class KeyManCapabilities {
 		this.keysService = keysService;
 		keyIdCache = new KeyIdCache();
 	}
-
+
+	/**
+	 * 
+	 * @param keyBundle
+	 */
+	private void cacheKid(KeyBundle keyBundle) {
+		Log.debug("Cache the key ID");
+		keyIdCache.setKid(keyBundle.getKey().getKid())
+			.setExp(keyBundle.getAttributes().getExp())
+			.setStoredAt(Instant.now().getEpochSecond());
+	}
+
 	/**
 	 * Creates a new key.
 	 * 
 	 * @param keyOps
-	 * @return key id (kid)
+	 * @return
 	 */
-	protected Uni<String> createKey(List<String> keyOps) {
+	protected Uni<KeyBundle> createKey(List<String> keyOps) {
 		Log.trace("Create e new key");
 		long now = Instant.now().getEpochSecond();
 		return keysService.createKey(
@@ -92,17 +105,9 @@ protected Uni<String> createKey(List<String> keyOps) {
 				.setTags(Map.of(it.pagopa.swclient.mil.azureservices.keyvault.keys.util.KeyUtils.DOMAIN_KEY, KeyUtils.DOMAIN_VALUE))
 				.setKeyOps(keyOps)
 				.setKeySize(keysize)
-				.setKty(JsonWebKeyType.RSA))
-			.map(keyBundle -> {
-				String kid = keyBundle.getKey().getKid();
-				Log.debug("Cache the key ID");
-				keyIdCache.setKid(kid)
-					.setExp(keyBundle.getAttributes().getExp())
-					.setStoredAt(Instant.now().getEpochSecond());
-				return kid;
-			});
+				.setKty(JsonWebKeyType.RSA));
 	}
-	
+
 	/**
 	 * Gets a key and if doesn't find it, creates a new one.
 	 * 
@@ -127,8 +132,19 @@ protected Uni<String> retrieveKey(List<String> keyOps) {
 					return createKey(keyOps);
 				} else {
 					Log.trace("Suitable key found");
-					return UniGenerator.item(keyBundle.get().getKey().getKid());
+					return UniGenerator.item(keyBundle.get());
 				}
+			})
+			.map(keyBundle -> {
+				cacheKid(keyBundle);
+				return keyBundle.getKey().getKid();
 			});
 	}
+
+	/**
+	 * 
+	 */
+	public void cleanCache() {
+		keyIdCache.clean();
+	}
 }
diff --git a/src/test/java/it/pagopa/swclient/mil/auth/service/ClaimEncryptorTest.java b/src/test/java/it/pagopa/swclient/mil/auth/service/ClaimEncryptorTest.java
@@ -90,6 +90,7 @@ void init(TestInfo testInfo) {
 		System.out.printf("* %s: START *%n", testInfo.getDisplayName());
 		System.out.println(frame);
 		keyBaseUrl = vaultBaseUrl + (vaultBaseUrl.endsWith("/") ? "keys/" : "/keys/");
+		claimEncryptor.cleanCache();
 	}
 
 	/**
@@ -106,6 +107,8 @@ void given_claimToEncrypt_when_allIsOk_then_getEncryptedClaim() {
 			List.of(JsonWebKeyType.RSA)))
 			.thenReturn(UniGenerator.item(
 				Optional.of(new KeyBundle()
+					.setAttributes(new KeyAttributes()
+						.setExp(Instant.now().plus(15, ChronoUnit.MINUTES).getEpochSecond()))
 					.setKey(new JsonWebKey()
 						.setKid(keyBaseUrl + "key_name/key_version")))));
 

diff --git a/src/test/java/it/pagopa/swclient/mil/auth/service/TokenSignerTest.java b/src/test/java/it/pagopa/swclient/mil/auth/service/TokenSignerTest.java
@@ -101,6 +101,7 @@ void init(TestInfo testInfo) {
 		System.out.printf("* %s: START *%n", testInfo.getDisplayName());
 		System.out.println(frame);
 		keyBaseUrl = vaultBaseUrl + (vaultBaseUrl.endsWith("/") ? "keys/" : "/keys/");
+		tokenSigner.cleanCache();
 	}
 
 	/**
@@ -117,6 +118,8 @@ void given_claimsSetToSign_when_suitableKeyExists_then_getSignedJwt() {
 			List.of(JsonWebKeyType.RSA)))
 			.thenReturn(UniGenerator.item(
 				Optional.of(new KeyBundle()
+					.setAttributes(new KeyAttributes()
+						.setExp(Instant.now().plus(15, ChronoUnit.MINUTES).getEpochSecond()))
 					.setKey(new JsonWebKey()
 						.setKid(keyBaseUrl + "key_name/key_version")))));
 
@@ -292,6 +295,8 @@ void given_claimsSetToSign_when_jwtCreateInstanceGoesWrong_then_getFailure() {
 			List.of(JsonWebKeyType.RSA)))
 			.thenReturn(UniGenerator.item(
 				Optional.of(new KeyBundle()
+					.setAttributes(new KeyAttributes()
+						.setExp(Instant.now().plus(15, ChronoUnit.MINUTES).getEpochSecond()))
 					.setKey(new JsonWebKey()
 						.setKid(keyBaseUrl + "key_name/key_version")))));
 
@@ -351,6 +356,8 @@ void given_claimsSetToSign_when_messageDigestGetInstanceGoesWrong_then_getFailur
 			List.of(JsonWebKeyType.RSA)))
 			.thenReturn(UniGenerator.item(
 				Optional.of(new KeyBundle()
+					.setAttributes(new KeyAttributes()
+						.setExp(Instant.now().plus(15, ChronoUnit.MINUTES).getEpochSecond()))
 					.setKey(new JsonWebKey()
 						.setKid(keyBaseUrl + "key_name/key_version")))));