NimbusJwtEncoder should simplify constructing with javax.security.Keys Signed-off-by: Suraj Bhadrike <[email protected]>

Signed-off-by: surajbh <[email protected]>

Author: surajbh123

oauth2/oauth2-jose/src/main/java/org/springframework/security/oauth2/jwt/NimbusJwtEncoder.java

@@ -16,42 +16,55 @@
 
 package org.springframework.security.oauth2.jwt;
 
-import java.net.URI;
-import java.net.URL;
-import java.time.Instant;
-import java.util.ArrayList;
-import java.util.Date;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.concurrent.ConcurrentHashMap;
-
 import com.nimbusds.jose.JOSEException;
 import com.nimbusds.jose.JOSEObjectType;
 import com.nimbusds.jose.JWSAlgorithm;
 import com.nimbusds.jose.JWSHeader;
 import com.nimbusds.jose.JWSSigner;
 import com.nimbusds.jose.crypto.factories.DefaultJWSSignerFactory;
+import com.nimbusds.jose.jwk.Curve;
+import com.nimbusds.jose.jwk.ECKey;
 import com.nimbusds.jose.jwk.JWK;
 import com.nimbusds.jose.jwk.JWKMatcher;
 import com.nimbusds.jose.jwk.JWKSelector;
+import com.nimbusds.jose.jwk.JWKSet;
 import com.nimbusds.jose.jwk.KeyType;
 import com.nimbusds.jose.jwk.KeyUse;
+import com.nimbusds.jose.jwk.OctetSequenceKey;
+import com.nimbusds.jose.jwk.RSAKey;
+import com.nimbusds.jose.jwk.source.ImmutableJWKSet;
 import com.nimbusds.jose.jwk.source.JWKSource;
 import com.nimbusds.jose.proc.SecurityContext;
 import com.nimbusds.jose.produce.JWSSignerFactory;
 import com.nimbusds.jose.util.Base64;
 import com.nimbusds.jose.util.Base64URL;
 import com.nimbusds.jwt.JWTClaimsSet;
 import com.nimbusds.jwt.SignedJWT;
-
 import org.springframework.core.convert.converter.Converter;
+import org.springframework.security.oauth2.jose.jws.MacAlgorithm;
 import org.springframework.security.oauth2.jose.jws.SignatureAlgorithm;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
 import org.springframework.util.StringUtils;
 
+import javax.crypto.SecretKey;
+import java.net.URI;
+import java.net.URL;
+import java.security.KeyPair;
+import java.security.interfaces.ECPrivateKey;
+import java.security.interfaces.ECPublicKey;
+import java.security.interfaces.RSAPrivateKey;
+import java.time.Instant;
+import java.util.ArrayList;
+import java.util.Date;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Stream;
+
 /**
  * An implementation of a {@link JwtEncoder} that encodes a JSON Web Token (JWT) using the
  * JSON Web Signature (JWS) Compact Serialization format. The private/secret key used for
@@ -79,7 +92,7 @@ public final class NimbusJwtEncoder implements JwtEncoder {
 
 	private static final String ENCODING_ERROR_MESSAGE_TEMPLATE = "An error occurred while attempting to encode the Jwt: %s";
 
-	private static final JwsHeader DEFAULT_JWS_HEADER = JwsHeader.with(SignatureAlgorithm.RS256).build();
+	private static JwsHeader DEFAULT_JWS_HEADER;
 
 	private static final JWSSignerFactory JWS_SIGNER_FACTORY = new DefaultJWSSignerFactory();
 
@@ -102,6 +115,12 @@ public final class NimbusJwtEncoder implements JwtEncoder {
 	public NimbusJwtEncoder(JWKSource<SecurityContext> jwkSource) {
 		Assert.notNull(jwkSource, "jwkSource cannot be null");
 		this.jwkSource = jwkSource;
+		this.DEFAULT_JWS_HEADER = JwsHeader.with(SignatureAlgorithm.RS256).build();
+	}
+	public NimbusJwtEncoder(JWKSource<SecurityContext> jwkSource, JwsHeader jwsHeader) {
+		Assert.notNull(jwkSource, "jwkSource cannot be null");
+		this.jwkSource = jwkSource;
+		this.DEFAULT_JWS_HEADER = jwsHeader;
 	}
 
 	/**
@@ -369,4 +388,295 @@ private static URI convertAsURI(String header, URL url) {
 		}
 	}
 
+	/**
+	 * Creates a builder for constructing a {@link NimbusJwtEncoder} using the provided
+	 * {@link SecretKey}.
+	 *
+	 * @param secretKey the {@link SecretKey} to use for signing JWTs
+	 * @return a {@link SecretKeyJwtEncoderBuilder} for further configuration
+	 * @since // TODO: Update version
+	 */
+	public static SecretKeyJwtEncoderBuilder withSecretKey(SecretKey secretKey) {
+		Assert.notNull(secretKey, "secretKey cannot be null");
+		return new SecretKeyJwtEncoderBuilder(secretKey);
+	}
+
+	/**
+	 * Creates a builder for constructing a {@link NimbusJwtEncoder} using the provided
+	 * {@link KeyPair}. The key pair must contain either an {@link RSAKey} or an
+	 * {@link ECKey}.
+	 *
+	 * @param keyPair the {@link KeyPair} to use for signing JWTs
+	 * @return a {@link KeyPairJwtEncoderBuilder} for further configuration
+	 * @since // TODO: Update version
+	 */
+	public static KeyPairJwtEncoderBuilder withKeyPair(KeyPair keyPair) {
+		Assert.notNull(keyPair, "keyPair cannot be null");
+		Assert.notNull(keyPair.getPrivate(), "keyPair must contain a private key");
+		Assert.notNull(keyPair.getPublic(), "keyPair must contain a public key");
+		Assert.isTrue(keyPair.getPrivate() instanceof java.security.interfaces.RSAKey || keyPair.getPrivate() instanceof java.security.interfaces.ECKey,
+				"keyPair must be an RSAKey or an ECKey");
+		return new KeyPairJwtEncoderBuilder(keyPair);
+	}
+
+	/**
+	 * A builder for creating {@link NimbusJwtEncoder} instances configured with a
+	 * {@link SecretKey}.
+	 *
+	 * @since // TODO: Update version
+	 */
+	public static final class SecretKeyJwtEncoderBuilder {
+
+		private final SecretKey secretKey;
+
+		private String keyId;
+
+		private JWSAlgorithm jwsAlgorithm = JWSAlgorithm.HS256;
+
+		private Converter<List<JWK>, JWK> jwkSelector;
+
+		private JwsHeader jwsHeader;
+
+		private SecretKeyJwtEncoderBuilder(SecretKey secretKey) {
+			this.secretKey = secretKey;
+		}
+
+		/**
+		 * Sets the JWS algorithm to use for signing. Defaults to {@link JWSAlgorithm#HS256}.
+		 * Must be an HMAC-based algorithm (HS256, HS384, or HS512).
+		 *
+		 * @param macAlgorithm the {@link MacAlgorithm} to use
+		 * @return this builder instance for method chaining
+		 */
+		public SecretKeyJwtEncoderBuilder macAlgorithm(MacAlgorithm macAlgorithm) {
+			Assert.notNull(macAlgorithm, "macAlgorithm cannot be null");
+			this.jwsAlgorithm = JWSAlgorithm.parse(macAlgorithm.getName());
+			return this;
+		}
+
+		/**
+		 * Sets the key ID ({@code kid}) to be included in the JWK and potentially the JWS
+		 * header.
+		 *
+		 * @param keyId the key identifier
+		 * @return this builder instance for method chaining
+		 */
+		public SecretKeyJwtEncoderBuilder keyId(String keyId) {
+			this.keyId = keyId;
+			return this;
+		}
+
+		/**
+		 * Configures the {@link Converter} used to select the JWK when multiple keys match
+		 * the header criteria. This is generally not needed for single-key setups but is
+		 * provided for consistency.
+		 *
+		 * @param jwkSelector the {@link Converter} to select a {@link JWK}
+		 * @return this builder instance for method chaining
+		 * @since // TODO: Update version
+		 */
+		public SecretKeyJwtEncoderBuilder jwkSelector(Converter<List<JWK>, JWK> jwkSelector) {
+			Assert.notNull(jwkSelector, "jwkSelector cannot be null");
+			this.jwkSelector = jwkSelector;
+			return this;
+		}
+
+
+		/**
+		 * Builds the {@link NimbusJwtEncoder} instance.
+		 *
+		 * @return the configured {@link NimbusJwtEncoder}
+		 * @throws IllegalStateException if the configured JWS algorithm is not compatible
+		 *                               with a {@link SecretKey}.
+		 */
+		public NimbusJwtEncoder build() {
+			this.jwsAlgorithm = this.jwsAlgorithm  != null ? this.jwsAlgorithm : JWSAlgorithm.HS256;
+			Assert.state(JWSAlgorithm.Family.HMAC_SHA.contains(this.jwsAlgorithm),
+					() -> "The algorithm '" + this.jwsAlgorithm + "' is not compatible with a SecretKey. "
+							+ "Please use one of the HS256, HS384, or HS512 algorithms.");
+
+			OctetSequenceKey.Builder builder = new OctetSequenceKey.Builder(this.secretKey)
+					.keyUse(KeyUse.SIGNATURE)
+					.algorithm(this.jwsAlgorithm);
+
+			this.jwsHeader = JwsHeader.with(MacAlgorithm.from(this.jwsAlgorithm.getName())).build();
+
+			if (StringUtils.hasText(this.keyId)) {
+				builder.keyID(this.keyId);
+			}
+
+			OctetSequenceKey jwk = builder.build();
+			JWKSource<SecurityContext> jwkSource = new ImmutableJWKSet<>(new JWKSet(jwk));
+			NimbusJwtEncoder encoder = this.jwsHeader  != null
+					? new NimbusJwtEncoder(jwkSource, this.jwsHeader)
+					: new NimbusJwtEncoder(jwkSource);
+			if (this.jwkSelector != null) {
+				encoder.setJwkSelector(this.jwkSelector);
+			} else {
+				encoder.setJwkSelector(jwkSet -> jwkSet.stream().findFirst().orElse(null));
+			}
+			return encoder;
+		}
+
+	}
+
+	/**
+	 * A builder for creating {@link NimbusJwtEncoder} instances configured with a
+	 * {@link KeyPair}.
+	 *
+	 * @since // TODO: Update version
+	 */
+	public static final class KeyPairJwtEncoderBuilder {
+
+		private final KeyPair keyPair;
+
+		private String keyId;
+
+		private JWSAlgorithm jwsAlgorithm;
+
+		private Converter<List<JWK>, JWK> jwkSelector;
+
+		private JwsHeader jwsHeader;
+
+		private static boolean isEcAlgorithm(SignatureAlgorithm algorithm) {
+			return Stream.of(SignatureAlgorithm.ES256, SignatureAlgorithm.ES384, SignatureAlgorithm.ES512)
+					.anyMatch(ecAlg -> ecAlg == algorithm);
+		}
+
+		private KeyPairJwtEncoderBuilder(KeyPair keyPair) {
+			this.keyPair = keyPair;
+		}
+
+		/**
+		 * Sets the JWS algorithm to use for signing. Must be compatible with the key type
+		 * (RSA or EC). If not set, a default algorithm will be chosen based on the key
+		 * type (e.g., RS256 for RSA, ES256 for EC).
+		 *
+		 * @param signatureAlgorithm the {@link SignatureAlgorithm} to use
+		 * @return this builder instance for method chaining
+		 */
+		public KeyPairJwtEncoderBuilder signatureAlgorithm(SignatureAlgorithm signatureAlgorithm) {
+			Assert.notNull(signatureAlgorithm, "signatureAlgorithm cannot be null");
+			this.jwsAlgorithm = JWSAlgorithm.parse(signatureAlgorithm.getName());
+			return this;
+		}
+
+		/**
+		 * Sets the key ID ({@code kid}) to be included in the JWK and potentially the JWS
+		 * header.
+		 *
+		 * @param keyId the key identifier
+		 * @return this builder instance for method chaining
+		 */
+		public KeyPairJwtEncoderBuilder keyId(String keyId) {
+			this.keyId = keyId;
+			return this;
+		}
+
+		/**
+		 * Configures the {@link Converter} used to select the JWK when multiple keys match
+		 * the header criteria. This is generally not needed for single-key setups but is
+		 * provided for consistency.
+		 *
+		 * @param jwkSelector the {@link Converter} to select a {@link JWK}
+		 * @return this builder instance for method chaining
+		 * @since // TODO: Update version
+		 */
+		public KeyPairJwtEncoderBuilder jwkSelector(Converter<List<JWK>, JWK> jwkSelector) {
+			Assert.notNull(jwkSelector, "jwkSelector cannot be null");
+			this.jwkSelector = jwkSelector;
+			return this;
+		}
+
+		/**
+		 * Builds the {@link NimbusJwtEncoder} instance.
+		 *
+		 * @return the configured {@link NimbusJwtEncoder}
+		 * @throws IllegalStateException if the key type is unsupported or the configured
+		 *                               JWS algorithm is not compatible with the key type.
+		 * @throws JwtEncodingException  if the key is invalid (e.g., EC key with unknown curve)
+		 */
+		public NimbusJwtEncoder build() {
+			this.keyId = this.keyId != null ? this.keyId : UUID.randomUUID().toString();
+			JWK jwk = buildJwk();
+			JWKSource<SecurityContext> jwkSource = new ImmutableJWKSet<>(new JWKSet(jwk));
+			NimbusJwtEncoder encoder = this.jwsHeader  == null ?
+					new NimbusJwtEncoder(jwkSource) :
+					new NimbusJwtEncoder(jwkSource, this.jwsHeader);
+			if (this.jwkSelector != null) {
+				encoder.setJwkSelector(this.jwkSelector);
+			} else {
+				encoder.setJwkSelector(jwkSet -> jwkSet.stream().findFirst().orElse(null));
+			}
+			return encoder;
+		}
+
+		private JWK buildJwk() {
+			if (this.keyPair.getPrivate() instanceof RSAPrivateKey) {
+				RSAKey rsaKey = buildRsaJwk();
+				this.jwsHeader = JwsHeader.with(SignatureAlgorithm.from(jwsAlgorithm.getName())) // Use Spring's enum
+						.keyId(rsaKey.getKeyID())
+						.build();
+				return rsaKey;
+			} else if (this.keyPair.getPrivate() instanceof ECPrivateKey) {
+				ECKey ecKey = buildEcJwk();
+				this.jwsHeader = JwsHeader.with(SignatureAlgorithm.from(jwsAlgorithm.getName())) // Use Spring's enum
+						.keyId(ecKey.getKeyID())
+						.build();
+				return ecKey;
+			} else {
+				throw new IllegalStateException("Unsupported key pair type: " + this.keyPair.getPrivate().getClass().getName());
+			}
+		}
+
+		private RSAKey buildRsaJwk() {
+			if (this.jwsAlgorithm == null) {
+				this.jwsAlgorithm = JWSAlgorithm.RS256;
+			}
+			Assert.state(JWSAlgorithm.Family.RSA.contains(this.jwsAlgorithm),
+					() -> "The algorithm '" + this.jwsAlgorithm + "' is not compatible with an RSAKey. "
+							+ "Please use one of the RS256, RS384, RS512, PS256, PS384, or PS512 algorithms.");
+
+			RSAKey.Builder builder = new RSAKey.Builder((java.security.interfaces.RSAPublicKey) this.keyPair.getPublic())
+					.privateKey(this.keyPair.getPrivate())
+					.keyUse(KeyUse.SIGNATURE)
+					.algorithm(this.jwsAlgorithm);
+
+			if (StringUtils.hasText(this.keyId)) {
+				builder.keyID(this.keyId);
+			}
+			return builder.build();
+		}
+
+		private com.nimbusds.jose.jwk.ECKey buildEcJwk() {
+			if (this.jwsAlgorithm == null) {
+				this.jwsAlgorithm = JWSAlgorithm.ES256;
+			}
+			Assert.state(JWSAlgorithm.Family.EC.contains(this.jwsAlgorithm),
+					() -> "The algorithm '" + this.jwsAlgorithm + "' is not compatible with an ECKey. "
+							+ "Please use one of the ES256, ES384, or ES512 algorithms.");
+
+			ECPublicKey publicKey = (ECPublicKey) this.keyPair.getPublic();
+			Curve curve = Curve.forECParameterSpec(publicKey.getParams());
+			if (curve == null) {
+				throw new JwtEncodingException("Unable to determine Curve for EC public key.");
+			}
+
+			com.nimbusds.jose.jwk.ECKey.Builder builder = new com.nimbusds.jose.jwk.ECKey.Builder(curve, publicKey)
+					.privateKey(this.keyPair.getPrivate())
+					.keyUse(KeyUse.SIGNATURE)
+					.algorithm(this.jwsAlgorithm);
+
+			if (StringUtils.hasText(this.keyId)) {
+				builder.keyID(this.keyId);
+			}
+
+			try {
+				return builder.build();
+			} catch (IllegalStateException ex) {
+				throw new JwtEncodingException("Failed to build ECKey: " + ex.getMessage(), ex);
+			}
+		}
+	}
+
 }