ENG-9494 expose roles in request context (#17)

d-trace · web-flow · commit a86db354f2b8 · 2021-06-14T14:10:47.000-07:00
expose roles in request context
diff --git a/grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/Jwt.java b/grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/Jwt.java
@@ -1,5 +1,6 @@
 package org.hypertrace.core.grpcutils.context;
 
+import java.util.List;
 import java.util.Optional;
 
 interface Jwt {
@@ -10,4 +11,6 @@ interface Jwt {
   Optional<String> getPictureUrl();
 
   Optional<String> getEmail();
+
+  List<String> getRoles(String rolesClaim);
 }
diff --git a/grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/JwtParser.java b/grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/JwtParser.java
@@ -4,6 +4,9 @@
 import com.auth0.jwt.interfaces.DecodedJWT;
 import com.google.common.cache.Cache;
 import com.google.common.cache.CacheBuilder;
+
+import java.util.Collections;
+import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
@@ -73,5 +76,14 @@ public Optional<String> getPictureUrl() {
     public Optional<String> getEmail() {
       return Optional.ofNullable(jwt.getClaim(EMAIL_CLAIM).asString());
     }
+
+    @Override
+    public List<String> getRoles(String rolesClaim) {
+      List<String> roles = jwt.getClaim(rolesClaim).asList(String.class);
+      if (roles == null || roles.isEmpty()) {
+        return Collections.emptyList();
+      }
+      return roles;
+    }
   }
 }
diff --git a/grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/RequestContext.java b/grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/RequestContext.java
@@ -1,7 +1,10 @@
 package org.hypertrace.core.grpcutils.context;
 
 import io.grpc.Context;
+
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.Callable;
@@ -44,6 +47,10 @@ public Optional<String> getEmail() {
     return getJwt().flatMap(Jwt::getEmail);
   }
 
+  public List<String> getRoles(String rolesClaim) {
+    return getJwt().map(jwt -> jwt.getRoles(rolesClaim)).orElse(Collections.emptyList());
+  }
+
   private Optional<Jwt> getJwt() {
     return get(RequestContextConstants.AUTHORIZATION_HEADER).flatMap(jwtParser::fromAuthHeader);
   }
diff --git a/grpc-context-utils/src/test/java/org/hypertrace/core/grpcutils/context/JwtParserTest.java b/grpc-context-utils/src/test/java/org/hypertrace/core/grpcutils/context/JwtParserTest.java
@@ -8,17 +8,24 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import java.util.Collections;
+import java.util.List;
 import java.util.Optional;
+
+import com.google.common.collect.ImmutableList;
 import org.junit.jupiter.api.Test;
 import org.mockito.ArgumentMatchers;
 
 class JwtParserTest {
-  private final String testJwt =
-      "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJPbmxpbmUgSldUIEJ1aWxkZXIiLCJpYXQiOjE2MjEzNjM1OTcsImV4cCI6MTY1Mjg5OTU5NywiYXVkIjoid3d3LmV4YW1wbGUuY29tIiwic3ViIjoianJvY2tldEBleGFtcGxlLmNvbSIsIkdpdmVuTmFtZSI6IkpvaG5ueSIsIlN1cm5hbWUiOiJSb2NrZXQiLCJuYW1lIjoiSm9obm55IFJvY2tldCIsImVtYWlsIjoianJvY2tldEBleGFtcGxlLmNvbSIsInBpY3R1cmUiOiJ3d3cuZXhhbXBsZS5jb20iLCJSb2xlIjpbIk1hbmFnZXIiLCJQcm9qZWN0IEFkbWluaXN0cmF0b3IiXX0.aesOuNIamZkTMR30CBt0J9NMZZt9iLRETa5ayN_EcVs";
+  private final String testJwt = "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJPbmxpbmUgSldUIEJ1aWxkZXIiLCJpYXQiOjE2MjEzNjM1OTcsImV4cCI6MTY1Mjg5OTU5NywiYXVkIjoid3d3LmV4YW1wbGUuY29tIiwic3ViIjoianJvY2tldEBleGFtcGxlLmNvbSIsIkdpdmVuTmFtZSI6IkpvaG5ueSIsIlN1cm5hbWUiOiJSb2NrZXQiLCJuYW1lIjoiSm9obm55IFJvY2tldCIsImVtYWlsIjoianJvY2tldEBleGFtcGxlLmNvbSIsInBpY3R1cmUiOiJ3d3cuZXhhbXBsZS5jb20iLCJyb2xlcyI6WyJzdXBlcl91c2VyIiwidXNlciIsImJpbGxpbmdfYWRtaW4iXX0.lEDjPPCjr-Epv6pNslq-HK9vmxfstp1sY85GstlbU1I";
+  private final String emptyRolesJwt = "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJPbmxpbmUgSldUIEJ1aWxkZXIiLCJpYXQiOjE2MjEzNjM1OTcsImV4cCI6MTY1Mjg5OTU5NywiYXVkIjoid3d3LmV4YW1wbGUuY29tIiwic3ViIjoianJvY2tldEBleGFtcGxlLmNvbSIsIkdpdmVuTmFtZSI6IkpvaG5ueSIsIlN1cm5hbWUiOiJSb2NrZXQiLCJuYW1lIjoiSm9obm55IFJvY2tldCIsImVtYWlsIjoianJvY2tldEBleGFtcGxlLmNvbSIsInBpY3R1cmUiOiJ3d3cuZXhhbXBsZS5jb20iLCJodHRwczovL3RyYWNlYWJsZS5haS9yb2xlcyI6W119.sFUMZNyypj379xy5P4kqTbBXBOR5XvX2nhpKx6YiiwU";
+  private final String noRolesJwt = "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJPbmxpbmUgSldUIEJ1aWxkZXIiLCJpYXQiOjE2MjEzNjM1OTcsImV4cCI6MTY1Mjg5OTU5NywiYXVkIjoid3d3LmV4YW1wbGUuY29tIiwic3ViIjoianJvY2tldEBleGFtcGxlLmNvbSIsIkdpdmVuTmFtZSI6IkpvaG5ueSIsIlN1cm5hbWUiOiJSb2NrZXQiLCJuYW1lIjoiSm9obm55IFJvY2tldCIsImVtYWlsIjoianJvY2tldEBleGFtcGxlLmNvbSIsInBpY3R1cmUiOiJ3d3cuZXhhbXBsZS5jb20ifQ.Ui1Z2RhiVe3tq6uJPgcyjsfDBdeOeINs_gXEHC6cdpU";
   private final String testJwtUserId = "jrocket@example.com";
   private final String testJwtName = "Johnny Rocket";
   private final String testJwtPictureUrl = "www.example.com";
   private final String testJwtEmail = "jrocket@example.com";
+  private final String testRolesClaim = "roles";
+  private final List<String> testRoles = ImmutableList.of("super_user", "user", "billing_admin");
 
   @Test
   void testGoodJwtParse() {
@@ -54,4 +61,25 @@ void testExtractBearerTokenReturnsEmptyOnMalformed() {
     assertEquals(Optional.empty(), parser.fromAuthHeader("Bad header"));
     verify(parser, times(0)).fromJwt(ArgumentMatchers.any());
   }
+
+  @Test
+  void testRolesCanBeParsedFromToken() {
+    JwtParser parser = new JwtParser();
+    Optional<Jwt> jwt = parser.fromJwt(testJwt);
+    assertEquals(Optional.of(testRoles), jwt.map(j -> j.getRoles(testRolesClaim)));
+  }
+
+  @Test
+  void testRolesAreEmptyIfRolesArrayIsEmptyInJwt() {
+    JwtParser parser = new JwtParser();
+    Optional<Jwt> jwt = parser.fromJwt(emptyRolesJwt);
+    assertEquals(Optional.of(Collections.emptyList()), jwt.map(j -> j.getRoles(testRolesClaim)));
+  }
+
+  @Test
+  void testRolesAreEmptyIfRolesIfNoRolesClaimInToken() {
+    JwtParser parser = new JwtParser();
+    Optional<Jwt> jwt = parser.fromJwt(noRolesJwt);
+    assertEquals(Optional.of(Collections.emptyList()), jwt.map(j -> j.getRoles(testRolesClaim)));
+  }
 }
diff --git a/grpc-context-utils/src/test/java/org/hypertrace/core/grpcutils/context/RequestContextTest.java b/grpc-context-utils/src/test/java/org/hypertrace/core/grpcutils/context/RequestContextTest.java
@@ -1,36 +1,40 @@
 package org.hypertrace.core.grpcutils.context;
 
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import org.junit.jupiter.api.Assertions;
+
+import com.google.common.collect.ImmutableList;
 import org.junit.jupiter.api.Test;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
 /** Unit tests for {@link RequestContext} and utility methods in it. */
 public class RequestContextTest {
   private static final String TENANT_ID = "example-tenant-id";
   private static final String TEST_AUTH_HEADER = "Bearer sample-auth-header";
 
   @Test
-  public void testTenantId() {
+  void testTenantId() {
     RequestContext requestContext = new RequestContext();
     requestContext.add(RequestContextConstants.TENANT_ID_HEADER_KEY, TENANT_ID);
     Optional<String> tenantId = requestContext.getTenantId();
-    Assertions.assertEquals(Optional.of(TENANT_ID), tenantId);
+    assertEquals(Optional.of(TENANT_ID), tenantId);
 
     requestContext = new RequestContext();
     tenantId = requestContext.getTenantId();
-    Assertions.assertEquals(Optional.empty(), tenantId);
+    assertEquals(Optional.empty(), tenantId);
   }
 
   @Test
-  public void testGetRequestHeaders() {
+  void testGetRequestHeaders() {
     RequestContext requestContext = new RequestContext();
     requestContext.add(RequestContextConstants.AUTHORIZATION_HEADER, TEST_AUTH_HEADER);
     requestContext.add("x-some-tenant-header", "v1");
 
     Map<String, String> requestHeaders = requestContext.getRequestHeaders();
 
-    Assertions.assertEquals(
+    assertEquals(
         Map.of(
             RequestContextConstants.AUTHORIZATION_HEADER,
             TEST_AUTH_HEADER,
@@ -40,12 +44,27 @@ public void testGetRequestHeaders() {
   }
 
   @Test
-  public void testCreateForTenantId() {
+  void testCreateForTenantId() {
     RequestContext requestContext = RequestContext.forTenantId(TENANT_ID);
-    Assertions.assertEquals(Optional.of(TENANT_ID), requestContext.getTenantId());
-    Assertions.assertEquals(
+    assertEquals(Optional.of(TENANT_ID), requestContext.getTenantId());
+    assertEquals(
         Optional.of(TENANT_ID), requestContext.get(RequestContextConstants.TENANT_ID_HEADER_KEY));
-    Assertions.assertEquals(
+    assertEquals(
         Map.of(RequestContextConstants.TENANT_ID_HEADER_KEY, TENANT_ID), requestContext.getAll());
   }
+
+  @Test
+  void testRolesArePropagatedInRequestContext() {
+    List<String> expectedRoles = ImmutableList.of("super_user", "user", "billing_admin");
+    String jwt = "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJPbmxpbmUgSldUIEJ1aWxkZXIiLCJpYXQiOjE2MjEzNjM1OTcsIm" +
+        "V4cCI6MTY1Mjg5OTU5NywiYXVkIjoid3d3LmV4YW1wbGUuY29tIiwic3ViIjoianJvY2tldEBleGFtcGxlLmNvbSIsIkdpdmVuTmFtZSI6Ik" +
+        "pvaG5ueSIsIlN1cm5hbWUiOiJSb2NrZXQiLCJuYW1lIjoiSm9obm55IFJvY2tldCIsImVtYWlsIjoianJvY2tldEBleGFtcGxlLmNvbSIsIn" +
+        "BpY3R1cmUiOiJ3d3cuZXhhbXBsZS5jb20iLCJyb2xlcyI6WyJzdXBlcl91c2VyIiwidXNlciIsImJpbGxpbmdfYWRtaW4iXX0.lEDjPPCjr-" +
+        "Epv6pNslq-HK9vmxfstp1sY85GstlbU1I";
+
+    RequestContext requestContext = new RequestContext();
+    requestContext.add("authorization", "Bearer " + jwt);
+    List<String> actualRoles = requestContext.getRoles("roles");
+    assertEquals(expectedRoles, actualRoles);
+  }
 }
