refactor: TokenGenerationService takes key ID instead of PrivateKey

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

@@ -0,0 +1,22 @@
+/*
+ *  Copyright (c) 2024 Bayerische Motoren Werke Aktiengesellschaft (BMW AG)
+ *
+ *  This program and the accompanying materials are made available under the
+ *  terms of the Apache License, Version 2.0 which is available at
+ *  https://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  SPDX-License-Identifier: Apache-2.0
+ *
+ *  Contributors:
+ *       Bayerische Motoren Werke Aktiengesellschaft (BMW AG) - initial API and implementation
+ *
+ */
+
+package org.eclipse.edc.token;
+
+import com.nimbusds.jose.JWSSigner;
+
+import java.util.function.Function;
+
+public interface JwsSignerProvider extends Function<String, JWSSigner> {
+}

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

@@ -18,6 +18,7 @@
 import com.nimbusds.jose.JOSEException;
 import com.nimbusds.jose.JWSAlgorithm;
 import com.nimbusds.jose.JWSHeader;
+import com.nimbusds.jose.JWSSigner;
 import com.nimbusds.jwt.JWTClaimsSet;
 import com.nimbusds.jwt.SignedJWT;
 import org.eclipse.edc.security.token.jwt.CryptoConverter;
@@ -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;
+import java.util.function.Function;
 
 public class JwtGenerationService implements TokenGenerationService {
 
+    private final Function<String, JWSSigner> jwsGeneratorFunction;
+
+    public JwtGenerationService(JwsSignerProvider jwsGeneratorFunction) {
+
+        this.jwsGeneratorFunction = jwsGeneratorFunction;
+    }
 
     @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 tokenSigner = jwsGeneratorFunction.apply(privateKeyId);
+        if (tokenSigner == null) {
+            return Result.failure("JWSSigner cannot be generated for private key '%s'".formatted(privateKeyId));
         }
 
-        var tokenSigner = CryptoConverter.createSignerFor(privateKey);
         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,12 @@
 
 package org.eclipse.edc.token;
 
+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.EdcException;
 import org.eclipse.edc.spi.system.ServiceExtension;
 import org.eclipse.edc.token.spi.TokenDecoratorRegistry;
 import org.eclipse.edc.token.spi.TokenValidationRulesRegistry;
@@ -31,6 +35,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 +52,12 @@ public TokenValidationService validationService() {
     public TokenDecoratorRegistry tokenDecoratorRegistry() {
         return new TokenDecoratorRegistryImpl();
     }
+
+    @Provider(isDefault = true)
+    public JwsSignerProvider defaultSignerProvider() {
+        return privateKeyId -> {
+            var pk = privateKeyResolver.resolvePrivateKey(privateKeyId).orElseThrow(f -> new EdcException("This EDC instance is not operational due to the following error: %s".formatted(f.getFailureDetail())));
+            return CryptoConverter.createSignerFor(pk);
+        };
+    }
 }

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

@@ -23,6 +23,7 @@
 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.token.spi.TokenDecorator;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -39,26 +40,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 CryptoConverter.createSignerFor(pk);
+                } catch (JOSEException e) {
+                    throw new RuntimeException(e);
+                }
+            }
+            return null;
+        });
     }
 
     @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 +90,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

@@ -24,15 +24,13 @@
 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.JwsSignerProvider;
 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.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.JwsSignerProvider;
 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,17 +16,18 @@
 
 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.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.security.Vault;
 import org.eclipse.edc.spi.system.ServiceExtension;
 import org.eclipse.edc.spi.system.ServiceExtensionContext;
+import org.eclipse.edc.token.JwsSignerProvider;
 import org.eclipse.edc.token.JwtGenerationService;
 import org.eclipse.edc.transaction.spi.TransactionContext;
 
@@ -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;