Changes to address #55. Added a new CognitoUserPoolPrincipal object t…

…hat exposes claims from the token. Also added support for custom claims.
aws · Nov 21, 2017 · ab7160b · ab7160b
1 parent 558565c
commit ab7160b
Show file tree

Hide file tree

Showing 5 changed files with 92 additions and 19 deletions.
diff --git a/...-core/src/main/java/com/amazonaws/serverless/exceptions/InvalidRequestEventException.java b/...-core/src/main/java/com/amazonaws/serverless/exceptions/InvalidRequestEventException.java
@@ -22,4 +22,8 @@ public class InvalidRequestEventException extends Exception {
     public InvalidRequestEventException(String message, Exception e) {
         super(message, e);
     }
+
+    public InvalidRequestEventException(String message) {
+        super(message);
+    }
 }
diff --git a/.../src/main/java/com/amazonaws/serverless/proxy/internal/jaxrs/AwsProxySecurityContext.java b/.../src/main/java/com/amazonaws/serverless/proxy/internal/jaxrs/AwsProxySecurityContext.java
@@ -13,6 +13,7 @@
 package com.amazonaws.serverless.proxy.internal.jaxrs;
 
 import com.amazonaws.serverless.proxy.internal.model.AwsProxyRequest;
+import com.amazonaws.serverless.proxy.internal.model.CognitoAuthorizerClaims;
 import com.amazonaws.services.lambda.runtime.Context;
 
 import javax.ws.rs.core.SecurityContext;
@@ -61,26 +62,33 @@ public AwsProxySecurityContext(final Context lambdaContext, final AwsProxyReques
     //-------------------------------------------------------------
 
     public Principal getUserPrincipal() {
-        return () -> {
-            if (getAuthenticationScheme() == null) {
-                return null;
-            }
-
-            if (getAuthenticationScheme().equals(AUTH_SCHEME_CUSTOM)) {
-                return event.getRequestContext().getAuthorizer().getPrincipalId();
-            } else if (getAuthenticationScheme().equals(AUTH_SCHEME_AWS_IAM)) {
-                // if we received credentials from Cognito Federated Identities then we return the identity id
-                if (event.getRequestContext().getIdentity().getCognitoIdentityId() != null) {
-                    return event.getRequestContext().getIdentity().getCognitoIdentityId();
-                } else { // otherwise the user arn from the credentials
-                    return event.getRequestContext().getIdentity().getUserArn();
+        if (getAuthenticationScheme() == null) {
+            return () -> null;
+        }
+
+        if (getAuthenticationScheme().equals(AUTH_SCHEME_CUSTOM) || getAuthenticationScheme().equals(AUTH_SCHEME_AWS_IAM)) {
+            return () -> {
+                if (getAuthenticationScheme().equals(AUTH_SCHEME_CUSTOM)) {
+                    return event.getRequestContext().getAuthorizer().getPrincipalId();
+                } else if (getAuthenticationScheme().equals(AUTH_SCHEME_AWS_IAM)) {
+                    // if we received credentials from Cognito Federated Identities then we return the identity id
+                    if (event.getRequestContext().getIdentity().getCognitoIdentityId() != null) {
+                        return event.getRequestContext().getIdentity().getCognitoIdentityId();
+                    } else { // otherwise the user arn from the credentials
+                        return event.getRequestContext().getIdentity().getUserArn();
+                    }
                 }
-            } else if (getAuthenticationScheme().equals(AUTH_SCHEME_COGNITO_POOL)) {
-                return event.getRequestContext().getAuthorizer().getClaims().getSubject();
-            }
 
-            return null;
-        };
+                // return null if we couldn't find a valid scheme
+                return null;
+            };
+        }
+
+        if (getAuthenticationScheme().equals(AUTH_SCHEME_COGNITO_POOL)) {
+            return new CognitoUserPoolPrincipal(event.getRequestContext().getAuthorizer().getClaims());
+        }
+
+        throw new RuntimeException("Cannot recognize authorization scheme in event");
     }
 
 
@@ -105,4 +113,27 @@ public String getAuthenticationScheme() {
             return null;
         }
     }
+
+
+    /**
+     * Custom object for request authorized with a Cognito User Pool authorizer. By casting the Principal
+     * object to this you can extract the Claims object included in the token.
+     */
+    public class CognitoUserPoolPrincipal implements Principal {
+
+        private CognitoAuthorizerClaims claims;
+
+        CognitoUserPoolPrincipal(CognitoAuthorizerClaims c) {
+            claims = c;
+        }
+
+        @Override
+        public String getName() {
+            return claims.getSubject();
+        }
+
+        public CognitoAuthorizerClaims getClaims() {
+            return claims;
+        }
+    }
 }
diff --git a/.../src/main/java/com/amazonaws/serverless/proxy/internal/model/CognitoAuthorizerClaims.java b/.../src/main/java/com/amazonaws/serverless/proxy/internal/model/CognitoAuthorizerClaims.java
@@ -13,10 +13,14 @@
 package com.amazonaws.serverless.proxy.internal.model;
 
 
+import com.fasterxml.jackson.annotation.JsonAnyGetter;
+import com.fasterxml.jackson.annotation.JsonAnySetter;
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 import com.fasterxml.jackson.annotation.JsonProperty;
 
 import java.time.format.DateTimeFormatter;
+import java.util.HashMap;
+import java.util.Map;
 
 
 /**
@@ -44,6 +48,8 @@ public class CognitoAuthorizerClaims {
     // Variables - Private
     //-------------------------------------------------------------
 
+    private Map<String, String> claims = new HashMap<>();
+
     @JsonProperty(value = "sub")
     private String subject;
     @JsonProperty(value = "aud")
@@ -69,6 +75,16 @@ public class CognitoAuthorizerClaims {
     // Methods - Getter/Setter
     //-------------------------------------------------------------
 
+    @JsonAnyGetter
+    public String getClaim(String claim) {
+        return claims.get(claim);
+    }
+
+    @JsonAnySetter
+    public void setClaim(String claim, String value) {
+        claims.put(claim, value);
+    }
+
     public String getSubject() { return this.subject; }
 
 

diff --git a/...c/main/java/com/amazonaws/serverless/proxy/internal/testutils/AwsProxyRequestBuilder.java b/...c/main/java/com/amazonaws/serverless/proxy/internal/testutils/AwsProxyRequestBuilder.java
@@ -189,6 +189,12 @@ public AwsProxyRequestBuilder cognitoUserPool(String identityId) {
         return this;
     }
 
+    public AwsProxyRequestBuilder claim(String claim, String value) {
+        this.request.getRequestContext().getAuthorizer().getClaims().setClaim(claim, value);
+
+        return this;
+    }
+
 
     public AwsProxyRequestBuilder cognitoIdentity(String identityId, String identityPoolId) {
         this.request.getRequestContext().getIdentity().setCognitoAuthenticationType("IDENTITY");

diff --git a/.../test/java/com/amazonaws/serverless/proxy/internal/jaxrs/AwsProxySecurityContextTest.java b/.../test/java/com/amazonaws/serverless/proxy/internal/jaxrs/AwsProxySecurityContextTest.java
@@ -4,13 +4,17 @@
 import com.amazonaws.serverless.proxy.internal.testutils.AwsProxyRequestBuilder;
 import org.junit.Test;
 
+import java.security.Principal;
+
 import static org.junit.Assert.*;
 
 public class AwsProxySecurityContextTest {
+    private static final String CLAIM_KEY = "custom:claim";
+    private static final String CLAIM_VALUE = "customClaimant";
     private static final String COGNITO_IDENTITY_ID = "us-east-2:123123123123";
     private static final AwsProxyRequest REQUEST_NO_AUTH = new AwsProxyRequestBuilder("/hello", "GET").build();
     private static final AwsProxyRequest REQUEST_COGNITO_USER_POOL = new AwsProxyRequestBuilder("/hello", "GET")
-            .cognitoUserPool(COGNITO_IDENTITY_ID).build();
+            .cognitoUserPool(COGNITO_IDENTITY_ID).claim(CLAIM_KEY, CLAIM_VALUE).build();
 
     @Test
     public void localVars_constructor_nullValues() {
@@ -39,4 +43,16 @@ public void authScheme_getPrincipal_userPool() {
         assertEquals("COGNITO_USER_POOL", context.getAuthenticationScheme());
         assertEquals(COGNITO_IDENTITY_ID, context.getUserPrincipal().getName());
     }
+
+    @Test
+    public void userPool_getClaims_retrieveCustomClaim() {
+        AwsProxySecurityContext context = new AwsProxySecurityContext(null, REQUEST_COGNITO_USER_POOL);
+        Principal userPrincipal = context.getUserPrincipal();
+        assertNotNull(userPrincipal.getName());
+        assertEquals(COGNITO_IDENTITY_ID, userPrincipal.getName());
+
+        assertTrue(userPrincipal instanceof AwsProxySecurityContext.CognitoUserPoolPrincipal);
+        assertNotNull(((AwsProxySecurityContext.CognitoUserPoolPrincipal)userPrincipal).getClaims().getClaim(CLAIM_KEY));
+        assertEquals(CLAIM_VALUE, ((AwsProxySecurityContext.CognitoUserPoolPrincipal)userPrincipal).getClaims().getClaim(CLAIM_KEY));
+    }
 }