feat: add jwt decoder (#15)

* feat: add jwt decoder

* chore: remove jwt verification and expose the required info via context

* test: add test for JwtParser

* chore:address review comments

Author: GurtejSohi

grpc-context-utils/build.gradle.kts

@@ -12,15 +12,15 @@ tasks.test {
 dependencies {
   // grpc
   implementation("io.grpc:grpc-core:1.36.0")
-  constraints {
-    implementation("com.google.guava:guava:30.0-jre") {
-      because("https://snyk.io/vuln/SNYK-JAVA-COMGOOGLEGUAVA-1015415")
-    }
-  }
+
+  implementation("com.auth0:java-jwt:3.14.0")
+  implementation("com.auth0:jwks-rsa:0.17.0")
+  implementation("com.google.guava:guava:30.1-jre")
 
   // Logging
   implementation("org.slf4j:slf4j-api:1.7.30")
   // End Logging
 
   testImplementation("org.junit.jupiter:junit-jupiter:5.7.0")
+  testImplementation("org.mockito:mockito-core:3.8.0")
 }

grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/Jwt.java

@@ -0,0 +1,13 @@
+package org.hypertrace.core.grpcutils.context;
+
+import java.util.Optional;
+
+interface Jwt {
+  Optional<String> getUserId();
+
+  Optional<String> getName();
+
+  Optional<String> getPictureUrl();
+
+  Optional<String> getEmail();
+}

grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/JwtParser.java

@@ -0,0 +1,77 @@
+package org.hypertrace.core.grpcutils.context;
+
+import com.auth0.jwt.JWT;
+import com.auth0.jwt.interfaces.DecodedJWT;
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
+import java.util.Optional;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+class JwtParser {
+  private static final Logger LOG = LoggerFactory.getLogger(JwtParser.class);
+  private static final String BEARER_TOKEN_PREFIX = "Bearer ";
+
+  private final Cache<String, Optional<Jwt>> jwtCache =
+      CacheBuilder.newBuilder().maximumSize(1000).expireAfterAccess(1, TimeUnit.HOURS).build();
+
+  Optional<Jwt> fromAuthHeader(String authHeaderValue) {
+    if (authHeaderValue.startsWith(BEARER_TOKEN_PREFIX)) {
+      return this.fromJwt(authHeaderValue.substring(BEARER_TOKEN_PREFIX.length()));
+    }
+    return Optional.empty();
+  }
+
+  Optional<Jwt> fromJwt(String jwtValue) {
+    try {
+      return this.jwtCache.get(jwtValue, () -> this.decode(jwtValue));
+    } catch (ExecutionException e) {
+      LOG.warn("Exception loading jwt from cache", e);
+      return Optional.empty();
+    }
+  }
+
+  private Optional<Jwt> decode(String jwtString) {
+    try {
+      DecodedJWT jwt = JWT.decode(jwtString);
+      return Optional.of(new DefaultJwt(jwt));
+    } catch (Throwable t) {
+      LOG.warn("Failed to verify JWT", t);
+      return Optional.empty();
+    }
+  }
+
+  private static final class DefaultJwt implements Jwt {
+    private final DecodedJWT jwt;
+    private static final String SUBJECT_CLAIM = "sub";
+    private static final String NAME_CLAIM = "name";
+    private static final String PICTURE_CLAIM = "picture";
+    private static final String EMAIL_CLAIM = "email";
+
+    private DefaultJwt(DecodedJWT jwt) {
+      this.jwt = jwt;
+    }
+
+    @Override
+    public Optional<String> getUserId() {
+      return Optional.ofNullable(jwt.getClaim(SUBJECT_CLAIM).asString());
+    }
+
+    @Override
+    public Optional<String> getName() {
+      return Optional.ofNullable(jwt.getClaim(NAME_CLAIM).asString());
+    }
+
+    @Override
+    public Optional<String> getPictureUrl() {
+      return Optional.ofNullable(jwt.getClaim(PICTURE_CLAIM).asString());
+    }
+
+    @Override
+    public Optional<String> getEmail() {
+      return Optional.ofNullable(jwt.getClaim(EMAIL_CLAIM).asString());
+    }
+  }
+}

grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/RequestContext.java

@@ -21,12 +21,33 @@ public static RequestContext forTenantId(String tenantId) {
   }
 
   private final Map<String, String> headers = new HashMap<>();
+  private final JwtParser jwtParser = new JwtParser();
 
   /** Reads tenant id from this RequestContext based on the tenant id http header and returns it. */
   public Optional<String> getTenantId() {
     return get(RequestContextConstants.TENANT_ID_HEADER_KEY);
   }
 
+  public Optional<String> getUserId() {
+    return getJwt().flatMap(Jwt::getUserId);
+  }
+
+  public Optional<String> getName() {
+    return getJwt().flatMap(Jwt::getName);
+  }
+
+  public Optional<String> getPictureUrl() {
+    return getJwt().flatMap(Jwt::getPictureUrl);
+  }
+
+  public Optional<String> getEmail() {
+    return getJwt().flatMap(Jwt::getEmail);
+  }
+
+  private Optional<Jwt> getJwt() {
+    return get(RequestContextConstants.AUTHORIZATION_HEADER).flatMap(jwtParser::fromAuthHeader);
+  }
+
   /** Method to read all GRPC request headers from this RequestContext. */
   public Map<String, String> getRequestHeaders() {
     return getAll();

grpc-context-utils/src/test/java/org/hypertrace/core/grpcutils/context/JwtParserTest.java

@@ -0,0 +1,57 @@
+package org.hypertrace.core.grpcutils.context;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import java.util.Optional;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentMatchers;
+
+class JwtParserTest {
+  private final String testJwt =
+      "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJPbmxpbmUgSldUIEJ1aWxkZXIiLCJpYXQiOjE2MjEzNjM1OTcsImV4cCI6MTY1Mjg5OTU5NywiYXVkIjoid3d3LmV4YW1wbGUuY29tIiwic3ViIjoianJvY2tldEBleGFtcGxlLmNvbSIsIkdpdmVuTmFtZSI6IkpvaG5ueSIsIlN1cm5hbWUiOiJSb2NrZXQiLCJuYW1lIjoiSm9obm55IFJvY2tldCIsImVtYWlsIjoianJvY2tldEBleGFtcGxlLmNvbSIsInBpY3R1cmUiOiJ3d3cuZXhhbXBsZS5jb20iLCJSb2xlIjpbIk1hbmFnZXIiLCJQcm9qZWN0IEFkbWluaXN0cmF0b3IiXX0.aesOuNIamZkTMR30CBt0J9NMZZt9iLRETa5ayN_EcVs";
+  private final String testJwtUserId = "[email protected]";
+  private final String testJwtName = "Johnny Rocket";
+  private final String testJwtPictureUrl = "www.example.com";
+  private final String testJwtEmail = "[email protected]";
+
+  @Test
+  void testGoodJwtParse() {
+    JwtParser parser = new JwtParser();
+    Optional<Jwt> jwt = parser.fromJwt(testJwt);
+    assertEquals(Optional.of(testJwtUserId), jwt.flatMap(Jwt::getUserId));
+    assertEquals(Optional.of(testJwtName), jwt.flatMap(Jwt::getName));
+    assertEquals(Optional.of(testJwtPictureUrl), jwt.flatMap(Jwt::getPictureUrl));
+    assertEquals(Optional.of(testJwtEmail), jwt.flatMap(Jwt::getEmail));
+  }
+
+  @Test
+  void testBadJwtParse() {
+    JwtParser parser = new JwtParser();
+    Optional<Jwt> jwt = parser.fromJwt("fake jwt");
+    assertTrue(jwt.isEmpty());
+  }
+
+  @Test
+  void testExtractBearerTokenPassesThrough() {
+    JwtParser parser = mock(JwtParser.class);
+    when(parser.fromAuthHeader(anyString())).thenCallRealMethod();
+
+    parser.fromAuthHeader("Bearer foobar");
+    verify(parser).fromJwt("foobar");
+  }
+
+  @Test
+  void testExtractBearerTokenReturnsEmptyOnMalformed() {
+    JwtParser parser = mock(JwtParser.class);
+    when(parser.fromAuthHeader(anyString())).thenCallRealMethod();
+
+    assertEquals(Optional.empty(), parser.fromAuthHeader("Bad header"));
+    verify(parser, times(0)).fromJwt(ArgumentMatchers.any());
+  }
+}