Skip to content

Commit

Permalink
P4ADEV-778 resolved changed request
Browse files Browse the repository at this point in the history
  • Loading branch information
macacia committed Nov 11, 2024
1 parent 9602e3c commit 0ac3bd4
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 33 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -18,17 +18,16 @@
public class ValidateJWTLegacyService {
public static final String TOKEN_TYPE_A2A = "a2a";

private final A2AClientLegacyPropConfig a2AClientLegacyPropConfig;
private final JWTValidator jwtValidator;
private final Map<String, PublicKey> clientApplicationsPublicKeyMap;

public ValidateJWTLegacyService(A2AClientLegacyPropConfig a2AClientLegacyPropConfig, JWTValidator jwtValidator) {
this.a2AClientLegacyPropConfig = a2AClientLegacyPropConfig;
this.clientApplicationsPublicKeyMap = a2AClientLegacyPropConfig.getPublicKeysAsMap();
this.jwtValidator = jwtValidator;
}

public Pair<String, Map<String, Claim>> validate(String token) {
Map<String, PublicKey> clientApplicationsPublicKeyMap = a2AClientLegacyPropConfig.getPublicKeysAsMap();
Pair<String, Map<String, Claim>> claims = validateToken(clientApplicationsPublicKeyMap, token);
Pair<String, Map<String, Claim>> claims = validateToken(token);
validateM2MType(claims.getRight());
validateClaims(claims.getRight());

Expand All @@ -53,14 +52,16 @@ private void validateClaims(Map<String, Claim> claims) {
}
}

private Pair<String, Map<String, Claim>> validateToken(Map<String, PublicKey> clientApplicationsPublicKeyMap, String token) {
try {
return clientApplicationsPublicKeyMap.keySet().stream()
.map(key -> Pair.of(key, jwtValidator.validate(token, clientApplicationsPublicKeyMap.get(key))))
.findFirst()
.orElseThrow(() -> new InvalidTokenException("Invalid token for A2A call"));
} catch (Exception e) {
return null;
private Pair<String, Map<String, Claim>> validateToken(String token) {
for (String key : clientApplicationsPublicKeyMap.keySet()) {
PublicKey publicKey = clientApplicationsPublicKeyMap.get(key);
try {
Map<String, Claim> claims = jwtValidator.validate(token, publicKey);
return Pair.of(key, claims);
} catch (Exception e) {
log.debug("continue cycling - validation failed with key {}", key);
}
}
throw new InvalidTokenException("Invalid token for A2A call");
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -31,24 +31,27 @@ class ValidateJWTLegacyServiceTest {
@Mock
private JWTValidator jwtValidatorMock;
private ValidateJWTLegacyService service;
private KeyPair keyPair;

private KeyPair keyPair1;
private KeyPair keyPair2;

@BeforeEach
void setup() throws Exception {
keyPair1 = JWTValidatorUtils.generateKeyPair();
keyPair2 = JWTValidatorUtils.generateKeyPair();
Map<String, PublicKey> publicKeyMap = Map.of(
"A2A-IPA_TEST_1", keyPair1.getPublic(),
"A2A-IPA_TEST_2", keyPair2.getPublic());
when(a2AClientLegacyPropConfig.getPublicKeysAsMap()).thenReturn(publicKeyMap);
service = new ValidateJWTLegacyService(a2AClientLegacyPropConfig, jwtValidatorMock);
keyPair = JWTValidatorUtils.generateKeyPair();
}

@Test
void GivenValidTokenThenOk() {
String appName = "A2A-IPA_TEST_1";
PublicKey publicKey = keyPair.getPublic();
String token = JWTValidatorUtils.generateLegacyToken(keyPair, "a2a", Instant.now(), Instant.now().plusSeconds(3_600_000L), "jti");
when(a2AClientLegacyPropConfig.getPublicKeysAsMap()).thenReturn(Map.of(appName, publicKey));
String appName = "A2A-IPA_TEST_2";
String token = JWTValidatorUtils.generateLegacyToken(keyPair2, "a2a", Instant.now(), Instant.now().plusSeconds(3_600_000L), "jti");

Map<String, Claim> claimsMap = JWT.decode(token).getClaims();
Mockito.when(jwtValidatorMock.validate(token, publicKey)).thenReturn(claimsMap);
Mockito.when(jwtValidatorMock.validate(token, keyPair2.getPublic())).thenReturn(claimsMap);

Pair<String, Map<String, Claim>> result = service.validate(token);

Expand All @@ -57,18 +60,18 @@ void GivenValidTokenThenOk() {

@Test
void GivenInvalidTokenThenInvalidTokenException() {
String appName = "A2A-IPA_TEST_1";
PublicKey publicKey = keyPair.getPublic();
String token = "invalidToken";
when(a2AClientLegacyPropConfig.getPublicKeysAsMap()).thenReturn(Map.of(appName, publicKey));
String token = JWTValidatorUtils.generateLegacyToken(keyPair1, "a2a", Instant.now(), Instant.now().plusSeconds(3_600_000L), "jti");

Map<String, Claim> claimsMap = JWT.decode(token).getClaims();
Mockito.when(jwtValidatorMock.validate(token, keyPair2.getPublic())).thenReturn(claimsMap);

assertThrows(Exception.class, () -> service.validate(token), "given an invalid token");
assertThrows(Exception.class, () -> service.validate("invalidToken"), "given an invalid token");
}

@Test
void GivenNonM2MAuthTokenThenInvalidTokenException() {
PublicKey publicKey = keyPair.getPublic();
String token = JWTValidatorUtils.generateLegacyToken(keyPair, "notA2A", Instant.now(), Instant.now().plusSeconds(3_600L), "jwtId");
PublicKey publicKey = keyPair2.getPublic();
String token = JWTValidatorUtils.generateLegacyToken(keyPair2, "notA2A", Instant.now(), Instant.now().plusSeconds(3_600L), "jwtId");
when(a2AClientLegacyPropConfig.getPublicKeysAsMap()).thenReturn(Map.of("A2A-IPA_TEST_1", publicKey));

Map<String, Claim> claimsMap = JWT.decode(token).getClaims();
Expand All @@ -79,8 +82,8 @@ void GivenNonM2MAuthTokenThenInvalidTokenException() {

@Test
void GivenInvalidIatThenInvalidTokenException() {
PublicKey publicKey = keyPair.getPublic();
String token = JWTValidatorUtils.generateLegacyToken(keyPair, "a2a", Instant.now().plusSeconds(3_600L), Instant.now().plusSeconds(3_600_000L), "jwtId");
PublicKey publicKey = keyPair2.getPublic();
String token = JWTValidatorUtils.generateLegacyToken(keyPair2, "a2a", Instant.now().plusSeconds(3_600L), Instant.now().plusSeconds(3_600_000L), "jwtId");


when(a2AClientLegacyPropConfig.getPublicKeysAsMap()).thenReturn(Map.of("A2A-IPA_TEST_1", publicKey));
Expand All @@ -93,8 +96,8 @@ void GivenInvalidIatThenInvalidTokenException() {

@Test
void GivenInvalidExpThenInvalidTokenException() {
PublicKey publicKey = keyPair.getPublic();
String token = JWTValidatorUtils.generateLegacyToken(keyPair, "a2a", Instant.now(), Instant.now().minusSeconds(3_600L), "jwtId");
PublicKey publicKey = keyPair2.getPublic();
String token = JWTValidatorUtils.generateLegacyToken(keyPair2, "a2a", Instant.now(), Instant.now().minusSeconds(3_600L), "jwtId");

when(a2AClientLegacyPropConfig.getPublicKeysAsMap()).thenReturn(Map.of("A2A-IPA_TEST_1", publicKey));

Expand All @@ -106,8 +109,8 @@ void GivenInvalidExpThenInvalidTokenException() {

@Test
void GivenInvalidJtiThenInvalidTokenException() {
PublicKey publicKey = keyPair.getPublic();
String token = JWTValidatorUtils.generateLegacyToken(keyPair, "a2a", Instant.now(), Instant.now().plusSeconds(3_600L), "");
PublicKey publicKey = keyPair2.getPublic();
String token = JWTValidatorUtils.generateLegacyToken(keyPair2, "a2a", Instant.now(), Instant.now().plusSeconds(3_600L), "");

when(a2AClientLegacyPropConfig.getPublicKeysAsMap()).thenReturn(Map.of("A2A-IPA_TEST_1", publicKey));

Expand Down

0 comments on commit 0ac3bd4

Please sign in to comment.