Add settings for websocket response mode

Author: yacchi

aws/adaptor.go

@@ -14,6 +14,11 @@ import (
 var DEBUGDumpPayload = os.Getenv("DEBUG_DUMP_PAYLOAD")
 var LambdaInvokeMode = os.Getenv("LAMBDA_INVOKE_MODE")
 
+// WebsocketResponseMode
+// * return - use lambda return value as response
+// * post_to_connection - use PostToConnection API to send response
+var WebsocketResponseMode = os.Getenv("WEBSOCKET_RESPONSE_MODE")
+
 type LambdaIntegrationType int
 
 const (

aws/lambda.go

@@ -166,7 +166,7 @@ func (l *LambdaHandler) InvokeWebsocketAPI(ctx context.Context, request *events.
 
 	routeKey := request.RequestContext.RouteKey
 
-	if routeKey == "$connect" || routeKey == "$disconnect" {
+	if routeKey == "$connect" || routeKey == "$disconnect" || WebsocketResponseMode == "return" {
 		w := NewResponseWriter()
 		l.httpHandler.ServeHTTP(w, req)
 		return RESTAPITargetResponse(w, multiValue)