Merge pull request #399 from deki/extensibility

Improve extensibility, allow other servlets besides DispatcherServlet

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

@@ -37,7 +37,7 @@
  * @param <ResponseType> The expected return type
  */
 public class SpringLambdaContainerHandler<RequestType, ResponseType> extends AwsLambdaServletContainerHandler<RequestType, ResponseType, HttpServletRequest, AwsHttpServletResponse> {
-    private ConfigurableWebApplicationContext appContext;
+    protected final ConfigurableWebApplicationContext appContext;
     private String[] profiles;
 
     // State vars
@@ -172,12 +172,20 @@ public void initialize()
             appContext.getEnvironment().setActiveProfiles(profiles);
         }
         appContext.setServletContext(getServletContext());
+        registerServlets();
+        // call initialize on AwsLambdaServletContainerHandler to initialize servlets that are set to load on startup
+        super.initialize();
+        Timer.stop("SPRING_COLD_START");
+    }
+
+    /**
+     * Overriding this method allows to customize the standard Spring DispatcherServlet
+     * or to register additional servlets
+     */
+    protected void registerServlets() {
         DispatcherServlet dispatcher = new DispatcherServlet(appContext);
         ServletRegistration.Dynamic reg = getServletContext().addServlet("dispatcherServlet", dispatcher);
         reg.addMapping("/");
         reg.setLoadOnStartup(1);
-        // call initialize on AwsLambdaServletContainerHandler to initialize servlets that are set to load on startup
-        super.initialize();
-        Timer.stop("SPRING_COLD_START");
     }
 }

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

@@ -13,17 +13,14 @@
 package com.amazonaws.serverless.proxy.spring;
 
 import com.amazonaws.serverless.exceptions.ContainerInitializationException;
-import com.amazonaws.serverless.proxy.internal.servlet.AwsProxyHttpServletRequest;
 import com.amazonaws.serverless.proxy.internal.servlet.ServletLambdaContainerHandlerBuilder;
-import com.amazonaws.serverless.proxy.model.AwsProxyRequest;
 import com.amazonaws.serverless.proxy.model.AwsProxyResponse;
-import org.springframework.web.WebApplicationInitializer;
 import org.springframework.web.context.ConfigurableWebApplicationContext;
 import org.springframework.web.context.support.AnnotationConfigWebApplicationContext;
 
 import javax.servlet.http.HttpServletRequest;
 
-public final class SpringProxyHandlerBuilder<RequestType> extends ServletLambdaContainerHandlerBuilder<
+public class SpringProxyHandlerBuilder<RequestType> extends ServletLambdaContainerHandlerBuilder<
             RequestType,
             AwsProxyResponse,
             HttpServletRequest,
@@ -69,22 +66,20 @@ public SpringLambdaContainerHandler<RequestType, AwsProxyResponse> build() throw
             }
         }
 
-        SpringLambdaContainerHandler<RequestType, AwsProxyResponse> handler =  new SpringLambdaContainerHandler<RequestType, AwsProxyResponse>(
-                requestTypeClass,
-                responseTypeClass,
-                requestReader,
-                responseWriter,
-                securityContextWriter,
-                exceptionHandler,
-                ctx,
-                initializationWrapper
-        );
+        SpringLambdaContainerHandler<RequestType, AwsProxyResponse> handler = createHandler(ctx);
         if (profiles != null) {
             handler.activateSpringProfiles(profiles);
         }
         return handler;
     }
 
+    protected SpringLambdaContainerHandler<RequestType, AwsProxyResponse> createHandler(ConfigurableWebApplicationContext ctx) {
+        return new SpringLambdaContainerHandler<>(
+                requestTypeClass, responseTypeClass, requestReader, responseWriter,
+                securityContextWriter, exceptionHandler, ctx, initializationWrapper
+        );
+    }
+
     @Override
     public SpringLambdaContainerHandler<RequestType, AwsProxyResponse> buildAndInitialize() throws ContainerInitializationException {
         SpringLambdaContainerHandler<RequestType, AwsProxyResponse> handler = build();

aws-serverless-java-container-spring/src/test/java/com/amazonaws/serverless/proxy/spring/extensibility/CustomServlet.java

@@ -0,0 +1,21 @@
+package com.amazonaws.serverless.proxy.spring.extensibility;
+
+import org.springframework.context.ApplicationContext;
+
+import javax.servlet.http.HttpServlet;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+import java.io.IOException;
+
+public class CustomServlet extends HttpServlet {
+    private ApplicationContext appCtx;
+
+    @Override
+    protected void doGet(HttpServletRequest req, HttpServletResponse resp) throws IOException {
+        resp.getWriter().print("Unittest " + (appCtx!=null ? appCtx.getDisplayName() : ""));
+    }
+
+    public void setAppCtx(ApplicationContext appCtx) {
+        this.appCtx = appCtx;
+    }
+}

aws-serverless-java-container-spring/src/test/java/com/amazonaws/serverless/proxy/spring/extensibility/CustomServletTest.java

@@ -0,0 +1,27 @@
+package com.amazonaws.serverless.proxy.spring.extensibility;
+
+import com.amazonaws.serverless.proxy.internal.testutils.AwsProxyRequestBuilder;
+import com.amazonaws.serverless.proxy.internal.testutils.MockLambdaContext;
+import org.junit.Test;
+
+import javax.ws.rs.HttpMethod;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+
+import static org.junit.Assert.assertTrue;
+
+public class CustomServletTest {
+
+    @Test
+    public void customServlet() throws IOException {
+        StreamLambdaHandler lambdaHandler = new StreamLambdaHandler();
+        InputStream requestStream = new AwsProxyRequestBuilder("/test", HttpMethod.GET)
+                .buildStream();
+        ByteArrayOutputStream responseStream = new ByteArrayOutputStream();
+        lambdaHandler.handleRequest(requestStream, responseStream, new MockLambdaContext());
+        assertTrue("response should contain value set in CustomServlet",
+                responseStream.toString().contains("Unittest"));
+    }
+
+}

aws-serverless-java-container-spring/src/test/java/com/amazonaws/serverless/proxy/spring/extensibility/CustomSpringLambdaContainerHandler.java

@@ -0,0 +1,92 @@
+package com.amazonaws.serverless.proxy.spring.extensibility;
+
+import com.amazonaws.serverless.exceptions.ContainerInitializationException;
+import com.amazonaws.serverless.proxy.*;
+import com.amazonaws.serverless.proxy.internal.servlet.AwsHttpServletResponse;
+import com.amazonaws.serverless.proxy.model.AwsProxyRequest;
+import com.amazonaws.serverless.proxy.model.AwsProxyResponse;
+import com.amazonaws.serverless.proxy.model.HttpApiV2ProxyRequest;
+import com.amazonaws.serverless.proxy.spring.SpringLambdaContainerHandler;
+import org.springframework.web.context.ConfigurableWebApplicationContext;
+
+import javax.servlet.ServletRegistration;
+import javax.servlet.http.HttpServletRequest;
+
+public class CustomSpringLambdaContainerHandler<RequestType, ResponseType> extends SpringLambdaContainerHandler<RequestType, ResponseType> {
+
+    /**
+     * Creates a default SpringLambdaContainerHandler initialized with the `AwsProxyRequest` and `AwsProxyResponse` objects
+     * @param config A set of classes annotated with the Spring @Configuration annotation
+     * @return An initialized instance of the `SpringLambdaContainerHandler`
+     * @throws ContainerInitializationException When the Spring framework fails to start.
+     */
+    public static SpringLambdaContainerHandler<AwsProxyRequest, AwsProxyResponse> getAwsProxyHandler(Class<?>... config) throws ContainerInitializationException {
+        return new CustomSpringProxyHandlerBuilder<AwsProxyRequest>()
+                .defaultProxy()
+                .initializationWrapper(new InitializationWrapper())
+                .configurationClasses(config)
+                .buildAndInitialize();
+    }
+
+    /**
+     * Creates a default SpringLambdaContainerHandler initialized with the `AwsProxyRequest` and `AwsProxyResponse` objects and sets the given profiles as active
+     * @param applicationContext A custom ConfigurableWebApplicationContext to be used
+     * @param profiles The spring profiles to activate
+     * @return An initialized instance of the `SpringLambdaContainerHandler`
+     * @throws ContainerInitializationException When the Spring framework fails to start.
+     */
+    public static SpringLambdaContainerHandler<AwsProxyRequest, AwsProxyResponse> getAwsProxyHandler(ConfigurableWebApplicationContext applicationContext, String... profiles)
+            throws ContainerInitializationException {
+        return new CustomSpringProxyHandlerBuilder<AwsProxyRequest>()
+                .defaultProxy()
+                .initializationWrapper(new InitializationWrapper())
+                .springApplicationContext(applicationContext)
+                .profiles(profiles)
+                .buildAndInitialize();
+    }
+
+    /**
+     * Creates a default SpringLambdaContainerHandler initialized with the `HttpApiV2ProxyRequest` and `AwsProxyResponse` objects
+     * @param config A set of classes annotated with the Spring @Configuration annotation
+     * @return An initialized instance of the `SpringLambdaContainerHandler`
+     * @throws ContainerInitializationException When the Spring framework fails to start.
+     */
+    public static SpringLambdaContainerHandler<HttpApiV2ProxyRequest, AwsProxyResponse> getHttpApiV2ProxyHandler(Class<?>... config) throws ContainerInitializationException {
+        return new CustomSpringProxyHandlerBuilder<HttpApiV2ProxyRequest>()
+                .defaultHttpApiV2Proxy()
+                .initializationWrapper(new InitializationWrapper())
+                .configurationClasses(config)
+                .buildAndInitialize();
+    }
+
+    /**
+     * Creates a new container handler with the given reader and writer objects
+     *
+     * @param requestTypeClass The class for the incoming Lambda event
+     * @param requestReader An implementation of `RequestReader`
+     * @param responseWriter An implementation of `ResponseWriter`
+     * @param securityContextWriter An implementation of `SecurityContextWriter`
+     * @param exceptionHandler An implementation of `ExceptionHandler`
+     */
+    public CustomSpringLambdaContainerHandler(Class<RequestType> requestTypeClass,
+                                              Class<ResponseType> responseTypeClass,
+                                              RequestReader<RequestType, HttpServletRequest> requestReader,
+                                              ResponseWriter<AwsHttpServletResponse, ResponseType> responseWriter,
+                                              SecurityContextWriter<RequestType> securityContextWriter,
+                                              ExceptionHandler<ResponseType> exceptionHandler,
+                                              ConfigurableWebApplicationContext applicationContext,
+                                              InitializationWrapper init) {
+        super(requestTypeClass, responseTypeClass, requestReader, responseWriter, securityContextWriter,
+                exceptionHandler, applicationContext, init);
+    }
+
+    @Override
+    protected void registerServlets() {
+        CustomServlet customServlet = new CustomServlet();
+        customServlet.setAppCtx(appContext);
+        ServletRegistration.Dynamic reg = getServletContext().addServlet("customServlet", customServlet);
+        reg.addMapping("/");
+        reg.setLoadOnStartup(1);
+    }
+
+}

aws-serverless-java-container-spring/src/test/java/com/amazonaws/serverless/proxy/spring/extensibility/CustomSpringProxyHandlerBuilder.java

@@ -0,0 +1,15 @@
+package com.amazonaws.serverless.proxy.spring.extensibility;
+
+import com.amazonaws.serverless.proxy.model.AwsProxyResponse;
+import com.amazonaws.serverless.proxy.spring.SpringLambdaContainerHandler;
+import com.amazonaws.serverless.proxy.spring.SpringProxyHandlerBuilder;
+import org.springframework.web.context.ConfigurableWebApplicationContext;
+
+public class CustomSpringProxyHandlerBuilder<RequestType> extends SpringProxyHandlerBuilder<RequestType> {
+
+    @Override
+    protected SpringLambdaContainerHandler<RequestType, AwsProxyResponse> createHandler(ConfigurableWebApplicationContext ctx) {
+        return new CustomSpringLambdaContainerHandler<>(requestTypeClass, responseTypeClass, requestReader, responseWriter,
+                securityContextWriter, exceptionHandler, ctx, initializationWrapper);
+    }
+}

aws-serverless-java-container-spring/src/test/java/com/amazonaws/serverless/proxy/spring/extensibility/StreamLambdaHandler.java

@@ -0,0 +1,32 @@
+package com.amazonaws.serverless.proxy.spring.extensibility;
+
+
+import com.amazonaws.serverless.exceptions.ContainerInitializationException;
+import com.amazonaws.serverless.proxy.model.AwsProxyRequest;
+import com.amazonaws.serverless.proxy.model.AwsProxyResponse;
+import com.amazonaws.serverless.proxy.spring.SpringLambdaContainerHandler;
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
+import org.springframework.web.context.support.GenericWebApplicationContext;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+
+
+public class StreamLambdaHandler implements RequestStreamHandler {
+    private static final SpringLambdaContainerHandler<AwsProxyRequest, AwsProxyResponse> handler;
+    static {
+        try {
+            handler = CustomSpringLambdaContainerHandler.getAwsProxyHandler(new GenericWebApplicationContext());
+        } catch (ContainerInitializationException e) {
+            throw new RuntimeException("Could not initialize Spring framework", e);
+        }
+    }
+
+    @Override
+    public void handleRequest(InputStream inputStream, OutputStream outputStream, Context context)
+            throws IOException {
+        handler.proxyStream(inputStream, outputStream, context);
+    }
+}