Merge pull request #910 from mbfreder/fix-alb-hostname

fix: Error with version 2.0.2 when multi-value header is active in alb
aws · Jun 19, 2024 · 7df3166 · 7df3166
2 parents 58084d9 + 9463113
commit 7df3166
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 12 deletions.
diff --git a/...main/java/com/amazonaws/serverless/proxy/internal/servlet/AwsProxyHttpServletRequest.java b/...main/java/com/amazonaws/serverless/proxy/internal/servlet/AwsProxyHttpServletRequest.java
@@ -437,23 +437,27 @@ public String getRemoteAddr() {
         if (request.getRequestContext() == null || request.getRequestContext().getIdentity() == null) {
             return "127.0.0.1";
         }
-        if (request.getRequestContext().getElb() != null) {
-            return request.getHeaders().get(CLIENT_IP_HEADER);
+        if (request.getRequestSource().equals(RequestSource.ALB)) {
+            return Objects.nonNull(request.getHeaders()) ?
+                    request.getHeaders().get(CLIENT_IP_HEADER) :
+                    request.getMultiValueHeaders().getFirst(CLIENT_IP_HEADER);
         }
         return request.getRequestContext().getIdentity().getSourceIp();
     }
 
 
     @Override
     public String getRemoteHost() {
-        if (Objects.nonNull(request.getRequestContext().getElb())) {
-            String hostHeader = request.getHeaders().get(HttpHeaders.HOST);
-
-            // the host header has the form host:port, so we split the string to get the host part
-            return Arrays.asList(hostHeader.split(":")).get(0);
+        String hostHeader;
+        if (request.getRequestSource().equals(RequestSource.ALB)) {
+            hostHeader = Objects.nonNull(request.getHeaders()) ?
+                    request.getHeaders().get(HttpHeaders.HOST) :
+                    request.getMultiValueHeaders().getFirst(HttpHeaders.HOST);
+        } else {
+            hostHeader = request.getMultiValueHeaders().getFirst(HttpHeaders.HOST);
         }
-
-        return request.getMultiValueHeaders().getFirst(HttpHeaders.HOST);
+        // the host header has the form host:port, so we split the string to get the host part
+        return Arrays.asList(hostHeader.split(":")).get(0);
     }
 
 
@@ -483,8 +487,11 @@ public RequestDispatcher getRequestDispatcher(String s) {
 
     @Override
     public int getRemotePort() {
-        if (Objects.nonNull(request.getRequestContext().getElb())) {
-            String portHeader = request.getHeaders().get(PORT_HEADER_NAME);
+        if (request.getRequestSource().equals(RequestSource.ALB)) {
+            String portHeader;
+            portHeader = Objects.nonNull(request.getHeaders()) ?
+                    request.getHeaders().get(PORT_HEADER_NAME) :
+                    request.getMultiValueHeaders().getFirst(PORT_HEADER_NAME);
             if (Objects.nonNull(portHeader)) {
                 return Integer.parseInt(portHeader);
             }

diff --git a/.../java/com/amazonaws/serverless/proxy/internal/servlet/AwsProxyHttpServletRequestTest.java b/.../java/com/amazonaws/serverless/proxy/internal/servlet/AwsProxyHttpServletRequestTest.java
@@ -652,17 +652,31 @@ void serverName_albHostHeader_returnsHostHeader() {
     }
 
     @Test
-    void getRemoteHost_albHostHeader_returnsHostHeader() {
+    void getRemoteHost_albHostHeader_singleValue_returnsHostHeader() {
         initAwsProxyHttpServletRequestTest("ALB");
         AwsProxyRequest proxyReq = new AwsProxyRequestBuilder("/test", "GET")
                 .alb().build();
+        proxyReq.setMultiValueHeaders(null);
         proxyReq.getHeaders().put(HttpHeaders.HOST, "testapi.us-east-1.elb.amazonaws.com");
         HttpServletRequest servletRequest = new AwsProxyHttpServletRequest(proxyReq, null, null);
 
         String host = servletRequest.getRemoteHost();
         assertEquals("testapi.us-east-1.elb.amazonaws.com", host);
     }
 
+    @Test
+    void getRemoteHost_albHostHeader_multiValue_returnsHostHeader() {
+        initAwsProxyHttpServletRequestTest("ALB");
+        AwsProxyRequest proxyReq = new AwsProxyRequestBuilder("/test", "GET")
+                .header(HttpHeaders.HOST, "testapi.us-east-1.elb.amazonaws.com")
+                .alb().build();
+        proxyReq.setHeaders(null);
+        HttpServletRequest servletRequest = new AwsProxyHttpServletRequest(proxyReq, null, null);
+
+        String host = servletRequest.getRemoteHost();
+        assertEquals("testapi.us-east-1.elb.amazonaws.com", host);
+    }
+
     private AwsProxyRequestBuilder getRequestWithHeaders() {
         return new AwsProxyRequestBuilder("/hello", "GET")
                 .header(CUSTOM_HEADER_KEY, CUSTOM_HEADER_VALUE)