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

@@ -171,12 +171,18 @@ func (l *LambdaHandler) InvokeWebsocketAPI(ctx context.Context, request *events.
 		l.httpHandler.ServeHTTP(w, req)
 		return RESTAPITargetResponse(w, multiValue)
 	} else {
-		if apiGW, err := l.ProvideAPIGatewayClient(ctx, request); err != nil {
-			return nil, err
+		if WebsocketResponseMode == "post_to_connection" {
+			if apiGW, err := l.ProvideAPIGatewayClient(ctx, request); err != nil {
+				return nil, err
+			} else {
+				w := NewWebsocketResponseWriter(ctx, apiGW, request)
+				l.httpHandler.ServeHTTP(w, req)
+				return WebsocketResponse(w, multiValue)
+			}
 		} else {
-			w := NewWebsocketResponseWriter(ctx, apiGW, request)
+			w := NewResponseWriter()
 			l.httpHandler.ServeHTTP(w, req)
-			return WebsocketResponse(w, multiValue)
+			return RESTAPITargetResponse(w, multiValue)
 		}
 	}
 }