diff --git a/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/authentication/utils/AuthTokenUtils.java b/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/authentication/utils/AuthTokenUtils.java index cb917a9e0bf00..e29b106f4f17e 100644 --- a/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/authentication/utils/AuthTokenUtils.java +++ b/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/authentication/utils/AuthTokenUtils.java @@ -35,6 +35,7 @@ import java.security.spec.PKCS8EncodedKeySpec; import java.security.spec.X509EncodedKeySpec; import java.util.Date; +import java.util.Map; import java.util.Optional; import javax.crypto.SecretKey; import lombok.experimental.UtilityClass; @@ -89,16 +90,22 @@ public static String encodeKeyBase64(Key key) { return Encoders.BASE64.encode(key.getEncoded()); } - public static String createToken(Key signingKey, String subject, Optional expiryTime) { + public static String createToken(Key signingKey, String subject, Optional expiryTime, + Optional> headers) { JwtBuilder builder = Jwts.builder() .setSubject(subject) .signWith(signingKey); expiryTime.ifPresent(builder::setExpiration); + headers.ifPresent(builder::setHeaderParams); return builder.compact(); } + public static String createToken(Key signingKey, String subject, Optional expiryTime) { + return createToken(signingKey, subject, expiryTime, Optional.empty()); + } + public static byte[] readKeyFromUrl(String keyConfUrl) throws IOException { if (keyConfUrl.startsWith("data:") || keyConfUrl.startsWith("file:")) { try { diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/utils/auth/tokens/TokensCliUtils.java b/pulsar-broker/src/main/java/org/apache/pulsar/utils/auth/tokens/TokensCliUtils.java index 4ae28b2c0bdb8..78268a6295c28 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/utils/auth/tokens/TokensCliUtils.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/utils/auth/tokens/TokensCliUtils.java @@ -34,6 +34,7 @@ import java.security.Key; import java.security.KeyPair; import java.util.Date; +import java.util.Map; import java.util.Optional; import java.util.concurrent.Callable; import javax.crypto.SecretKey; @@ -140,6 +141,11 @@ public static class CommandCreateToken implements Callable { description = "Pass the private key for signing the token. This can either be: data:, file:, etc..") private String privateKey; + @Option(names = {"-hs", + "--headers"}, + description = "Additional headers to token. Format: --headers key1=value1") + private Map headers; + @Override public Integer call() throws Exception { if (secretKey == null && privateKey == null) { @@ -166,7 +172,7 @@ public Integer call() throws Exception { ? Optional.empty() : Optional.of(new Date(System.currentTimeMillis() + expiryTime)); - String token = AuthTokenUtils.createToken(signingKey, subject, optExpiryTime); + String token = AuthTokenUtils.createToken(signingKey, subject, optExpiryTime, Optional.ofNullable(headers)); System.out.println(token); return 0; diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/utils/auth/tokens/TokensCliUtilsTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/utils/auth/tokens/TokensCliUtilsTest.java index d5dc259438ea8..65c5d9981bfd2 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/utils/auth/tokens/TokensCliUtilsTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/utils/auth/tokens/TokensCliUtilsTest.java @@ -18,7 +18,12 @@ */ package org.apache.pulsar.utils.auth.tokens; +import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertTrue; +import io.jsonwebtoken.JwsHeader; +import io.jsonwebtoken.Jwt; +import io.jsonwebtoken.Jwts; +import io.jsonwebtoken.io.Decoders; import java.io.ByteArrayOutputStream; import java.io.PrintStream; import java.lang.reflect.Field; @@ -31,6 +36,45 @@ */ public class TokensCliUtilsTest { + @Test + public void testCreateToken() { + PrintStream oldStream = System.out; + try { + ByteArrayOutputStream baoStream = new ByteArrayOutputStream(); + System.setOut(new PrintStream(baoStream)); + + new TokensCliUtils().execute(new String[]{"create-secret-key", "--base64"}); + String secretKey = baoStream.toString(); + + baoStream.reset(); + + String[] command = {"create", "--secret-key", + "data:;base64," + secretKey, + "--subject", "test", + "--headers", "kid=test", + "--headers", "my-k=my-v" + }; + + new TokensCliUtils().execute(command); + String token = baoStream.toString(); + + Jwt jwt = Jwts.parserBuilder() + .setSigningKey(Decoders.BASE64.decode(secretKey)) + .build() + .parseClaimsJws(token); + + JwsHeader header = (JwsHeader) jwt.getHeader(); + String keyId = header.getKeyId(); + assertEquals(keyId, "test"); + assertEquals(header.get("my-k"), "my-v"); + + } catch (Exception e) { + throw new RuntimeException(e); + } finally { + System.setOut(oldStream); + } + } + /** * Test tokens generate docs. *