Fix for async context and setting response in proxy servlet request before processing

Author: sapessi

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

@@ -11,6 +11,7 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * Async context for Serverless Java Container. This is used to support reactive embedded servers for our support for
@@ -23,6 +24,8 @@ public class AwsAsyncContext implements AsyncContext {
     private AwsLambdaServletContainerHandler handler;
     private List<AsyncListenerHolder> listeners;
     private long timeout;
+    private AtomicBoolean dispatched;
+    private AtomicBoolean completed;
 
     private Logger log = LoggerFactory.getLogger(AwsAsyncContext.class);
 
@@ -33,6 +36,8 @@ public AwsAsyncContext(HttpServletRequest request, HttpServletResponse response,
         handler = servletHandler;
         listeners = new ArrayList<>();
         timeout = 3000;
+        dispatched = new AtomicBoolean(false);
+        completed = new AtomicBoolean(false);
     }
 
     @Override
@@ -57,6 +62,7 @@ public void dispatch() {
             notifyListeners(NotificationType.START_ASYNC, null);
             Servlet servlet = ((AwsServletContext)handler.getServletContext()).getServletForPath(req.getPathInfo());
             handler.doFilter(req, res, servlet);
+            dispatched.set(true);
         } catch (ServletException | IOException e) {
             notifyListeners(NotificationType.ERROR, e);
         }
@@ -83,6 +89,7 @@ public void complete() {
             log.debug("Completing request");
             notifyListeners(NotificationType.COMPLETE, null);
             res.flushBuffer();
+            completed.set(true);
         } catch (IOException e) {
             log.error("Could not flush response buffer", e);
             throw new RuntimeException(e);
@@ -126,6 +133,14 @@ public long getTimeout() {
         return timeout;
     }
 
+    public boolean isDispatched() {
+        return dispatched.get();
+    }
+
+    public boolean isCompleted() {
+        return completed.get();
+    }
+
     private void notifyListeners(NotificationType type, Throwable t) {
         listeners.forEach((h) -> {
             try {

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

@@ -76,7 +76,7 @@ public class AwsProxyHttpServletRequest extends AwsHttpServletRequest {
 
     private AwsProxyRequest request;
     private SecurityContext securityContext;
-    private AsyncContext asyncContext;
+    private AwsAsyncContext asyncContext;
     private Map<String, List<String>> urlEncodedFormParameters;
     private Map<String, Part> multipartFormParameters;
     private static Logger log = LoggerFactory.getLogger(AwsProxyHttpServletRequest.class);
@@ -670,7 +670,13 @@ public boolean isAsyncSupported() {
 
     @Override
     public boolean isAsyncStarted() {
-        return asyncContext != null;
+        if (asyncContext == null) {
+            return false;
+        }
+        if (asyncContext.isCompleted() || asyncContext.isDispatched()) {
+            return false;
+        }
+        return true;
     }
 
 

aws-serverless-java-container-spring/src/main/java/com/amazonaws/serverless/proxy/spring/SpringBootLambdaContainerHandler.java

@@ -142,6 +142,7 @@ protected void handleRequest(AwsProxyHttpServletRequest containerRequest, AwsHtt
 
         // process filters & invoke servlet
         Servlet reqServlet = ((AwsServletContext)getServletContext()).getServletForPath(containerRequest.getPathInfo());
+        containerRequest.setResponse(containerResponse);
         doFilter(containerRequest, containerResponse, reqServlet);
         Timer.stop("SPRINGBOOT_HANDLE_REQUEST");
     }

aws-serverless-java-container-spring/src/main/java/com/amazonaws/serverless/proxy/spring/SpringLambdaContainerHandler.java

@@ -143,6 +143,7 @@ protected void handleRequest(AwsProxyHttpServletRequest containerRequest, AwsHtt
 
         // process filters
         Servlet reqServlet = ((AwsServletContext)getServletContext()).getServletForPath(containerRequest.getPathInfo());
+        containerRequest.setResponse(containerResponse);
         doFilter(containerRequest, containerResponse, reqServlet);
         Timer.stop("SPRING_HANDLE_REQUEST");
     }

aws-serverless-java-container-springboot2/src/main/java/com/amazonaws/serverless/proxy/spring/SpringBootLambdaContainerHandler.java

@@ -138,6 +138,7 @@ protected void handleRequest(AwsProxyHttpServletRequest containerRequest, AwsHtt
 
         // process filters & invoke servlet
         Servlet reqServlet = ((AwsServletContext)getServletContext()).getServletForPath(containerRequest.getPathInfo());
+        containerRequest.setResponse(containerResponse);
         doFilter(containerRequest, containerResponse, reqServlet);
         Timer.stop("SPRINGBOOT2_HANDLE_REQUEST");
     }