Merge pull request #491 from driverpt/add-alb-to-http-v2

Add ALB Context to API GW Payload V2

Author: Dennis Kieselhorst

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

@@ -14,15 +14,13 @@
 
 import com.amazonaws.serverless.proxy.model.AwsProxyRequest;
 import com.amazonaws.serverless.proxy.model.CognitoAuthorizerClaims;
+import com.amazonaws.serverless.proxy.model.RequestSource;
 import com.amazonaws.services.lambda.runtime.Context;
 
 import javax.ws.rs.core.SecurityContext;
 
 import java.security.Principal;
 
-import static com.amazonaws.serverless.proxy.model.AwsProxyRequest.RequestSource.API_GATEWAY;
-
-
 /**
  * default implementation of the <code>SecurityContext</code> object. This class supports 3 API Gateway's authorization methods:
  * AWS_IAM, CUSTOM_AUTHORIZER, and COGNITO_USER_POOL (oidc). The Principal returned by the object depends on the authorization

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

@@ -18,6 +18,7 @@
 import com.amazonaws.serverless.proxy.model.AwsProxyRequest;
 import com.amazonaws.serverless.proxy.model.ContainerConfig;
 import com.amazonaws.serverless.proxy.model.Headers;
+import com.amazonaws.serverless.proxy.model.RequestSource;
 import com.amazonaws.services.lambda.runtime.Context;
 
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
@@ -49,7 +50,6 @@
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
-
 /**
  * Implementation of the <code>HttpServletRequest</code> interface that supports <code>AwsProxyRequest</code> object.
  * This object is initialized with an <code>AwsProxyRequest</code> event and a <code>SecurityContext</code> generated
@@ -204,7 +204,7 @@ public String getQueryString() {
             return this.generateQueryString(
                     request.getMultiValueQueryStringParameters(),
                     // ALB does not automatically decode parameters, so we don't want to re-encode them
-                    request.getRequestSource() != AwsProxyRequest.RequestSource.ALB,
+                    request.getRequestSource() != RequestSource.ALB,
                     config.getUriEncoding());
         } catch (ServletException e) {
             log.error("Could not generate query string", e);
@@ -580,7 +580,7 @@ private List<String> getHeaderValues(String key) {
         // special cases for referer and user agent headers
         List<String> values = new ArrayList<>();
 
-        if (request.getRequestSource() == AwsProxyRequest.RequestSource.API_GATEWAY) {
+        if (request.getRequestSource() == RequestSource.API_GATEWAY) {
             if ("referer".equals(key.toLowerCase(Locale.ENGLISH))) {
                 values.add(request.getRequestContext().getIdentity().getCaller());
                 return values;

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

@@ -20,6 +20,7 @@
 import com.amazonaws.serverless.proxy.model.AwsProxyRequest;
 import com.amazonaws.serverless.proxy.model.AwsProxyResponse;
 import com.amazonaws.serverless.proxy.model.Headers;
+import com.amazonaws.serverless.proxy.model.RequestSource;
 import com.amazonaws.services.lambda.runtime.Context;
 
 import javax.ws.rs.core.Response;
@@ -73,7 +74,7 @@ public AwsProxyResponse writeResponse(AwsHttpServletResponse containerResponse,
 
         awsProxyResponse.setStatusCode(containerResponse.getStatus());
 
-        if (containerResponse.getAwsProxyRequest() != null && containerResponse.getAwsProxyRequest().getRequestSource() == AwsProxyRequest.RequestSource.ALB) {
+        if (containerResponse.getAwsProxyRequest() != null && containerResponse.getAwsProxyRequest().getRequestSource() == RequestSource.ALB) {
             awsProxyResponse.setStatusDescription(containerResponse.getStatus() + " " + Response.Status.fromStatusCode(containerResponse.getStatus()).getReasonPhrase());
         }
 

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

@@ -224,11 +224,11 @@ public AwsProxyRequestBuilder queryString(String key, String value) {
             this.request.setMultiValueQueryStringParameters(new MultiValuedTreeMap<>());
         }
 
-        if (request.getRequestSource() == AwsProxyRequest.RequestSource.API_GATEWAY) {
+        if (request.getRequestSource() == RequestSource.API_GATEWAY) {
             this.request.getMultiValueQueryStringParameters().add(key, value);
         }
         // ALB does not decode parameters automatically like API Gateway.
-        if (request.getRequestSource() == AwsProxyRequest.RequestSource.ALB) {
+        if (request.getRequestSource() == RequestSource.ALB) {
             try {
                 //if (URLDecoder.decode(value, ContainerConfig.DEFAULT_CONTENT_CHARSET).equals(value)) {
                 // TODO: Assume we are always given an unencoded value, smarter check here to encode
@@ -285,7 +285,7 @@ public AwsProxyRequestBuilder binaryBody(InputStream is)
 
 
     public AwsProxyRequestBuilder authorizerPrincipal(String principal) {
-        if (this.request.getRequestSource() == AwsProxyRequest.RequestSource.API_GATEWAY) {
+        if (this.request.getRequestSource() == RequestSource.API_GATEWAY) {
             if (this.request.getRequestContext().getAuthorizer() == null) {
                 this.request.getRequestContext().setAuthorizer(new ApiGatewayAuthorizerContext());
             }
@@ -295,7 +295,7 @@ public AwsProxyRequestBuilder authorizerPrincipal(String principal) {
             }
             this.request.getRequestContext().getAuthorizer().getClaims().setSubject(principal);
         }
-        if (this.request.getRequestSource() == AwsProxyRequest.RequestSource.ALB) {
+        if (this.request.getRequestSource() == RequestSource.ALB) {
             header("x-amzn-oidc-identity", principal);
             try {
                 header(

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

@@ -192,9 +192,4 @@ public boolean isBase64Encoded() {
     public void setIsBase64Encoded(boolean base64Encoded) {
         isBase64Encoded = base64Encoded;
     }
-
-    public static enum RequestSource {
-        ALB,
-        API_GATEWAY
-    }
 }

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

@@ -13,9 +13,12 @@
 package com.amazonaws.serverless.proxy.model;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.annotation.JsonIgnore;
 
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
 
 public class HttpApiV2ProxyRequest {
     private String version;
@@ -127,4 +130,12 @@ public HttpApiV2ProxyRequestContext getRequestContext() {
     public void setRequestContext(HttpApiV2ProxyRequestContext requestContext) {
         this.requestContext = requestContext;
     }
+
+    @JsonIgnore
+    public RequestSource getRequestSource() {
+        return Optional.ofNullable(getRequestContext())
+            .map(HttpApiV2ProxyRequestContext::getElb)
+            .map(albContext -> RequestSource.ALB)
+            .orElse(RequestSource.API_GATEWAY);
+    }
 }

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

@@ -25,6 +25,7 @@ public class HttpApiV2ProxyRequestContext {
     private String stage;
     private String time;
     private long timeEpoch;
+    private AlbContext elb;
 
     private HttpApiV2HttpContext http;
     private HttpApiV2AuthorizerMap authorizer;
@@ -117,4 +118,13 @@ public void setAuthorizer(HttpApiV2AuthorizerMap authorizer) {
         this.authorizer = authorizer;
     }
 
+    public AlbContext getElb() {
+        return this.elb;
+    }
+
+    public void setElb(AlbContext context) {
+        this.elb = context;
+    }
+
+    
 }

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

@@ -0,0 +1,18 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package com.amazonaws.serverless.proxy.model;
+
+public enum RequestSource {
+    ALB,
+    API_GATEWAY
+}

aws-serverless-java-container-core/src/test/java/com/amazonaws/serverless/proxy/model/HttpApiV2ProxyRequestTest.java

@@ -136,6 +136,7 @@ void deserialize_fromJsonString_authorizerPopulatedCorrectly() {
             HttpApiV2ProxyRequest req = LambdaContainerHandler.getObjectMapper().readValue(BASE_PROXY_REQUEST, HttpApiV2ProxyRequest.class);
             assertTrue(req.getRequestContext().getAuthorizer().getJwtAuthorizer().getClaims().containsKey("claim1"));
             assertEquals(2, req.getRequestContext().getAuthorizer().getJwtAuthorizer().getScopes().size());
+            assertEquals(RequestSource.API_GATEWAY, req.getRequestSource());
         } catch (JsonProcessingException e) {
             e.printStackTrace();
             fail("Exception while parsing request" + e.getMessage());
@@ -177,6 +178,7 @@ void deserialize_fromJsonString_isBase64EncodedPopulates() {
             assertFalse(req.isBase64Encoded());
             req = LambdaContainerHandler.getObjectMapper().readValue(NO_AUTH_PROXY, HttpApiV2ProxyRequest.class);
             assertTrue(req.isBase64Encoded());
+            assertEquals(RequestSource.API_GATEWAY, req.getRequestSource());
         } catch (JsonProcessingException e) {
             e.printStackTrace();
             fail("Exception while parsing request" + e.getMessage());