[improve][cli] Supports creating tokens with additional headers (#22690)

apache · May 20, 2024 · 9c60134 · 9c60134
1 parent 6372b9c
commit 9c60134
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 2 deletions.
diff --git a/...er-common/src/main/java/org/apache/pulsar/broker/authentication/utils/AuthTokenUtils.java b/...er-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<Date> expiryTime) {
+    public static String createToken(Key signingKey, String subject, Optional<Date> expiryTime,
+                                     Optional<Map<String, Object>> 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<Date> 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
@@ -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<Integer> {
                 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<String, Object> 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
@@ -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.
      *