From a7c843f11ac572c0ce8d5a9590021059abecc02d Mon Sep 17 00:00:00 2001 From: Matt Collins Date: Wed, 17 Jun 2026 16:13:40 +1000 Subject: [PATCH 1/3] UID2-4739: Refactor JWTTokenProvider to accept KmsClient directly Replace the `Supplier` + `JsonObject` constructor with a pre-built `KmsClient`, and extract `buildKmsClient(JsonObject)` as a static factory. Make `JwtSigningException` a static nested class so callers no longer need an enclosing instance to construct it. Update `OperatorJWTTokenProvider` to accept explicit `issuerUrl` and `optOutUrl` strings instead of a raw config object. Wire everything up in `Main`. --- src/main/java/com/uid2/core/Main.java | 8 ++- .../uid2/core/service/JWTTokenProvider.java | 56 +++++++------------ .../service/OperatorJWTTokenProvider.java | 19 +++---- .../core/service/JWTTokenProviderTest.java | 44 ++++++--------- .../com/uid2/core/vertx/CoreVerticleTest.java | 2 +- 5 files changed, 53 insertions(+), 76 deletions(-) diff --git a/src/main/java/com/uid2/core/Main.java b/src/main/java/com/uid2/core/Main.java index 012a4aae..1401024d 100644 --- a/src/main/java/com/uid2/core/Main.java +++ b/src/main/java/com/uid2/core/Main.java @@ -4,7 +4,9 @@ import com.uid2.core.model.Constants; import com.uid2.core.model.SecretStore; import com.uid2.core.service.AttestationService; +import com.uid2.core.service.JWTTokenProvider; import com.uid2.core.service.OperatorJWTTokenProvider; +import software.amazon.awssdk.services.kms.KmsClient; import com.uid2.core.vertx.CoreVerticle; import com.uid2.core.vertx.Endpoints; import com.uid2.shared.Const; @@ -136,7 +138,11 @@ public static void main(String[] args) { attestationService.with("gcp-oidc", new GcpOidcCoreAttestationService(corePublicUrl)); - OperatorJWTTokenProvider operatorJWTTokenProvider = new OperatorJWTTokenProvider(config); + KmsClient kmsClient = JWTTokenProvider.buildKmsClient(config); + OperatorJWTTokenProvider operatorJWTTokenProvider = new OperatorJWTTokenProvider( + config.getString(Const.Config.CorePublicUrlProp), + config.getString(Const.Config.OptOutUrlProp), + new JWTTokenProvider(kmsClient)); IAttestationTokenService attestationTokenService = new AttestationTokenService( SecretStore.Global.get(Constants.AttestationEncryptionKeyName), diff --git a/src/main/java/com/uid2/core/service/JWTTokenProvider.java b/src/main/java/com/uid2/core/service/JWTTokenProvider.java index 38bf6335..05514c84 100644 --- a/src/main/java/com/uid2/core/service/JWTTokenProvider.java +++ b/src/main/java/com/uid2/core/service/JWTTokenProvider.java @@ -13,7 +13,6 @@ import java.util.Base64; import java.util.Map; import java.util.Optional; -import java.util.function.Supplier; import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider; @@ -33,12 +32,10 @@ public class JWTTokenProvider { private static final Logger LOGGER = LoggerFactory.getLogger(JWTTokenProvider.class); private static final Base64.Encoder encoder = Base64.getUrlEncoder().withoutPadding(); - private final Supplier kmsClientBuilderSupplier; - private final JsonObject config; + private final KmsClient kmsClient; - public JWTTokenProvider(JsonObject config, Supplier kmsClientBuilderSupplier) { - this.config = config; - this.kmsClientBuilderSupplier = kmsClientBuilderSupplier; + public JWTTokenProvider(KmsClient kmsClient) { + this.kmsClient = kmsClient; } public String getJWT(Instant expiresAt, Instant issuedAt, Map customClaims) throws JwtSigningException { @@ -62,13 +59,7 @@ public String getJWT(Instant expiresAt, Instant issuedAt, Map he .append(encoder.encodeToString(claimsJson.encode().getBytes(StandardCharsets.UTF_8))) .toString(); - KmsClient client = null; - try { - client = getKmsClient(this.kmsClientBuilderSupplier.get(), this.config); - } catch (URISyntaxException e) { - throw new JwtSigningException(Optional.of("Unable to get KMS Client"), e); - } - String signature = signJwtContent(client, jwtContent); + String signature = signJwtContent(this.kmsClient, jwtContent); if (signature != null && !signature.isBlank()) { return new StringBuilder() .append(jwtContent) @@ -128,44 +119,35 @@ private void addMapToJsonObject(JsonObject jsonObject, Map map) } } - private static KmsClient getKmsClient(KmsClientBuilder kmsClientBuilder, JsonObject config) throws URISyntaxException { - KmsClient client; - + public static KmsClient buildKmsClient(JsonObject config) throws URISyntaxException { String region = config.getString(KmsRegionProp, config.getString(Const.Config.AwsRegionProp)); String accessKeyId = config.getString(KmsAccessKeyIdProp); String secretAccessKey = config.getString(KmsSecretAccessKeyProp); String endpoint = config.getString(KmsEndpointProp); - if (accessKeyId != null && !accessKeyId.isBlank() && secretAccessKey != null && !secretAccessKey.isBlank()) { - AwsBasicCredentials basicCredentials = AwsBasicCredentials.create(accessKeyId, secretAccessKey); - - StaticCredentialsProvider.create(basicCredentials); + KmsClientBuilder kmsClientBuilder = KmsClient.builder(); + if (endpoint != null && !endpoint.isBlank()) { try { - if (endpoint != null && !endpoint.isBlank()) { - kmsClientBuilder.endpointOverride(new URI(endpoint)); - } - - client = kmsClientBuilder - .region(Region.of(region)) - .credentialsProvider(StaticCredentialsProvider.create(basicCredentials)) - .build(); + kmsClientBuilder.endpointOverride(new URI(endpoint)); } catch (URISyntaxException e) { - LOGGER.error("Error creating KMS Client Builder using static credentials.", e); + LOGGER.error("Error creating KMS Client Builder.", e); throw e; } - } else { - DefaultCredentialsProvider credentialsProvider = DefaultCredentialsProvider.create(); + } - client = kmsClientBuilder - .region(Region.of(region)) - .credentialsProvider(credentialsProvider) - .build(); + kmsClientBuilder.region(Region.of(region)); + + if (accessKeyId != null && !accessKeyId.isBlank() && secretAccessKey != null && !secretAccessKey.isBlank()) { + kmsClientBuilder.credentialsProvider(StaticCredentialsProvider.create( + AwsBasicCredentials.create(accessKeyId, secretAccessKey))); + } else { + kmsClientBuilder.credentialsProvider(DefaultCredentialsProvider.create()); } - return client; + return kmsClientBuilder.build(); } - public class JwtSigningException extends Exception { + public static class JwtSigningException extends Exception { public JwtSigningException(Optional message) { this(message, null); } diff --git a/src/main/java/com/uid2/core/service/OperatorJWTTokenProvider.java b/src/main/java/com/uid2/core/service/OperatorJWTTokenProvider.java index da8f5a67..9852d295 100644 --- a/src/main/java/com/uid2/core/service/OperatorJWTTokenProvider.java +++ b/src/main/java/com/uid2/core/service/OperatorJWTTokenProvider.java @@ -1,12 +1,9 @@ package com.uid2.core.service; -import com.uid2.shared.Const; import com.uid2.shared.Utils; import com.uid2.shared.auth.Role; -import io.vertx.core.json.JsonObject; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import software.amazon.awssdk.services.kms.KmsClient; import java.time.Clock; import java.time.Instant; @@ -20,16 +17,18 @@ public class OperatorJWTTokenProvider { private static final Logger LOGGER = LoggerFactory.getLogger(OperatorJWTTokenProvider.class); - private final JsonObject config; + private final String issuerUrl; + private final String optOutUrl; private final JWTTokenProvider jwtTokenProvider; private final Clock clock; - public OperatorJWTTokenProvider(JsonObject config) { - this(config, new JWTTokenProvider(config, KmsClient::builder), Clock.systemUTC()); + public OperatorJWTTokenProvider(String issuerUrl, String optOutUrl, JWTTokenProvider jwtTokenProvider) { + this(issuerUrl, optOutUrl, jwtTokenProvider, Clock.systemUTC()); } - public OperatorJWTTokenProvider(JsonObject config, JWTTokenProvider jwtTokenProvider, Clock clock) { - this.config = config; + public OperatorJWTTokenProvider(String issuerUrl, String optOutUrl, JWTTokenProvider jwtTokenProvider, Clock clock) { + this.issuerUrl = issuerUrl; + this.optOutUrl = optOutUrl; this.jwtTokenProvider = jwtTokenProvider; this.clock = clock; } @@ -45,7 +44,7 @@ public OperatorJWTTokenProvider(JsonObject config, JWTTokenProvider jwtTokenProv "iat" : the current date time */ public String getOptOutJWTToken(String operatorKey, String name, Set roles, Integer siteId, String enclaveId, String enclaveType, String operatorVersion, Instant expiresAt) throws JWTTokenProvider.JwtSigningException { - return this.getJWTToken(this.config.getString(Const.Config.CorePublicUrlProp), this.config.getString(Const.Config.OptOutUrlProp), operatorKey, name, roles, siteId, enclaveId, enclaveType, operatorVersion, expiresAt); + return this.getJWTToken(this.issuerUrl, this.optOutUrl, operatorKey, name, roles, siteId, enclaveId, enclaveType, operatorVersion, expiresAt); } /* @@ -59,7 +58,7 @@ public String getOptOutJWTToken(String operatorKey, String name, Set roles "iat" : the current date time */ public String getCoreJWTToken(String operatorKey, String name, Set roles, Integer siteId, String enclaveId, String enclaveType, String operatorVersion, Instant expiresAt) throws JWTTokenProvider.JwtSigningException { - return this.getJWTToken(this.config.getString(Const.Config.CorePublicUrlProp), this.config.getString(Const.Config.CorePublicUrlProp), operatorKey, name, roles, siteId, enclaveId, enclaveType, operatorVersion, expiresAt); + return this.getJWTToken(this.issuerUrl, this.issuerUrl, operatorKey, name, roles, siteId, enclaveId, enclaveType, operatorVersion, expiresAt); } private String getJWTToken(String issuer, String audience, String operatorKey, String name, Set roles, Integer siteId, String enclaveId, String enclaveType, String operatorVersion, Instant expiresAt) throws JWTTokenProvider.JwtSigningException { diff --git a/src/test/java/com/uid2/core/service/JWTTokenProviderTest.java b/src/test/java/com/uid2/core/service/JWTTokenProviderTest.java index a3426643..de115756 100644 --- a/src/test/java/com/uid2/core/service/JWTTokenProviderTest.java +++ b/src/test/java/com/uid2/core/service/JWTTokenProviderTest.java @@ -6,12 +6,9 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; -import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; import software.amazon.awssdk.core.SdkBytes; import software.amazon.awssdk.http.SdkHttpResponse; -import software.amazon.awssdk.regions.Region; import software.amazon.awssdk.services.kms.KmsClient; -import software.amazon.awssdk.services.kms.KmsClientBuilder; import software.amazon.awssdk.services.kms.model.KmsException; import software.amazon.awssdk.services.kms.model.SignRequest; import software.amazon.awssdk.services.kms.model.SignResponse; @@ -27,7 +24,6 @@ import static com.uid2.shared.Utils.readToEndAsString; import static org.junit.jupiter.api.Assertions.*; -import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -36,11 +32,10 @@ public class JWTTokenProviderTest { private KmsClient mockClient; private ArgumentCaptor capturedSignRequest; - private JsonObject config; @BeforeEach void setUp() throws IOException { - this.config = ((JsonObject) Json.decodeValue(openFile("/com.uid2.core/service/jwt-token-provider-test-config.json"))); + JsonObject config = (JsonObject) Json.decodeValue(openFile("/com.uid2.core/service/jwt-token-provider-test-config.json")); ConfigStore.Global.load(config); defaultHeaders.put("typ", "JWT"); defaultHeaders.put("alg", "RS256"); @@ -59,8 +54,8 @@ void getJwtReturnsValidToken() throws JWTTokenProvider.JwtSigningException { content.put("iss", "issuer"); content.put("jti", jti); - var builder = getBuilder(true, "TestSignature"); - JWTTokenProvider provider = new JWTTokenProvider(config, () -> builder); + var kmsClient = getKmsClient(true, "TestSignature"); + JWTTokenProvider provider = new JWTTokenProvider(kmsClient); Instant i = Clock.systemUTC().instant(); @@ -86,9 +81,9 @@ void getJwtReturnsValidToken() throws JWTTokenProvider.JwtSigningException { @Test void getJwtEmptySignatureThrowsException() { - var builder = getBuilder(false, ""); + var kmsClient = getKmsClient(false, ""); - JWTTokenProvider provider = new JWTTokenProvider(config, () -> builder); + JWTTokenProvider provider = new JWTTokenProvider(kmsClient); JWTTokenProvider.JwtSigningException e = assertThrows( JWTTokenProvider.JwtSigningException.class, @@ -99,9 +94,9 @@ void getJwtEmptySignatureThrowsException() { @Test void getJwtEmptySignatureEmptyResponseText() { - var builder = getBuilder(false, "", Optional.empty()); + var kmsClient = getKmsClient(false, "", Optional.empty()); - JWTTokenProvider provider = new JWTTokenProvider(config, () -> builder); + JWTTokenProvider provider = new JWTTokenProvider(kmsClient); JWTTokenProvider.JwtSigningException e = assertThrows( JWTTokenProvider.JwtSigningException.class, @@ -112,9 +107,9 @@ void getJwtEmptySignatureEmptyResponseText() { @Test void getJwtEmptySignatureNullResponseText() { - var builder = getBuilder(false, "", null); + var kmsClient = getKmsClient(false, "", null); - JWTTokenProvider provider = new JWTTokenProvider(config, () -> builder); + JWTTokenProvider provider = new JWTTokenProvider(kmsClient); JWTTokenProvider.JwtSigningException e = assertThrows( JWTTokenProvider.JwtSigningException.class, @@ -125,9 +120,9 @@ void getJwtEmptySignatureNullResponseText() { @Test void getJwtSignatureThrowsKmsException() { - var builder = getBuilder(false, "", Optional.empty()); + var kmsClient = getKmsClient(false, "", Optional.empty()); - JWTTokenProvider provider = new JWTTokenProvider(config, () -> builder); + JWTTokenProvider provider = new JWTTokenProvider(kmsClient); var ex = KmsException.builder().message("Test Error").build(); when(mockClient.sign(capturedSignRequest.capture())).thenThrow(ex); @@ -146,9 +141,9 @@ void getJwtMissingKeyInConfig() throws IOException { ConfigStore.Global.load(data); - var builder = getBuilder(false, "", Optional.empty()); + var kmsClient = getKmsClient(false, "", Optional.empty()); - JWTTokenProvider provider = new JWTTokenProvider(config, () -> builder); + JWTTokenProvider provider = new JWTTokenProvider(kmsClient); JWTTokenProvider.JwtSigningException e = assertThrows( JWTTokenProvider.JwtSigningException.class, @@ -161,11 +156,11 @@ String openFile(String filePath) throws IOException { return readToEndAsString(JWTTokenProviderTest.class.getResourceAsStream(filePath)); } - private KmsClientBuilder getBuilder(boolean isSuccessful, String signature) { - return getBuilder(isSuccessful, signature, Optional.of("Test status text")); + private KmsClient getKmsClient(boolean isSuccessful, String signature) { + return getKmsClient(isSuccessful, signature, Optional.of("Test status text")); } - private KmsClientBuilder getBuilder(boolean isSuccessful, String signature, Optional statusText) { + private KmsClient getKmsClient(boolean isSuccessful, String signature, Optional statusText) { SdkHttpResponse sdkHttpResponse = mock(SdkHttpResponse.class); when(sdkHttpResponse.isSuccessful()).thenReturn(isSuccessful); when(sdkHttpResponse.statusText()).thenReturn(statusText); @@ -178,12 +173,7 @@ private KmsClientBuilder getBuilder(boolean isSuccessful, String signature, Opti capturedSignRequest = ArgumentCaptor.forClass(SignRequest.class); when(mockClient.sign(capturedSignRequest.capture())).thenReturn(response); - KmsClientBuilder builder = mock(KmsClientBuilder.class); - when(builder.region(any(Region.class))).thenReturn(builder); - when(builder.credentialsProvider(any(AwsCredentialsProvider.class))).thenReturn(builder); - when(builder.build()).thenReturn(mockClient); - - return builder; + return mockClient; } private void assertJWT(String expectedHeader, String expectedContent, String expectedSignature, String jwt) { diff --git a/src/test/java/com/uid2/core/vertx/CoreVerticleTest.java b/src/test/java/com/uid2/core/vertx/CoreVerticleTest.java index bd0c4d7e..725e2822 100644 --- a/src/test/java/com/uid2/core/vertx/CoreVerticleTest.java +++ b/src/test/java/com/uid2/core/vertx/CoreVerticleTest.java @@ -591,7 +591,7 @@ void attestOptOutJWTCalledReturns500OnError(Vertx vertx, VertxTestContext testCo EncryptedAttestationToken encryptedAttestationToken = new EncryptedAttestationToken("test-attestation-token", Instant.ofEpochMilli(111)); when(attestationTokenService.createToken(any())).thenReturn(encryptedAttestationToken); - when(operatorJWTTokenProvider.getCoreJWTToken(anyString(), anyString(), any(), anyInt(), anyString(), any(), anyString(), any())).thenThrow(new JWTTokenProvider(null, null).new JwtSigningException(Optional.of("Test error"))); + when(operatorJWTTokenProvider.getCoreJWTToken(anyString(), anyString(), any(), anyInt(), anyString(), any(), anyString(), any())).thenThrow(new JWTTokenProvider.JwtSigningException(Optional.of("Test error"))); post(vertx, "attest", makeAttestationRequestJson("xxx", null), ar -> { assertTrue(ar.succeeded()); HttpResponse response = ar.result(); From c02781a311589f9576ecff0be65aea35a717a7b8 Mon Sep 17 00:00:00 2001 From: Matt Collins Date: Wed, 17 Jun 2026 16:44:34 +1000 Subject: [PATCH 2/3] UID2-4739: Fix JWT comment inaccuracies for sub and aud claims --- .../com/uid2/core/service/OperatorJWTTokenProvider.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/main/java/com/uid2/core/service/OperatorJWTTokenProvider.java b/src/main/java/com/uid2/core/service/OperatorJWTTokenProvider.java index 9852d295..d070ef18 100644 --- a/src/main/java/com/uid2/core/service/OperatorJWTTokenProvider.java +++ b/src/main/java/com/uid2/core/service/OperatorJWTTokenProvider.java @@ -38,7 +38,7 @@ public OperatorJWTTokenProvider(String issuerUrl, String optOutUrl, JWTTokenProv OptOut when the operator makes calls to OptOut. The claims we will add are: "iss" : the config value for issuer, something like https://core-prod.uidapi.com - "sub" : the name of the operator as registered in the Admin site + "sub" : the base64-encoded SHA-512 hash of the operator key "aud" : the url of the optout service that this token can be used with https://optout-prod.uidapi.com "exp" : the expiry date time of the token, set to be the same as the expiry of the attestation token "iat" : the current date time @@ -49,11 +49,11 @@ public String getOptOutJWTToken(String operatorKey, String name, Set roles /* Returns a JWT that is given to the operator. This is then presented by the operator to - OptOut when the operator makes calls to Core. + Core when the operator makes calls to Core. The claims we will add are: "iss" : the config value for issuer, something like https://core-prod.uidapi.com - "sub" : the name of the operator as registered in the Admin site - "aud" : the url of the optout service that this token can be used with https://core-prod.uidapi.com + "sub" : the base64-encoded SHA-512 hash of the operator key + "aud" : the url of the core service that this token can be used with https://core-prod.uidapi.com "exp" : the expiry date time of the token, set to be the same as the expiry of the attestation token "iat" : the current date time */ From 5e5e7464eac4e1e3e1465b904943f98a0b7e13f4 Mon Sep 17 00:00:00 2001 From: Matt Collins Date: Wed, 17 Jun 2026 16:44:50 +1000 Subject: [PATCH 3/3] UID2-4739: Fix swapped issuer/audience in debug log --- .../java/com/uid2/core/service/OperatorJWTTokenProvider.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main/java/com/uid2/core/service/OperatorJWTTokenProvider.java b/src/main/java/com/uid2/core/service/OperatorJWTTokenProvider.java index d070ef18..2f2be7d1 100644 --- a/src/main/java/com/uid2/core/service/OperatorJWTTokenProvider.java +++ b/src/main/java/com/uid2/core/service/OperatorJWTTokenProvider.java @@ -82,7 +82,7 @@ private String getJWTToken(String issuer, String audience, String operatorKey, S claims.put("operatorVersion", operatorVersion); claims.put("jti", UUID.randomUUID().toString()); - LOGGER.debug(String.format("Creating token with: Issuer: %s, Audience: %s, Roles: %s, SiteId: %s, EnclaveId: %s, EnclaveType: %s, OperatorVersion: %s", audience, issuer, roleString, siteId, enclaveId, enclaveType, operatorVersion)); + LOGGER.debug(String.format("Creating token with: Issuer: %s, Audience: %s, Roles: %s, SiteId: %s, EnclaveId: %s, EnclaveType: %s, OperatorVersion: %s", issuer, audience, roleString, siteId, enclaveId, enclaveType, operatorVersion)); return this.jwtTokenProvider.getJWT(expiresAt, this.clock.instant(), claims); } }