Add support websocket return response

Author: yacchi

aws/api_gateway_websocket.go

@@ -164,6 +164,14 @@ func WebsocketResponse(w *WebsocketResponseWriter, multiValue bool) (r *events.A
 		r.Headers = utils.SemicolonSeparatedHeaderMap(w.Header())
 	}
 
+	if 0 < w.buf.Len() {
+		if r.IsBase64Encoded {
+			r.Body = base64.StdEncoding.EncodeToString(w.buf.Bytes())
+		} else {
+			r.Body = w.buf.String()
+		}
+	}
+
 	w.Done()
 	return
 }

aws/websocket_response_writer.go

@@ -1,6 +1,7 @@
 package aws
 
 import (
+	"bytes"
 	"context"
 	"github.com/aws/aws-lambda-go/events"
 	"net/http"
@@ -14,6 +15,7 @@ type WebsocketResponseWriter struct {
 	headers     http.Header
 	wroteHeader bool
 	closeCh     chan bool
+	buf         *bytes.Buffer
 }
 
 func NewWebsocketResponseWriter(ctx context.Context, client APIGatewayManagementAPI, request *events.APIGatewayWebsocketProxyRequest) *WebsocketResponseWriter {
@@ -23,6 +25,7 @@ func NewWebsocketResponseWriter(ctx context.Context, client APIGatewayManagement
 		req:     request,
 		headers: map[string][]string{},
 		closeCh: make(chan bool, 1),
+		buf:     bytes.NewBuffer(nil),
 	}
 }
 
@@ -35,6 +38,8 @@ func (w *WebsocketResponseWriter) Write(i []byte) (int, error) {
 		w.WriteHeader(http.StatusOK)
 	}
 
+	w.buf.Write(i)
+
 	err := w.client.PostToConnection(w.ctx, w.req.RequestContext.ConnectionID, i)
 	if err != nil {
 		return 0, err