Extract identity building to bearer authenticators

trinodb · Jan 24, 2022 · 3d64dfb · 3d64dfb
1 parent b1e66fc
commit 3d64dfb
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 32 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/server/security/AbstractBearerAuthenticator.java b/core/trino-main/src/main/java/io/trino/server/security/AbstractBearerAuthenticator.java
@@ -18,24 +18,15 @@
 
 import javax.ws.rs.container.ContainerRequestContext;
 
-import java.security.Principal;
 import java.util.List;
 import java.util.Optional;
 
 import static com.google.common.net.HttpHeaders.AUTHORIZATION;
 import static java.lang.String.format;
-import static java.util.Objects.requireNonNull;
 
 public abstract class AbstractBearerAuthenticator
         implements Authenticator
 {
-    private final UserMapping userMapping;
-
-    protected AbstractBearerAuthenticator(UserMapping userMapping)
-    {
-        this.userMapping = requireNonNull(userMapping, "userMapping is null");
-    }
-
     @Override
     public Identity authenticate(ContainerRequestContext request)
             throws AuthenticationException
@@ -47,15 +38,7 @@ public Identity authenticate(ContainerRequestContext request, String token)
             throws AuthenticationException
     {
         try {
-            Optional<Principal> principal = extractPrincipalFromToken(token);
-            if (principal.isEmpty()) {
-                throw needAuthentication(request, "Invalid credentials");
-            }
-
-            String authenticatedUser = userMapping.mapUser(principal.get().getName());
-            return Identity.forUser(authenticatedUser)
-                    .withPrincipal(principal.get())
-                    .build();
+            return createIdentity(token).orElseThrow(() -> needAuthentication(request, "Invalid credentials"));
         }
         catch (JwtException | UserMappingException e) {
             throw needAuthentication(request, e.getMessage());
@@ -88,7 +71,8 @@ public String extractToken(ContainerRequestContext request)
         return token;
     }
 
-    protected abstract Optional<Principal> extractPrincipalFromToken(String token);
+    protected abstract Optional<Identity> createIdentity(String token)
+            throws UserMappingException;
 
     protected abstract AuthenticationException needAuthentication(ContainerRequestContext request, String message);
 }
diff --git a/core/trino-main/src/main/java/io/trino/server/security/jwt/JwtAuthenticator.java b/core/trino-main/src/main/java/io/trino/server/security/jwt/JwtAuthenticator.java
@@ -19,12 +19,14 @@
 import io.jsonwebtoken.SigningKeyResolver;
 import io.trino.server.security.AbstractBearerAuthenticator;
 import io.trino.server.security.AuthenticationException;
+import io.trino.server.security.UserMapping;
+import io.trino.server.security.UserMappingException;
 import io.trino.spi.security.BasicPrincipal;
+import io.trino.spi.security.Identity;
 
 import javax.inject.Inject;
 import javax.ws.rs.container.ContainerRequestContext;
 
-import java.security.Principal;
 import java.util.Optional;
 
 import static io.trino.server.security.UserMapping.createUserMapping;
@@ -34,11 +36,11 @@ public class JwtAuthenticator
 {
     private final JwtParser jwtParser;
     private final String principalField;
+    private final UserMapping userMapping;
 
     @Inject
     public JwtAuthenticator(JwtAuthenticatorConfig config, SigningKeyResolver signingKeyResolver)
     {
-        super(createUserMapping(config.getUserMappingPattern(), config.getUserMappingFile()));
         principalField = config.getPrincipalField();
 
         JwtParserBuilder jwtParser = Jwts.parserBuilder()
@@ -51,15 +53,22 @@ public JwtAuthenticator(JwtAuthenticatorConfig config, SigningKeyResolver signin
             jwtParser.requireAudience(config.getRequiredAudience());
         }
         this.jwtParser = jwtParser.build();
+        userMapping = createUserMapping(config.getUserMappingPattern(), config.getUserMappingFile());
     }
 
     @Override
-    protected Optional<Principal> extractPrincipalFromToken(String token)
+    protected Optional<Identity> createIdentity(String token)
+            throws UserMappingException
     {
-        return Optional.ofNullable(jwtParser.parseClaimsJws(token)
+        Optional<String> principal = Optional.ofNullable(jwtParser.parseClaimsJws(token)
                 .getBody()
-                .get(principalField, String.class))
-                .map(BasicPrincipal::new);
+                .get(principalField, String.class));
+        if (principal.isEmpty()) {
+            return Optional.empty();
+        }
+        return Optional.of(Identity.forUser(userMapping.mapUser(principal.get()))
+                .withPrincipal(new BasicPrincipal(principal.get()))
+                .build());
     }
 
     @Override

diff --git a/core/trino-main/src/main/java/io/trino/server/security/oauth2/OAuth2Authenticator.java b/core/trino-main/src/main/java/io/trino/server/security/oauth2/OAuth2Authenticator.java
@@ -15,13 +15,16 @@
 
 import io.trino.server.security.AbstractBearerAuthenticator;
 import io.trino.server.security.AuthenticationException;
+import io.trino.server.security.UserMapping;
+import io.trino.server.security.UserMappingException;
 import io.trino.spi.security.BasicPrincipal;
+import io.trino.spi.security.Identity;
 
 import javax.inject.Inject;
 import javax.ws.rs.container.ContainerRequestContext;
 
 import java.net.URI;
-import java.security.Principal;
+import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
 
@@ -36,23 +39,29 @@ public class OAuth2Authenticator
 {
     private final OAuth2Service service;
     private final String principalField;
+    private final UserMapping userMapping;
 
     @Inject
     public OAuth2Authenticator(OAuth2Service service, OAuth2Config config)
     {
-        super(createUserMapping(config.getUserMappingPattern(), config.getUserMappingFile()));
         this.service = requireNonNull(service, "service is null");
         this.principalField = config.getPrincipalField();
+        userMapping = createUserMapping(config.getUserMappingPattern(), config.getUserMappingFile());
     }
 
     @Override
-    protected Optional<Principal> extractPrincipalFromToken(String token)
+    protected Optional<Identity> createIdentity(String token)
+            throws UserMappingException
     {
         try {
-            return service.convertTokenToClaims(token)
-                    .map(claims -> claims.get(principalField))
-                    .map(String.class::cast)
-                    .map(BasicPrincipal::new);
+            Optional<Map<String, Object>> claims = service.convertTokenToClaims(token);
+            if (claims.isEmpty()) {
+                return Optional.empty();
+            }
+            String principal = (String) claims.get().get(principalField);
+            return Optional.of(Identity.forUser(userMapping.mapUser(principal))
+                    .withPrincipal(new BasicPrincipal(principal))
+                    .build());
         }
         catch (ChallengeFailedException e) {
             return Optional.empty();