Skip to content

Commit ab7160b

Browse files
committed
Changes to address #55. Added a new CognitoUserPoolPrincipal object that exposes claims from the token. Also added support for custom claims.
1 parent 558565c commit ab7160b

File tree

5 files changed

+92
-19
lines changed

5 files changed

+92
-19
lines changed

aws-serverless-java-container-core/src/main/java/com/amazonaws/serverless/exceptions/InvalidRequestEventException.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,4 +22,8 @@ public class InvalidRequestEventException extends Exception {
2222
public InvalidRequestEventException(String message, Exception e) {
2323
super(message, e);
2424
}
25+
26+
public InvalidRequestEventException(String message) {
27+
super(message);
28+
}
2529
}

aws-serverless-java-container-core/src/main/java/com/amazonaws/serverless/proxy/internal/jaxrs/AwsProxySecurityContext.java

Lines changed: 49 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
package com.amazonaws.serverless.proxy.internal.jaxrs;
1414

1515
import com.amazonaws.serverless.proxy.internal.model.AwsProxyRequest;
16+
import com.amazonaws.serverless.proxy.internal.model.CognitoAuthorizerClaims;
1617
import com.amazonaws.services.lambda.runtime.Context;
1718

1819
import javax.ws.rs.core.SecurityContext;
@@ -61,26 +62,33 @@ public AwsProxySecurityContext(final Context lambdaContext, final AwsProxyReques
6162
//-------------------------------------------------------------
6263

6364
public Principal getUserPrincipal() {
64-
return () -> {
65-
if (getAuthenticationScheme() == null) {
66-
return null;
67-
}
68-
69-
if (getAuthenticationScheme().equals(AUTH_SCHEME_CUSTOM)) {
70-
return event.getRequestContext().getAuthorizer().getPrincipalId();
71-
} else if (getAuthenticationScheme().equals(AUTH_SCHEME_AWS_IAM)) {
72-
// if we received credentials from Cognito Federated Identities then we return the identity id
73-
if (event.getRequestContext().getIdentity().getCognitoIdentityId() != null) {
74-
return event.getRequestContext().getIdentity().getCognitoIdentityId();
75-
} else { // otherwise the user arn from the credentials
76-
return event.getRequestContext().getIdentity().getUserArn();
65+
if (getAuthenticationScheme() == null) {
66+
return () -> null;
67+
}
68+
69+
if (getAuthenticationScheme().equals(AUTH_SCHEME_CUSTOM) || getAuthenticationScheme().equals(AUTH_SCHEME_AWS_IAM)) {
70+
return () -> {
71+
if (getAuthenticationScheme().equals(AUTH_SCHEME_CUSTOM)) {
72+
return event.getRequestContext().getAuthorizer().getPrincipalId();
73+
} else if (getAuthenticationScheme().equals(AUTH_SCHEME_AWS_IAM)) {
74+
// if we received credentials from Cognito Federated Identities then we return the identity id
75+
if (event.getRequestContext().getIdentity().getCognitoIdentityId() != null) {
76+
return event.getRequestContext().getIdentity().getCognitoIdentityId();
77+
} else { // otherwise the user arn from the credentials
78+
return event.getRequestContext().getIdentity().getUserArn();
79+
}
7780
}
78-
} else if (getAuthenticationScheme().equals(AUTH_SCHEME_COGNITO_POOL)) {
79-
return event.getRequestContext().getAuthorizer().getClaims().getSubject();
80-
}
8181

82-
return null;
83-
};
82+
// return null if we couldn't find a valid scheme
83+
return null;
84+
};
85+
}
86+
87+
if (getAuthenticationScheme().equals(AUTH_SCHEME_COGNITO_POOL)) {
88+
return new CognitoUserPoolPrincipal(event.getRequestContext().getAuthorizer().getClaims());
89+
}
90+
91+
throw new RuntimeException("Cannot recognize authorization scheme in event");
8492
}
8593

8694

@@ -105,4 +113,27 @@ public String getAuthenticationScheme() {
105113
return null;
106114
}
107115
}
116+
117+
118+
/**
119+
* Custom object for request authorized with a Cognito User Pool authorizer. By casting the Principal
120+
* object to this you can extract the Claims object included in the token.
121+
*/
122+
public class CognitoUserPoolPrincipal implements Principal {
123+
124+
private CognitoAuthorizerClaims claims;
125+
126+
CognitoUserPoolPrincipal(CognitoAuthorizerClaims c) {
127+
claims = c;
128+
}
129+
130+
@Override
131+
public String getName() {
132+
return claims.getSubject();
133+
}
134+
135+
public CognitoAuthorizerClaims getClaims() {
136+
return claims;
137+
}
138+
}
108139
}

aws-serverless-java-container-core/src/main/java/com/amazonaws/serverless/proxy/internal/model/CognitoAuthorizerClaims.java

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,14 @@
1313
package com.amazonaws.serverless.proxy.internal.model;
1414

1515

16+
import com.fasterxml.jackson.annotation.JsonAnyGetter;
17+
import com.fasterxml.jackson.annotation.JsonAnySetter;
1618
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
1719
import com.fasterxml.jackson.annotation.JsonProperty;
1820

1921
import java.time.format.DateTimeFormatter;
22+
import java.util.HashMap;
23+
import java.util.Map;
2024

2125

2226
/**
@@ -44,6 +48,8 @@ public class CognitoAuthorizerClaims {
4448
// Variables - Private
4549
//-------------------------------------------------------------
4650

51+
private Map<String, String> claims = new HashMap<>();
52+
4753
@JsonProperty(value = "sub")
4854
private String subject;
4955
@JsonProperty(value = "aud")
@@ -69,6 +75,16 @@ public class CognitoAuthorizerClaims {
6975
// Methods - Getter/Setter
7076
//-------------------------------------------------------------
7177

78+
@JsonAnyGetter
79+
public String getClaim(String claim) {
80+
return claims.get(claim);
81+
}
82+
83+
@JsonAnySetter
84+
public void setClaim(String claim, String value) {
85+
claims.put(claim, value);
86+
}
87+
7288
public String getSubject() { return this.subject; }
7389

7490

aws-serverless-java-container-core/src/main/java/com/amazonaws/serverless/proxy/internal/testutils/AwsProxyRequestBuilder.java

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,12 @@ public AwsProxyRequestBuilder cognitoUserPool(String identityId) {
189189
return this;
190190
}
191191

192+
public AwsProxyRequestBuilder claim(String claim, String value) {
193+
this.request.getRequestContext().getAuthorizer().getClaims().setClaim(claim, value);
194+
195+
return this;
196+
}
197+
192198

193199
public AwsProxyRequestBuilder cognitoIdentity(String identityId, String identityPoolId) {
194200
this.request.getRequestContext().getIdentity().setCognitoAuthenticationType("IDENTITY");

aws-serverless-java-container-core/src/test/java/com/amazonaws/serverless/proxy/internal/jaxrs/AwsProxySecurityContextTest.java

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,17 @@
44
import com.amazonaws.serverless.proxy.internal.testutils.AwsProxyRequestBuilder;
55
import org.junit.Test;
66

7+
import java.security.Principal;
8+
79
import static org.junit.Assert.*;
810

911
public class AwsProxySecurityContextTest {
12+
private static final String CLAIM_KEY = "custom:claim";
13+
private static final String CLAIM_VALUE = "customClaimant";
1014
private static final String COGNITO_IDENTITY_ID = "us-east-2:123123123123";
1115
private static final AwsProxyRequest REQUEST_NO_AUTH = new AwsProxyRequestBuilder("/hello", "GET").build();
1216
private static final AwsProxyRequest REQUEST_COGNITO_USER_POOL = new AwsProxyRequestBuilder("/hello", "GET")
13-
.cognitoUserPool(COGNITO_IDENTITY_ID).build();
17+
.cognitoUserPool(COGNITO_IDENTITY_ID).claim(CLAIM_KEY, CLAIM_VALUE).build();
1418

1519
@Test
1620
public void localVars_constructor_nullValues() {
@@ -39,4 +43,16 @@ public void authScheme_getPrincipal_userPool() {
3943
assertEquals("COGNITO_USER_POOL", context.getAuthenticationScheme());
4044
assertEquals(COGNITO_IDENTITY_ID, context.getUserPrincipal().getName());
4145
}
46+
47+
@Test
48+
public void userPool_getClaims_retrieveCustomClaim() {
49+
AwsProxySecurityContext context = new AwsProxySecurityContext(null, REQUEST_COGNITO_USER_POOL);
50+
Principal userPrincipal = context.getUserPrincipal();
51+
assertNotNull(userPrincipal.getName());
52+
assertEquals(COGNITO_IDENTITY_ID, userPrincipal.getName());
53+
54+
assertTrue(userPrincipal instanceof AwsProxySecurityContext.CognitoUserPoolPrincipal);
55+
assertNotNull(((AwsProxySecurityContext.CognitoUserPoolPrincipal)userPrincipal).getClaims().getClaim(CLAIM_KEY));
56+
assertEquals(CLAIM_VALUE, ((AwsProxySecurityContext.CognitoUserPoolPrincipal)userPrincipal).getClaims().getClaim(CLAIM_KEY));
57+
}
4258
}

0 commit comments

Comments
 (0)