refactor: TokenGenerationService takes key ID instead of PrivateKey (#…

…4395)

* refactor: TokenGenerationService takes key ID instead of PrivateKey

* move JwsSignerProvider to its own SPI

* Update spi/common/jwt-signer-spi/src/main/java/org/eclipse/edc/jwt/signer/spi/JwsSignerProvider.java

Co-authored-by: Enrico Risa <enrico.risa@gmail.com>

* Update core/common/token-core/src/main/java/org/eclipse/edc/token/JwtGenerationService.java

Co-authored-by: ndr_brt <andrea.bertagnolli@gmail.com>

* pr remarks

* fix tests

---------

Co-authored-by: Enrico Risa <enrico.risa@gmail.com>
Co-authored-by: ndr_brt <andrea.bertagnolli@gmail.com>

core/common/token-core/build.gradle.kts

@@ -21,6 +21,7 @@ dependencies {
     api(project(":spi:common:keys-spi"))
     api(project(":spi:common:token-spi"))
     api(project(":spi:common:jwt-spi"))
+    api(project(":spi:common:jwt-signer-spi"))
 
     implementation(project(":core:common:lib:crypto-common-lib")) // for the CryptoConverter
     implementation(libs.nimbus.jwt)

core/common/token-core/src/main/java/org/eclipse/edc/token/JwtGenerationService.java

@@ -20,6 +20,7 @@
 import com.nimbusds.jose.JWSHeader;
 import com.nimbusds.jwt.JWTClaimsSet;
 import com.nimbusds.jwt.SignedJWT;
+import org.eclipse.edc.jwt.signer.spi.JwsSignerProvider;
 import org.eclipse.edc.security.token.jwt.CryptoConverter;
 import org.eclipse.edc.spi.EdcException;
 import org.eclipse.edc.spi.iam.TokenParameters;
@@ -29,25 +30,29 @@
 import org.eclipse.edc.token.spi.TokenGenerationService;
 import org.jetbrains.annotations.NotNull;
 
-import java.security.PrivateKey;
 import java.text.ParseException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Map;
-import java.util.function.Supplier;
 
 public class JwtGenerationService implements TokenGenerationService {
 
+    private final JwsSignerProvider jwsGeneratorFunction;
+
+    public JwtGenerationService(JwsSignerProvider jwsSignerProvider) {
+
+        this.jwsGeneratorFunction = jwsSignerProvider;
+    }
 
     @Override
-    public Result<TokenRepresentation> generate(Supplier<PrivateKey> privateKeySupplier, @NotNull TokenDecorator... decorators) {
+    public Result<TokenRepresentation> generate(String privateKeyId, @NotNull TokenDecorator... decorators) {
 
-        var privateKey = privateKeySupplier.get();
-        if (privateKey == null) {
-            return Result.failure("PrivateKey cannot be resolved.");
+        var tokenSignerResult = jwsGeneratorFunction.createJwsSigner(privateKeyId);
+        if (tokenSignerResult.failed()) {
+            return Result.failure("JWSSigner cannot be generated for private key '%s': %s".formatted(privateKeyId, tokenSignerResult.getFailureDetail()));
         }
 
-        var tokenSigner = CryptoConverter.createSignerFor(privateKey);
+        var tokenSigner = tokenSignerResult.getContent();
         var jwsAlgorithm = CryptoConverter.getRecommendedAlgorithm(tokenSigner);
 
         var bldr = TokenParameters.Builder.newInstance();

core/common/token-core/src/main/java/org/eclipse/edc/token/TokenServicesExtension.java

@@ -14,8 +14,13 @@
 
 package org.eclipse.edc.token;
 
+import org.eclipse.edc.jwt.signer.spi.JwsSignerProvider;
+import org.eclipse.edc.keys.spi.PrivateKeyResolver;
 import org.eclipse.edc.runtime.metamodel.annotation.Extension;
+import org.eclipse.edc.runtime.metamodel.annotation.Inject;
 import org.eclipse.edc.runtime.metamodel.annotation.Provider;
+import org.eclipse.edc.security.token.jwt.CryptoConverter;
+import org.eclipse.edc.spi.result.Result;
 import org.eclipse.edc.spi.system.ServiceExtension;
 import org.eclipse.edc.token.spi.TokenDecoratorRegistry;
 import org.eclipse.edc.token.spi.TokenValidationRulesRegistry;
@@ -31,6 +36,9 @@
 public class TokenServicesExtension implements ServiceExtension {
     public static final String NAME = "Token Services Extension";
 
+    @Inject
+    private PrivateKeyResolver privateKeyResolver;
+
     @Provider
     public TokenValidationRulesRegistry tokenValidationRulesRegistry() {
         return new TokenValidationRulesRegistryImpl();
@@ -45,4 +53,11 @@ public TokenValidationService validationService() {
     public TokenDecoratorRegistry tokenDecoratorRegistry() {
         return new TokenDecoratorRegistryImpl();
     }
+
+    @Provider(isDefault = true)
+    public JwsSignerProvider defaultSignerProvider() {
+        // default implementation: resolve the private key (from vault of config) and create a JWSSigner based on its algorithm
+        return privateKeyId -> privateKeyResolver.resolvePrivateKey(privateKeyId)
+                .compose(pk -> Result.ofThrowable(() -> CryptoConverter.createSignerFor(pk)));
+    }
 }

core/common/token-core/src/test/java/org/eclipse/edc/token/JwtGenerationServiceTest.java

@@ -23,6 +23,8 @@
 import com.nimbusds.jose.jwk.gen.RSAKeyGenerator;
 import com.nimbusds.jose.util.Base64URL;
 import com.nimbusds.jwt.SignedJWT;
+import org.eclipse.edc.security.token.jwt.CryptoConverter;
+import org.eclipse.edc.spi.result.Result;
 import org.eclipse.edc.token.spi.TokenDecorator;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -39,26 +41,31 @@
 
 class JwtGenerationServiceTest {
 
+    public static final String TEST_KEY_ID = "test-key-id";
     private RSAKey keys;
     private JwtGenerationService tokenGenerationService;
 
     @BeforeEach
     void setUp() throws JOSEException {
         keys = testKey();
-        tokenGenerationService = new JwtGenerationService();
+        tokenGenerationService = new JwtGenerationService(keyId -> {
+            if (TEST_KEY_ID.equals(keyId)) {
+                try {
+                    var pk = keys.toPrivateKey();
+                    return Result.success(CryptoConverter.createSignerFor(pk));
+                } catch (JOSEException e) {
+                    return Result.failure(e.getMessage());
+                }
+            }
+            return Result.failure("key not found");
+        });
     }
 
     @Test
     void verifyTokenGeneration() throws ParseException, JOSEException {
         var decorator = testDecorator();
 
-        var result = tokenGenerationService.generate(() -> {
-            try {
-                return keys.toPrivateKey();
-            } catch (JOSEException e) {
-                throw new RuntimeException(e);
-            }
-        }, decorator);
+        var result = tokenGenerationService.generate(TEST_KEY_ID, decorator);
 
         assertThat(result.succeeded()).isTrue();
         var token = result.getContent().getToken();
@@ -84,7 +91,7 @@ void verifyTokenGeneration() throws ParseException, JOSEException {
     void shouldFail_whenPrivateKeyCannotBeResolved() {
         var decorator = testDecorator();
 
-        var result = tokenGenerationService.generate(() -> null, decorator);
+        var result = tokenGenerationService.generate("not-exist-key", decorator);
 
         assertThat(result.failed()).isTrue();
     }

core/data-plane/data-plane-core/src/main/java/org/eclipse/edc/connector/dataplane/framework/DataPlaneDefaultIamServicesExtension.java

@@ -18,21 +18,19 @@
 import org.eclipse.edc.connector.dataplane.spi.iam.DataPlaneAccessControlService;
 import org.eclipse.edc.connector.dataplane.spi.iam.DataPlaneAccessTokenService;
 import org.eclipse.edc.connector.dataplane.spi.store.AccessTokenDataStore;
+import org.eclipse.edc.jwt.signer.spi.JwsSignerProvider;
 import org.eclipse.edc.keys.spi.LocalPublicKeyService;
 import org.eclipse.edc.keys.spi.PrivateKeyResolver;
 import org.eclipse.edc.runtime.metamodel.annotation.Extension;
 import org.eclipse.edc.runtime.metamodel.annotation.Inject;
 import org.eclipse.edc.runtime.metamodel.annotation.Provider;
 import org.eclipse.edc.runtime.metamodel.annotation.Setting;
-import org.eclipse.edc.spi.monitor.Monitor;
 import org.eclipse.edc.spi.result.Result;
 import org.eclipse.edc.spi.system.ServiceExtension;
 import org.eclipse.edc.spi.system.ServiceExtensionContext;
 import org.eclipse.edc.token.JwtGenerationService;
 import org.eclipse.edc.token.spi.TokenValidationService;
-import org.jetbrains.annotations.NotNull;
 
-import java.security.PrivateKey;
 import java.util.function.Supplier;
 
 
@@ -54,6 +52,8 @@ public class DataPlaneDefaultIamServicesExtension implements ServiceExtension {
     private PrivateKeyResolver privateKeyResolver;
     @Inject
     private LocalPublicKeyService localPublicKeyService;
+    @Inject
+    private JwsSignerProvider jwsSignerProvider;
 
     @Override
     public String name() {
@@ -71,21 +71,13 @@ public DataPlaneAccessTokenService defaultAccessTokenService(ServiceExtensionCon
         var tokenVerifierPublicKeyAlias = context.getConfig().getString(TOKEN_VERIFIER_PUBLIC_KEY_ALIAS);
         var tokenSignerPrivateKeyAlias = context.getConfig().getString(TOKEN_SIGNER_PRIVATE_KEY_ALIAS);
         var monitor = context.getMonitor().withPrefix("DataPlane IAM");
-        return new DefaultDataPlaneAccessTokenServiceImpl(new JwtGenerationService(),
-                accessTokenDataStore, monitor, getPrivateKeySupplier(tokenSignerPrivateKeyAlias, monitor),
+        return new DefaultDataPlaneAccessTokenServiceImpl(new JwtGenerationService(jwsSignerProvider),
+                accessTokenDataStore, monitor, () -> tokenSignerPrivateKeyAlias,
                 publicKeyIdSupplier(tokenVerifierPublicKeyAlias), tokenValidationService, localPublicKeyService);
     }
 
     private Supplier<String> publicKeyIdSupplier(String tokenVerifierPublicKeyAlias) {
         return () -> tokenVerifierPublicKeyAlias;
     }
 
-    @NotNull
-    private Supplier<PrivateKey> getPrivateKeySupplier(String tokenSignerPrivateKeyAlias, Monitor monitor) {
-        return () -> privateKeyResolver.resolvePrivateKey(tokenSignerPrivateKeyAlias)
-                .orElse(f -> {
-                    monitor.warning("Cannot resolve private key: " + f.getFailureDetail());
-                    return null;
-                });
-    }
 }

core/data-plane/data-plane-core/src/main/java/org/eclipse/edc/connector/dataplane/framework/iam/DefaultDataPlaneAccessTokenServiceImpl.java

@@ -32,7 +32,6 @@
 import org.eclipse.edc.token.spi.TokenValidationRule;
 import org.eclipse.edc.token.spi.TokenValidationService;
 
-import java.security.PrivateKey;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -56,22 +55,22 @@ public class DefaultDataPlaneAccessTokenServiceImpl implements DataPlaneAccessTo
     private final TokenGenerationService tokenGenerationService;
     private final AccessTokenDataStore accessTokenDataStore;
     private final Monitor monitor;
-    private final Supplier<PrivateKey> privateKeySupplier;
+    private final Supplier<String> privateKeyIdSupplier;
     private final Supplier<String> publicKeyIdSupplier;
     private final TokenValidationService tokenValidationService;
     private final PublicKeyResolver publicKeyResolver;
 
     public DefaultDataPlaneAccessTokenServiceImpl(TokenGenerationService tokenGenerationService,
                                                   AccessTokenDataStore accessTokenDataStore,
                                                   Monitor monitor,
-                                                  Supplier<PrivateKey> privateKeySupplier,
+                                                  Supplier<String> privateKeyIdSupplier,
                                                   Supplier<String> publicKeyIdSupplier,
                                                   TokenValidationService tokenValidationService,
                                                   PublicKeyResolver publicKeyResolver) {
         this.tokenGenerationService = tokenGenerationService;
         this.accessTokenDataStore = accessTokenDataStore;
         this.monitor = monitor;
-        this.privateKeySupplier = privateKeySupplier;
+        this.privateKeyIdSupplier = privateKeyIdSupplier;
         this.publicKeyIdSupplier = publicKeyIdSupplier;
         this.tokenValidationService = tokenValidationService;
         this.publicKeyResolver = publicKeyResolver;
@@ -106,9 +105,9 @@ public Result<TokenRepresentation> obtainToken(TokenParameters parameters, DataA
             allDecorators.add(tokenIdDecorator);
         }
 
-        var tokenResult = tokenGenerationService.generate(privateKeySupplier, allDecorators.toArray(new TokenDecorator[0]));
+        var tokenResult = tokenGenerationService.generate(privateKeyIdSupplier.get(), allDecorators.toArray(new TokenDecorator[0]));
         if (tokenResult.failed()) {
-            return tokenResult.mapTo();
+            return tokenResult.mapEmpty();
         }
 
         // store a record of the token for future reference. We'll need that when we resolve the AccessTokenData later.

extensions/common/iam/identity-trust/identity-trust-core/src/main/java/org/eclipse/edc/iam/identitytrust/core/DcpDefaultServicesExtension.java

@@ -25,24 +25,21 @@
 import org.eclipse.edc.iam.identitytrust.spi.verification.SignatureSuiteRegistry;
 import org.eclipse.edc.iam.identitytrust.sts.embedded.EmbeddedSecureTokenService;
 import org.eclipse.edc.iam.verifiablecredentials.spi.validation.TrustedIssuerRegistry;
-import org.eclipse.edc.keys.spi.PrivateKeyResolver;
+import org.eclipse.edc.jwt.signer.spi.JwsSignerProvider;
 import org.eclipse.edc.runtime.metamodel.annotation.Extension;
 import org.eclipse.edc.runtime.metamodel.annotation.Inject;
 import org.eclipse.edc.runtime.metamodel.annotation.Provider;
 import org.eclipse.edc.runtime.metamodel.annotation.Setting;
-import org.eclipse.edc.spi.EdcException;
 import org.eclipse.edc.spi.iam.AudienceResolver;
 import org.eclipse.edc.spi.iam.ClaimToken;
 import org.eclipse.edc.spi.result.Result;
 import org.eclipse.edc.spi.system.ServiceExtension;
 import org.eclipse.edc.spi.system.ServiceExtensionContext;
 import org.eclipse.edc.token.JwtGenerationService;
 
-import java.security.PrivateKey;
 import java.time.Clock;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
-import java.util.function.Supplier;
 
 import static org.eclipse.edc.spi.result.Result.failure;
 import static org.eclipse.edc.spi.result.Result.success;
@@ -63,7 +60,7 @@ public class DcpDefaultServicesExtension implements ServiceExtension {
     @Inject
     private Clock clock;
     @Inject
-    private PrivateKeyResolver privateKeyResolver;
+    private JwsSignerProvider externalSigner;
 
     @Provider(isDefault = true)
     public SecureTokenService createDefaultTokenService(ServiceExtensionContext context) {
@@ -78,10 +75,9 @@ public SecureTokenService createDefaultTokenService(ServiceExtensionContext cont
 
 
         var publicKeyId = context.getSetting(STS_PUBLIC_KEY_ID, null);
-        var privKeyAlias = context.getSetting(STS_PRIVATE_KEY_ALIAS, null);
+        var privateKeyAlias = context.getSetting(STS_PRIVATE_KEY_ALIAS, null);
 
-        Supplier<PrivateKey> supplier = () -> privateKeyResolver.resolvePrivateKey(privKeyAlias).orElseThrow(f -> new EdcException("This EDC instance is not operational due to the following error: %s".formatted(f.getFailureDetail())));
-        return new EmbeddedSecureTokenService(new JwtGenerationService(), supplier, () -> publicKeyId, clock, TimeUnit.MINUTES.toSeconds(tokenExpiration));
+        return new EmbeddedSecureTokenService(new JwtGenerationService(externalSigner), () -> privateKeyAlias, () -> publicKeyId, clock, TimeUnit.MINUTES.toSeconds(tokenExpiration));
     }
 
     @Provider(isDefault = true)

extensions/common/iam/identity-trust/identity-trust-sts/identity-trust-sts-core/build.gradle.kts

@@ -14,6 +14,7 @@ dependencies {
 
     testImplementation(testFixtures(project(":spi:common:identity-trust-sts-spi")))
     testImplementation(project(":core:common:lib:boot-lib"))
+    testImplementation(project(":core:common:lib:crypto-common-lib"))
     testImplementation(project(":core:common:lib:keys-lib"))
     testImplementation(project(":core:common:junit"))
     testImplementation(libs.nimbus.jwt)

extensions/common/iam/identity-trust/identity-trust-sts/identity-trust-sts-core/src/main/java/org/eclipse/edc/iam/identitytrust/sts/defaults/StsDefaultServicesExtension.java

@@ -16,10 +16,11 @@
 
 import org.eclipse.edc.iam.identitytrust.sts.defaults.service.StsClientServiceImpl;
 import org.eclipse.edc.iam.identitytrust.sts.defaults.service.StsClientTokenGeneratorServiceImpl;
+import org.eclipse.edc.iam.identitytrust.sts.spi.model.StsClient;
 import org.eclipse.edc.iam.identitytrust.sts.spi.service.StsClientService;
 import org.eclipse.edc.iam.identitytrust.sts.spi.service.StsClientTokenGeneratorService;
 import org.eclipse.edc.iam.identitytrust.sts.spi.store.StsClientStore;
-import org.eclipse.edc.keys.spi.PrivateKeyResolver;
+import org.eclipse.edc.jwt.signer.spi.JwsSignerProvider;
 import org.eclipse.edc.runtime.metamodel.annotation.Extension;
 import org.eclipse.edc.runtime.metamodel.annotation.Inject;
 import org.eclipse.edc.runtime.metamodel.annotation.Provider;
@@ -53,7 +54,7 @@ public class StsDefaultServicesExtension implements ServiceExtension {
     private Vault vault;
 
     @Inject
-    private PrivateKeyResolver privateKeyResolver;
+    private JwsSignerProvider jwsSignerProvider;
 
     @Inject
     private Clock clock;
@@ -67,8 +68,8 @@ public String name() {
     public StsClientTokenGeneratorService clientTokenService(ServiceExtensionContext context) {
         var tokenExpiration = context.getSetting(STS_TOKEN_EXPIRATION, DEFAULT_STS_TOKEN_EXPIRATION_MIN);
         return new StsClientTokenGeneratorServiceImpl(
-                (client) -> new JwtGenerationService(),
-                (client) -> privateKeyResolver.resolvePrivateKey(client.getPrivateKeyAlias()).orElse(null),
+                (client) -> new JwtGenerationService(jwsSignerProvider),
+                StsClient::getPrivateKeyAlias,
                 clock,
                 TimeUnit.MINUTES.toSeconds(tokenExpiration));
     }

extensions/common/iam/identity-trust/identity-trust-sts/identity-trust-sts-core/src/main/java/org/eclipse/edc/iam/identitytrust/sts/defaults/service/StsClientTokenGeneratorServiceImpl.java

@@ -22,7 +22,6 @@
 import org.eclipse.edc.spi.iam.TokenRepresentation;
 import org.eclipse.edc.spi.result.ServiceResult;
 
-import java.security.PrivateKey;
 import java.time.Clock;
 import java.util.Collections;
 import java.util.HashMap;
@@ -42,10 +41,10 @@ public class StsClientTokenGeneratorServiceImpl implements StsClientTokenGenerat
 
     private final long tokenExpiration;
     private final StsTokenGenerationProvider tokenGenerationProvider;
-    private final Function<StsClient, PrivateKey> keyFunction;
+    private final Function<StsClient, String> keyFunction;
     private final Clock clock;
 
-    public StsClientTokenGeneratorServiceImpl(StsTokenGenerationProvider tokenGenerationProvider, Function<StsClient, PrivateKey> keyFunction, Clock clock, long tokenExpiration) {
+    public StsClientTokenGeneratorServiceImpl(StsTokenGenerationProvider tokenGenerationProvider, Function<StsClient, String> keyFunction, Clock clock, long tokenExpiration) {
         this.tokenGenerationProvider = tokenGenerationProvider;
         this.keyFunction = keyFunction;
         this.clock = clock;

extensions/common/iam/identity-trust/identity-trust-sts/identity-trust-sts-core/src/test/java/org/eclipse/edc/iam/identitytrust/sts/defaults/StsClientTokenIssuanceIntegrationTest.java

@@ -20,6 +20,7 @@
 import org.eclipse.edc.iam.identitytrust.sts.defaults.service.StsClientServiceImpl;
 import org.eclipse.edc.iam.identitytrust.sts.defaults.service.StsClientTokenGeneratorServiceImpl;
 import org.eclipse.edc.iam.identitytrust.sts.defaults.store.InMemoryStsClientStore;
+import org.eclipse.edc.iam.identitytrust.sts.spi.model.StsClient;
 import org.eclipse.edc.iam.identitytrust.sts.spi.model.StsClientTokenAdditionalParams;
 import org.eclipse.edc.junit.annotations.ComponentTest;
 import org.eclipse.edc.keys.KeyParserRegistryImpl;
@@ -28,6 +29,8 @@
 import org.eclipse.edc.keys.keyparsers.PemParser;
 import org.eclipse.edc.keys.spi.KeyParserRegistry;
 import org.eclipse.edc.keys.spi.PrivateKeyResolver;
+import org.eclipse.edc.security.token.jwt.CryptoConverter;
+import org.eclipse.edc.spi.result.Result;
 import org.eclipse.edc.spi.security.Vault;
 import org.eclipse.edc.token.JwtGenerationService;
 import org.eclipse.edc.transaction.spi.NoopTransactionContext;
@@ -70,8 +73,9 @@ void setup() {
         privateKeyResolver = new VaultPrivateKeyResolver(keyParserRegistry, vault, mock(), mock());
 
         tokenGeneratorService = new StsClientTokenGeneratorServiceImpl(
-                client -> new JwtGenerationService(),
-                stsClient -> privateKeyResolver.resolvePrivateKey(stsClient.getPrivateKeyAlias()).orElse(null),
+                client -> new JwtGenerationService(keyId -> privateKeyResolver.resolvePrivateKey(keyId)
+                        .compose(pk -> Result.ofThrowable(() -> CryptoConverter.createSignerFor(pk)))),
+                StsClient::getPrivateKeyAlias,
                 Clock.systemUTC(), 60 * 5);
 
     }

extensions/common/iam/identity-trust/identity-trust-sts/identity-trust-sts-core/src/test/java/org/eclipse/edc/iam/identitytrust/sts/defaults/service/StsClientTokenGeneratorServiceImplTest.java

@@ -43,7 +43,7 @@ public class StsClientTokenGeneratorServiceImplTest {
 
     @BeforeEach
     void setup() {
-        clientTokenService = new StsClientTokenGeneratorServiceImpl(tokenGenerationProvider, (client) -> mock(), Clock.systemUTC(), TOKEN_EXPIRATION);
+        clientTokenService = new StsClientTokenGeneratorServiceImpl(tokenGenerationProvider, (client) -> "test-key-id", Clock.systemUTC(), TOKEN_EXPIRATION);
     }
 
     @Test