Better JWT token processing (#733)

Signed-off-by: Dmitrii Tikhomirov <[email protected]>

Author: treblereel

impl/http/src/main/java/io/serverlessworkflow/impl/executors/http/OAuth2AuthProvider.java

@@ -52,10 +52,14 @@ public Builder build(
   @Override
   public void preRequest(
       Invocation.Builder builder, WorkflowContext workflow, TaskContext task, WorkflowModel model) {
-    JWT token = requestBuilder.build(workflow, task, model).validateAndGet();
-    String tokenType = (String) token.getClaim("typ");
+    JWT jwt = requestBuilder.build(workflow, task, model).validateAndGet();
+    String type =
+        jwt.claim("typ", String.class)
+            .map(String::trim)
+            .filter(t -> !t.isEmpty())
+            .orElseThrow(() -> new IllegalStateException("Token type is not present"));
+
     builder.header(
-        AuthProviderFactory.AUTH_HEADER_NAME,
-        String.format(BEARER_TOKEN, tokenType, token.getToken()));
+        AuthProviderFactory.AUTH_HEADER_NAME, String.format(BEARER_TOKEN, type, jwt.token()));
   }
 }

impl/http/src/main/java/io/serverlessworkflow/impl/executors/http/oauth/AccessTokenProvider.java

@@ -47,10 +47,13 @@ public JWT validateAndGet() {
     Map<String, Object> token = tokenResponseHandler.apply(invocation, context);
     JWT jwt = jwtConverter.fromToken((String) token.get("access_token"));
     if (!(issuers == null || issuers.isEmpty())) {
-      String tokenIssuer = (String) jwt.getClaim("iss");
-      if (tokenIssuer == null || tokenIssuer.isEmpty() || !issuers.contains(tokenIssuer)) {
-        throw new IllegalStateException("Token issuer is not valid: " + tokenIssuer);
-      }
+      jwt.issuer()
+          .ifPresent(
+              issuer -> {
+                if (!issuers.contains(issuer)) {
+                  throw new IllegalStateException("Token issuer is not valid: " + issuer);
+                }
+              });
     }
     return jwt;
   }

impl/jwt-impl/pom.xml

@@ -21,6 +21,26 @@
             <groupId>com.fasterxml.jackson.core</groupId>
             <artifactId>jackson-databind</artifactId>
         </dependency>
+        <dependency>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter-api</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter-engine</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter-params</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.assertj</groupId>
+            <artifactId>assertj-core</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
 </project>

impl/jwt-impl/src/main/java/io/serverlessworkflow/impl/http/jwt/JacksonJWTConverter.java

@@ -29,14 +29,24 @@ public class JacksonJWTConverter implements JWTConverter {
 
   @Override
   public JWT fromToken(String token) throws IllegalArgumentException {
+    if (token == null || token.isBlank()) {
+      throw new IllegalArgumentException("JWT token must not be null or blank");
+    }
+
     String[] parts = token.split("\\.");
     if (parts.length < 2) {
       throw new IllegalArgumentException("Invalid JWT token format");
     }
     try {
+      String headerJson =
+          new String(Base64.getUrlDecoder().decode(parts[0]), StandardCharsets.UTF_8);
       String payloadJson =
           new String(Base64.getUrlDecoder().decode(parts[1]), StandardCharsets.UTF_8);
-      return new JacksonJWTImpl(token, MAPPER.readValue(payloadJson, Map.class));
+
+      Map<String, Object> header = MAPPER.readValue(headerJson, Map.class);
+      Map<String, Object> claims = MAPPER.readValue(payloadJson, Map.class);
+
+      return new JacksonJWTImpl(token, header, claims);
     } catch (JsonProcessingException e) {
       throw new IllegalArgumentException("Failed to parse JWT token payload: " + e.getMessage(), e);
     }

impl/jwt-impl/src/main/java/io/serverlessworkflow/impl/http/jwt/JacksonJWTImpl.java

@@ -16,28 +16,132 @@
 package io.serverlessworkflow.impl.http.jwt;
 
 import io.serverlessworkflow.http.jwt.JWT;
+import java.time.Instant;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
 import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
 
 public class JacksonJWTImpl implements JWT {
 
+  private final Map<String, Object> header;
   private final Map<String, Object> claims;
+
   private final String token;
 
-  JacksonJWTImpl(String token, Map<String, Object> claims) {
+  JacksonJWTImpl(String token, Map<String, Object> header, Map<String, Object> claims) {
+    Objects.requireNonNull(token, "token");
+    Objects.requireNonNull(header, "header");
+    Objects.requireNonNull(claims, "claims");
     this.token = token;
-    this.claims = claims;
+    this.header = Collections.unmodifiableMap(header);
+    this.claims = Collections.unmodifiableMap(claims);
   }
 
   @Override
-  public String getToken() {
+  public String token() {
     return token;
   }
 
   @Override
-  public Object getClaim(String name) {
-    if (claims == null || claims.isEmpty()) {
+  public Map<String, Object> header() {
+    return header;
+  }
+
+  @Override
+  public Map<String, Object> claims() {
+    return claims;
+  }
+
+  @Override
+  public <T> Optional<T> claim(String name, Class<T> type) {
+    if (name == null || type == null) return Optional.empty();
+    Object value = claims.get(name);
+    return type.isInstance(value) ? Optional.of(type.cast(value)) : Optional.empty();
+  }
+
+  @Override
+  public Optional<String> issuer() {
+    return Optional.ofNullable(asString(claims.get("iss")));
+  }
+
+  @Override
+  public Optional<String> subject() {
+    return Optional.ofNullable(asString(claims.get("sub")));
+  }
+
+  @Override
+  public List<String> audience() {
+    Object aud = claims.get("aud");
+    if (aud == null) {
+      return List.of();
+    } else if (aud instanceof String asString) {
+      return List.of(asString);
+    } else if (aud instanceof String[] asStringArray) {
+      return Arrays.asList(asStringArray);
+    }
+    return ((Collection<?>) aud).stream().map(String::valueOf).toList();
+  }
+
+  @Override
+  public Optional<Instant> issuedAt() {
+    return Optional.ofNullable(toInstant(claims.get("iat")));
+  }
+
+  @Override
+  public Optional<Instant> expiresAt() {
+    return Optional.ofNullable(toInstant(claims.get("exp")));
+  }
+
+  @Override
+  public Optional<String> type() {
+    if (header.containsKey("typ")) {
+      return Optional.of(asString(header.get("typ")));
+    }
+    return Optional.ofNullable(asString(claims.get("typ")));
+  }
+
+  private static Instant toInstant(Object v) {
+    if (v == null) {
       return null;
     }
-    return claims.get(name);
+    if (v instanceof Instant i) {
+      return i;
+    }
+    if (v instanceof Number n) {
+      return Instant.ofEpochSecond(n.longValue());
+    }
+    if (v instanceof String s) {
+      try {
+        long sec = Long.parseLong(s.trim());
+        return Instant.ofEpochSecond(sec);
+      } catch (NumberFormatException ignored) {
+        try {
+          return Instant.parse(s.trim());
+        } catch (Exception e) {
+          throw new IllegalArgumentException("Cannot parse time claim: " + s, e);
+        }
+      }
+    }
+    throw new IllegalArgumentException("Unsupported time claim type: " + v.getClass());
+  }
+
+  private static String asString(Object v) {
+    return (v == null) ? null : String.valueOf(v).trim();
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) return true;
+    if (!(o instanceof JacksonJWTImpl that)) return false;
+    return Objects.equals(token, that.token);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(token);
   }
 }

impl/jwt-impl/src/test/java/io/serverlessworkflow/impl/http/jwt/JacksonJWTImplTest.java

@@ -0,0 +1,131 @@
+/*
+ * Copyright 2020-Present The Serverless Workflow Specification Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.serverlessworkflow.impl.http.jwt;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import io.serverlessworkflow.http.jwt.JWT;
+import java.time.Instant;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import org.junit.jupiter.api.Test;
+
+public class JacksonJWTImplTest {
+
+  private static String JWT_TOKEN =
+      "eyJhbGciOiJSUzI1NiIsInR5cCIgOiAiSldUIiwia2lkIiA6ICJTY1c4RGI5RThtSnlvTUNqZHVtek5VR21FcG9MYURDb19ralZkVHJ2NDdVIn0.eyJleHAiOjE3NTYxNDgyNTcsImlhdCI6MTc1NjE0Nzk1NywianRpIjoiOWU4YzZjMWItZDBmZS00NGNhLThlOTgtNzNkZTY4OTdjYzE4IiwiaXNzIjoiaHR0cDovL2xvY2FsaG9zdDo4MDg4L3JlYWxtcy90ZXN0LXJlYWxtIiwiYXVkIjoiYWNjb3VudCIsInN1YiI6ImExZDdlMWVlLTVmZjMtNDc3Yi1hMmRhLTNhNDZkYThlZmRkMSIsInR5cCI6IkJlYXJlciIsImF6cCI6InNlcnZlcmxlc3Mtd29ya2Zsb3ciLCJzZXNzaW9uX3N0YXRlIjoiNDdiYzIxZDEtYTMyYi00OGJlLWI2OWQtZDM5MjE5MjViZTlmIiwiYWNyIjoiMSIsImFsbG93ZWQtb3JpZ2lucyI6WyJodHRwOi8vbG9jYWxob3N0OjgwODAiXSwicmVhbG1fYWNjZXNzIjp7InJvbGVzIjpbImRlZmF1bHQtcm9sZXMtdGVzdC1yZWFsbSIsIm9mZmxpbmVfYWNjZXNzIiwidW1hX2F1dGhvcml6YXRpb24iXX0sInJlc291cmNlX2FjY2VzcyI6eyJhY2NvdW50Ijp7InJvbGVzIjpbIm1hbmFnZS1hY2NvdW50IiwibWFuYWdlLWFjY291bnQtbGlua3MiLCJ2aWV3LXByb2ZpbGUiXX19LCJzY29wZSI6InByb2ZpbGUgZW1haWwiLCJzaWQiOiI0N2JjMjFkMS1hMzJiLTQ4YmUtYjY5ZC1kMzkyMTkyNWJlOWYiLCJlbWFpbF92ZXJpZmllZCI6ZmFsc2UsIm5hbWUiOiJEbWl0cmlpIFRpa2hvbWlyb3YiLCJwcmVmZXJyZWRfdXNlcm5hbWUiOiJzZXJ2ZXJsZXNzLXdvcmtmbG93LXRlc3QiLCJnaXZlbl9uYW1lIjoiRG1pdHJpaSIsImZhbWlseV9uYW1lIjoiVGlraG9taXJvdiIsImVtYWlsIjoiYW55QGV4YW1wbGUub3JnIn0.bY9DqVt5jPcfsY5-tL4JbXXptnyFBBLLKuViBH3xQBXpHSmtXKK78BjBdUOcCAg6SMatdQ4XK13RJJNPpjPDCkuBeLDMwsdCrzKcfaUWY7JCMfam4gIiC989m_S8y8jtJovDst8sjIkn95f5izuuiAxRYw69_IcY__cfBw7k0OhL7Y_YXCcwwz7l2yrbplmpLiakTTuzhbiCER5EPohguy9BkYG_u2RDUZg0Rvy3EzbyVhTQQyfKRFjoAxTa1Se484n2lXqgMn1JHTZLwrgXAjcMDVaktktLb_cn276ygAeuqPj7dOsQGEoLR-8enRz1eZKBPgO70LNqGkkkyHiyOA";
+
+  private static JacksonJWTConverter converter = new JacksonJWTConverter();
+
+  @Test
+  void tokenRoundTrip() {
+    JWT jwt = converter.fromToken(JWT_TOKEN);
+    assertEquals(JWT_TOKEN, jwt.token());
+  }
+
+  @Test
+  void headerAndType() {
+    JWT jwt = converter.fromToken(JWT_TOKEN);
+
+    Map<String, Object> h = jwt.header();
+    assertEquals("RS256", h.get("alg"));
+    assertTrue(h.containsKey("kid"));
+    assertEquals("JWT", h.get("typ"));
+
+    assertEquals(Optional.of("JWT"), jwt.type(), "type() must be a header.typ");
+  }
+
+  @Test
+  void payloadTypIsArbitraryClaim() {
+    JWT jwt = converter.fromToken(JWT_TOKEN);
+
+    assertEquals(Optional.of("Bearer"), jwt.claim("typ", String.class));
+    assertEquals("JWT", jwt.header().get("typ"));
+  }
+
+  @Test
+  void standardClaims() {
+    JWT jwt = converter.fromToken(JWT_TOKEN);
+
+    assertEquals(Optional.of("http://localhost:8088/realms/test-realm"), jwt.issuer());
+    assertEquals(Optional.of("a1d7e1ee-5ff3-477b-a2da-3a46da8efdd1"), jwt.subject());
+    assertEquals(List.of("account"), jwt.audience());
+  }
+
+  @Test
+  void timeClaims() {
+    JWT jwt = converter.fromToken(JWT_TOKEN);
+
+    assertEquals(Instant.ofEpochSecond(1756148257L), jwt.expiresAt().orElseThrow());
+    assertEquals(Instant.ofEpochSecond(1756147957L), jwt.issuedAt().orElseThrow());
+  }
+
+  @Test
+  void typedClaimAccess() {
+    JWT jwt = converter.fromToken(JWT_TOKEN);
+
+    assertEquals(Optional.of("[email protected]"), jwt.claim("email", String.class));
+    assertEquals(Optional.of(false), jwt.claim("email_verified", Boolean.class));
+    assertTrue(jwt.claim("nonexistent", String.class).isEmpty());
+    assertTrue(jwt.claim("email_verified", String.class).isEmpty());
+  }
+
+  @Test
+  void mapsAreUnmodifiable() {
+    JWT jwt = converter.fromToken(JWT_TOKEN);
+
+    assertThrows(UnsupportedOperationException.class, () -> jwt.claims().put("x", 1));
+    assertThrows(UnsupportedOperationException.class, () -> jwt.header().put("x", 1));
+  }
+
+  @Test
+  void nullToken() {
+    assertThrows(IllegalArgumentException.class, () -> converter.fromToken(null));
+  }
+
+  @Test
+  void blankToken() {
+    assertThrows(IllegalArgumentException.class, () -> converter.fromToken("   "));
+  }
+
+  @Test
+  void wrongPartsCount() {
+    assertThrows(IllegalArgumentException.class, () -> converter.fromToken("a.b"));
+    assertThrows(IllegalArgumentException.class, () -> converter.fromToken("a.b.c.d"));
+  }
+
+  @Test
+  void badBase64Header() {
+    String bad = "###." + JWT_TOKEN.split("\\.")[1] + "." + JWT_TOKEN.split("\\.")[2];
+    assertThrows(IllegalArgumentException.class, () -> converter.fromToken(bad));
+  }
+
+  @Test
+  void badBase64Payload() {
+    String[] p = JWT_TOKEN.split("\\.");
+    String bad = p[0] + ".###." + p[2];
+    assertThrows(IllegalArgumentException.class, () -> converter.fromToken(bad));
+  }
+
+  @Test
+  void badJson() {
+    String[] p = JWT_TOKEN.split("\\.");
+    String notJsonB64url = "bm90LWpzb24";
+    String bad = p[0] + "." + notJsonB64url + "." + p[2];
+    assertThrows(IllegalArgumentException.class, () -> converter.fromToken(bad));
+  }
+}

impl/jwt/src/main/java/io/serverlessworkflow/http/jwt/JWT.java

@@ -15,9 +15,30 @@
  */
 package io.serverlessworkflow.http.jwt;
 
+import java.time.Instant;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
 public interface JWT {
 
-  String getToken();
+  String token();
+
+  List<String> audience();
+
+  Map<String, Object> claims();
+
+  <T> Optional<T> claim(String name, Class<T> type);
+
+  Optional<Instant> expiresAt();
+
+  Map<String, Object> header();
+
+  Optional<Instant> issuedAt();
+
+  Optional<String> issuer();
+
+  Optional<String> subject();
 
-  Object getClaim(String name);
+  Optional<String> type();
 }