Skip to content

Commit d3cc501

Browse files
sapessisapessi
committed
Experimental ALB support to address issue #214 
1 parent 3d8b655 commit d3cc501

File tree

18 files changed

+209
-57
lines changed

18 files changed

+209
-57
lines changed

‎aws-serverless-java-container-core/src/main/java/com/amazonaws/serverless/proxy/RequestReader.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,11 @@ public abstract class RequestReader<RequestType, ContainerRequestType> {
4545
*/
4646
public static final String API_GATEWAY_STAGE_VARS_PROPERTY = "com.amazonaws.apigateway.stage.variables";
4747

48+
/**
49+
* The key for the <strong>ALB context</strong> property in the PropertiesDelegate object
50+
*/
51+
public static final String ALB_CONTEXT_PROPERTY = "com.amazonaws.alb.request.context";
52+
4853
/**
4954
* The key to store the entire API Gateway event
5055
*/

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

Lines changed: 37 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,23 @@
1212
*/
1313
package com.amazonaws.serverless.proxy.internal.jaxrs;
1414

15+
import com.amazonaws.serverless.proxy.internal.LambdaContainerHandler;
1516
import com.amazonaws.serverless.proxy.model.AwsProxyRequest;
1617
import com.amazonaws.serverless.proxy.model.CognitoAuthorizerClaims;
1718
import com.amazonaws.services.lambda.runtime.Context;
1819

20+
import com.fasterxml.jackson.core.type.TypeReference;
21+
1922
import javax.ws.rs.core.SecurityContext;
23+
24+
import java.io.IOException;
2025
import java.security.Principal;
26+
import java.util.Base64;
27+
import java.util.Map;
28+
29+
import static com.amazonaws.serverless.proxy.model.AwsProxyRequest.*;
30+
import static com.amazonaws.serverless.proxy.model.AwsProxyRequest.RequestSource.API_GATEWAY;
31+
2132

2233
/**
2334
* default implementation of the <code>SecurityContext</code> object. This class supports 3 API Gateway's authorization methods:
@@ -38,6 +49,9 @@ public class AwsProxySecurityContext
3849
private static final String AUTH_SCHEME_COGNITO_POOL = "COGNITO_USER_POOL";
3950
private static final String AUTH_SCHEME_AWS_IAM = "AWS_IAM";
4051

52+
private static final String ALB_ACESS_TOKEN_HEADER = "x-amzn-oidc-accesstoken";
53+
private static final String ALB_IDENTITY_HEADER = "x-amzn-oidc-identity";
54+
4155

4256
//-------------------------------------------------------------
4357
// Variables - Private
@@ -78,7 +92,12 @@ public Principal getUserPrincipal() {
7892
if (getAuthenticationScheme().equals(AUTH_SCHEME_CUSTOM) || getAuthenticationScheme().equals(AUTH_SCHEME_AWS_IAM)) {
7993
return () -> {
8094
if (getAuthenticationScheme().equals(AUTH_SCHEME_CUSTOM)) {
81-
return event.getRequestContext().getAuthorizer().getPrincipalId();
95+
switch (event.getRequestSource()) {
96+
case API_GATEWAY:
97+
return event.getRequestContext().getAuthorizer().getPrincipalId();
98+
case ALB:
99+
return event.getMultiValueHeaders().getFirst(ALB_IDENTITY_HEADER);
100+
}
82101
} else if (getAuthenticationScheme().equals(AUTH_SCHEME_AWS_IAM)) {
83102
// if we received credentials from Cognito Federated Identities then we return the identity id
84103
if (event.getRequestContext().getIdentity().getCognitoIdentityId() != null) {
@@ -112,15 +131,24 @@ public boolean isSecure() {
112131

113132

114133
public String getAuthenticationScheme() {
115-
if (event.getRequestContext().getAuthorizer() != null && event.getRequestContext().getAuthorizer().getClaims() != null && event.getRequestContext().getAuthorizer().getClaims().getSubject() != null) {
116-
return AUTH_SCHEME_COGNITO_POOL;
117-
} else if (event.getRequestContext().getAuthorizer() != null) {
118-
return AUTH_SCHEME_CUSTOM;
119-
} else if (event.getRequestContext().getIdentity().getAccessKey() != null) {
120-
return AUTH_SCHEME_AWS_IAM;
121-
} else {
122-
return null;
134+
switch (event.getRequestSource()) {
135+
case API_GATEWAY:
136+
if (event.getRequestContext().getAuthorizer() != null && event.getRequestContext().getAuthorizer().getClaims() != null
137+
&& event.getRequestContext().getAuthorizer().getClaims().getSubject() != null) {
138+
return AUTH_SCHEME_COGNITO_POOL;
139+
} else if (event.getRequestContext().getAuthorizer() != null) {
140+
return AUTH_SCHEME_CUSTOM;
141+
} else if (event.getRequestContext().getIdentity().getAccessKey() != null) {
142+
return AUTH_SCHEME_AWS_IAM;
143+
} else {
144+
return null;
145+
}
146+
case ALB:
147+
if (event.getMultiValueHeaders().containsKey(ALB_ACESS_TOKEN_HEADER)) {
148+
return AUTH_SCHEME_CUSTOM;
149+
}
123150
}
151+
return null;
124152
}
125153

126154

‎aws-serverless-java-container-core/src/main/java/com/amazonaws/serverless/proxy/internal/servlet/ApacheCombinedServletLogFormatter.java

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,8 @@
22

33

44
import com.amazonaws.serverless.proxy.LogFormatter;
5-
import com.amazonaws.serverless.proxy.model.ApiGatewayRequestContext;
5+
import com.amazonaws.serverless.proxy.model.AwsProxyRequest;
6+
import com.amazonaws.serverless.proxy.model.AwsProxyRequestContext;
67

78
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
89

@@ -19,6 +20,7 @@
1920
import java.util.Locale;
2021

2122
import static com.amazonaws.serverless.proxy.RequestReader.API_GATEWAY_CONTEXT_PROPERTY;
23+
import static com.amazonaws.serverless.proxy.RequestReader.API_GATEWAY_EVENT_PROPERTY;
2224
import static java.time.temporal.ChronoField.DAY_OF_MONTH;
2325
import static java.time.temporal.ChronoField.HOUR_OF_DAY;
2426
import static java.time.temporal.ChronoField.MINUTE_OF_HOUR;
@@ -64,14 +66,15 @@ public ApacheCombinedServletLogFormatter() {
6466
public String format(ContainerRequestType servletRequest, ContainerResponseType servletResponse, SecurityContext ctx) {
6567
//LogFormat "%h %l %u %t \"%r\" %>s %b \"%{Referer}i\" \"%{User-agent}i\"" combined
6668
StringBuilder logLineBuilder = new StringBuilder();
67-
ApiGatewayRequestContext gatewayContext = (ApiGatewayRequestContext)servletRequest.getAttribute(API_GATEWAY_CONTEXT_PROPERTY);
69+
AwsProxyRequest req = (AwsProxyRequest)servletRequest.getAttribute(API_GATEWAY_EVENT_PROPERTY);
70+
AwsProxyRequestContext gatewayContext = req.getRequestContext();
6871

6972
// %h
7073
logLineBuilder.append(servletRequest.getRemoteAddr());
7174
logLineBuilder.append(" ");
7275

7376
// %l
74-
if (gatewayContext != null) {
77+
if (gatewayContext != null && req.getRequestSource() == AwsProxyRequest.RequestSource.API_GATEWAY) {
7578
if (gatewayContext.getIdentity().getUserArn() != null) {
7679
logLineBuilder.append(gatewayContext.getIdentity().getUserArn());
7780
} else {

‎aws-serverless-java-container-core/src/main/java/com/amazonaws/serverless/proxy/internal/servlet/AwsHttpServletRequest.java

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,12 +14,11 @@
1414

1515
import com.amazonaws.serverless.proxy.RequestReader;
1616
import com.amazonaws.serverless.proxy.internal.SecurityUtils;
17-
import com.amazonaws.serverless.proxy.model.ApiGatewayRequestContext;
17+
import com.amazonaws.serverless.proxy.model.AwsProxyRequestContext;
1818
import com.amazonaws.serverless.proxy.model.ContainerConfig;
1919
import com.amazonaws.serverless.proxy.model.MultiValuedTreeMap;
2020
import com.amazonaws.services.lambda.runtime.Context;
2121

22-
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
2322
import org.slf4j.Logger;
2423
import org.slf4j.LoggerFactory;
2524

@@ -34,16 +33,13 @@
3433
import java.io.UnsupportedEncodingException;
3534
import java.net.URLDecoder;
3635
import java.net.URLEncoder;
37-
import java.nio.charset.StandardCharsets;
3836
import java.time.format.DateTimeFormatter;
39-
import java.util.AbstractMap;
4037
import java.util.ArrayList;
4138
import java.util.Collections;
4239
import java.util.Enumeration;
4340
import java.util.HashMap;
4441
import java.util.List;
4542
import java.util.Map;
46-
import java.util.stream.Collectors;
4743

4844

4945
/**
@@ -116,7 +112,7 @@ public String getRequestedSessionId() {
116112
public HttpSession getSession(boolean b) {
117113
log.debug("Trying to access session. Lambda functions are stateless and should not rely on the session");
118114
if (b && null == this.session) {
119-
ApiGatewayRequestContext requestContext = (ApiGatewayRequestContext) getAttribute(RequestReader.API_GATEWAY_CONTEXT_PROPERTY);
115+
AwsProxyRequestContext requestContext = (AwsProxyRequestContext) getAttribute(RequestReader.API_GATEWAY_CONTEXT_PROPERTY);
120116
this.session = new AwsHttpSession(requestContext.getRequestId());
121117
}
122118
return this.session;

‎aws-serverless-java-container-core/src/main/java/com/amazonaws/serverless/proxy/internal/servlet/AwsHttpServletResponse.java

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

1515
import com.amazonaws.serverless.proxy.internal.SecurityUtils;
16+
import com.amazonaws.serverless.proxy.model.AwsProxyRequest;
1617
import com.amazonaws.serverless.proxy.model.MultiValuedTreeMap;
1718

1819
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
@@ -33,6 +34,9 @@
3334
import java.util.*;
3435
import java.util.concurrent.CountDownLatch;
3536

37+
import static com.amazonaws.serverless.proxy.RequestReader.API_GATEWAY_EVENT_PROPERTY;
38+
39+
3640
/**
3741
* Basic implementation of the <code>HttpServletResponse</code> object. This is used by the <code>AwsProxyHttpServletResponseWriter</code>
3842
* to generate an <code>AwsProxyResponse</code> object. We have an additional <code>getAwsResponseHeaders()</code> method
@@ -443,6 +447,10 @@ MultiValuedTreeMap<String, String> getAwsResponseHeaders() {
443447
return headers;
444448
}
445449

450+
AwsProxyRequest getAwsProxyRequest() {
451+
return (AwsProxyRequest)request.getAttribute(API_GATEWAY_EVENT_PROPERTY);
452+
}
453+
446454

447455
//-------------------------------------------------------------
448456
// Methods - Private

‎aws-serverless-java-container-core/src/main/java/com/amazonaws/serverless/proxy/internal/servlet/AwsProxyHttpServletRequest.java

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -800,13 +800,15 @@ private List<String> getHeaderValues(String key) {
800800
// special cases for referer and user agent headers
801801
List<String> values = new ArrayList<>();
802802

803-
if ("referer".equals(key.toLowerCase(Locale.ENGLISH))) {
804-
values.add(request.getRequestContext().getIdentity().getCaller());
805-
return values;
806-
}
807-
if ("user-agent".equals(key.toLowerCase(Locale.ENGLISH))) {
808-
values.add(request.getRequestContext().getIdentity().getUserAgent());
809-
return values;
803+
if (request.getRequestSource() == AwsProxyRequest.RequestSource.API_GATEWAY) {
804+
if ("referer".equals(key.toLowerCase(Locale.ENGLISH))) {
805+
values.add(request.getRequestContext().getIdentity().getCaller());
806+
return values;
807+
}
808+
if ("user-agent".equals(key.toLowerCase(Locale.ENGLISH))) {
809+
values.add(request.getRequestContext().getIdentity().getUserAgent());
810+
return values;
811+
}
810812
}
811813

812814
if (request.getMultiValueHeaders() == null) {

‎aws-serverless-java-container-core/src/main/java/com/amazonaws/serverless/proxy/internal/servlet/AwsProxyHttpServletRequestReader.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ public class AwsProxyHttpServletRequestReader extends RequestReader<AwsProxyRequ
3535
@Override
3636
public AwsProxyHttpServletRequest readRequest(AwsProxyRequest request, SecurityContext securityContext, Context lambdaContext, ContainerConfig config)
3737
throws InvalidRequestEventException {
38+
3839
request.setPath(stripBasePath(request.getPath(), config));
3940
if (request.getMultiValueHeaders().getFirst(HttpHeaders.CONTENT_TYPE) != null) {
4041
String contentType = request.getMultiValueHeaders().getFirst(HttpHeaders.CONTENT_TYPE);
@@ -45,6 +46,7 @@ public AwsProxyHttpServletRequest readRequest(AwsProxyRequest request, SecurityC
4546
servletRequest.setAttribute(API_GATEWAY_CONTEXT_PROPERTY, request.getRequestContext());
4647
servletRequest.setAttribute(API_GATEWAY_STAGE_VARS_PROPERTY, request.getStageVariables());
4748
servletRequest.setAttribute(API_GATEWAY_EVENT_PROPERTY, request);
49+
servletRequest.setAttribute(ALB_CONTEXT_PROPERTY, request.getRequestContext().getElb());
4850
servletRequest.setAttribute(LAMBDA_CONTEXT_PROPERTY, lambdaContext);
4951
servletRequest.setAttribute(JAX_SECURITY_CONTEXT_PROPERTY, securityContext);
5052

‎aws-serverless-java-container-core/src/main/java/com/amazonaws/serverless/proxy/internal/servlet/AwsProxyHttpServletResponseWriter.java

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,18 @@
1717
import com.amazonaws.serverless.proxy.ResponseWriter;
1818
import com.amazonaws.serverless.proxy.internal.LambdaContainerHandler;
1919
import com.amazonaws.serverless.proxy.internal.testutils.Timer;
20+
import com.amazonaws.serverless.proxy.model.AwsProxyRequest;
2021
import com.amazonaws.serverless.proxy.model.AwsProxyResponse;
2122
import com.amazonaws.services.lambda.runtime.Context;
2223

24+
import org.apache.http.HttpStatus;
25+
26+
import javax.ws.rs.core.Response;
27+
2328
import java.util.Base64;
29+
import java.util.HashMap;
30+
import java.util.Map;
31+
2432

2533
/**
2634
* Creates an <code>AwsProxyResponse</code> object given an <code>AwsHttpServletResponse</code> object. If the
@@ -53,6 +61,10 @@ public AwsProxyResponse writeResponse(AwsHttpServletResponse containerResponse,
5361

5462
awsProxyResponse.setStatusCode(containerResponse.getStatus());
5563

64+
if (containerResponse.getAwsProxyRequest().getRequestSource() == AwsProxyRequest.RequestSource.ALB) {
65+
awsProxyResponse.setStatusDescription(containerResponse.getStatus() + " " + Response.Status.fromStatusCode(containerResponse.getStatus()).getReasonPhrase());
66+
}
67+
5668
Timer.stop("SERVLET_RESPONSE_WRITE");
5769
return awsProxyResponse;
5870
}

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

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -13,19 +13,18 @@
1313
package com.amazonaws.serverless.proxy.internal.testutils;
1414

1515
import com.amazonaws.serverless.proxy.internal.LambdaContainerHandler;
16+
import com.amazonaws.serverless.proxy.model.AlbContext;
1617
import com.amazonaws.serverless.proxy.model.ApiGatewayAuthorizerContext;
17-
import com.amazonaws.serverless.proxy.model.ApiGatewayRequestContext;
18+
import com.amazonaws.serverless.proxy.model.AwsProxyRequestContext;
1819
import com.amazonaws.serverless.proxy.model.ApiGatewayRequestIdentity;
1920
import com.amazonaws.serverless.proxy.model.AwsProxyRequest;
2021
import com.amazonaws.serverless.proxy.model.CognitoAuthorizerClaims;
2122
import com.amazonaws.serverless.proxy.model.MultiValuedTreeMap;
2223

2324
import com.fasterxml.jackson.core.JsonProcessingException;
24-
import com.fasterxml.jackson.databind.ObjectMapper;
2525
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
2626
import org.apache.commons.io.IOUtils;
2727
import org.apache.http.HttpEntity;
28-
import org.apache.http.entity.mime.FormBodyPart;
2928
import org.apache.http.entity.mime.MultipartEntityBuilder;
3029
import org.apache.http.entity.mime.content.ByteArrayBody;
3130

@@ -39,7 +38,6 @@
3938
import java.nio.charset.Charset;
4039
import java.nio.charset.StandardCharsets;
4140
import java.util.Base64;
42-
import java.util.HashMap;
4341
import java.util.UUID;
4442

4543

@@ -70,13 +68,12 @@ public AwsProxyRequestBuilder(String path) {
7068

7169

7270
public AwsProxyRequestBuilder(String path, String httpMethod) {
73-
7471
this.request = new AwsProxyRequest();
7572
this.request.setMultiValueHeaders(new MultiValuedTreeMap<>(String.CASE_INSENSITIVE_ORDER)); // avoid NPE
7673
this.request.setHttpMethod(httpMethod);
7774
this.request.setPath(path);
7875
this.request.setMultiValueQueryStringParameters(new MultiValuedTreeMap<>());
79-
this.request.setRequestContext(new ApiGatewayRequestContext());
76+
this.request.setRequestContext(new AwsProxyRequestContext());
8077
this.request.getRequestContext().setRequestId(UUID.randomUUID().toString());
8178
this.request.getRequestContext().setExtendedRequestId(UUID.randomUUID().toString());
8279
this.request.getRequestContext().setStage("test");
@@ -92,6 +89,15 @@ public AwsProxyRequestBuilder(String path, String httpMethod) {
9289
// Methods - Public
9390
//-------------------------------------------------------------
9491

92+
public AwsProxyRequestBuilder alb() {
93+
this.request.setRequestContext(new AwsProxyRequestContext());
94+
this.request.getRequestContext().setElb(new AlbContext());
95+
this.request.getRequestContext().getElb().setTargetGroupArn(
96+
"arn:aws:elasticloadbalancing:us-east-1:123456789012:targetgroup/lambda-target/d6190d154bc908a5"
97+
);
98+
return this;
99+
}
100+
95101
public AwsProxyRequestBuilder stage(String stageName) {
96102
this.request.getRequestContext().setStage(stageName);
97103
return this;
@@ -294,7 +300,7 @@ public AwsProxyRequestBuilder serverName(String serverName) {
294300

295301
public AwsProxyRequestBuilder userAgent(String agent) {
296302
if (request.getRequestContext() == null) {
297-
request.setRequestContext(new ApiGatewayRequestContext());
303+
request.setRequestContext(new AwsProxyRequestContext());
298304
}
299305
if (request.getRequestContext().getIdentity() == null) {
300306
request.getRequestContext().setIdentity(new ApiGatewayRequestIdentity());
@@ -306,7 +312,7 @@ public AwsProxyRequestBuilder userAgent(String agent) {
306312

307313
public AwsProxyRequestBuilder referer(String referer) {
308314
if (request.getRequestContext() == null) {
309-
request.setRequestContext(new ApiGatewayRequestContext());
315+
request.setRequestContext(new AwsProxyRequestContext());
310316
}
311317
if (request.getRequestContext().getIdentity() == null) {
312318
request.getRequestContext().setIdentity(new ApiGatewayRequestIdentity());

‎aws-serverless-java-container-core/src/main/java/com/amazonaws/serverless/proxy/model/AlbContext.java

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
package com.amazonaws.serverless.proxy.model;
2+
3+
4+
/***
5+
* Context passed by ALB proxy events
6+
*/
7+
public class AlbContext {
8+
private String targetGroupArn;
9+
10+
11+
public String getTargetGroupArn() {
12+
return targetGroupArn;
13+
}
14+
15+
16+
public void setTargetGroupArn(String targetGroupArn) {
17+
this.targetGroupArn = targetGroupArn;
18+
}
19+
}

0 commit comments

Comments
 (0)