API Gateway data wrapper added, triggermesh/ghbackend#104

Author: tzununbekov

main.go

@@ -25,6 +25,8 @@ import (
 	"strings"
 	"sync"
 	"time"
+
+	"github.com/triggermesh/aws-custom-runtime/pkg/events/apiGateway"
 )
 
 const (
@@ -39,6 +41,11 @@ type message struct {
 	data     []byte
 }
 
+type responseWrapper struct {
+	http.ResponseWriter
+	Body []byte
+}
+
 var (
 	tasks   chan message
 	results map[string]chan message
@@ -60,6 +67,11 @@ var (
 	}
 )
 
+func (rw *responseWrapper) Write(data []byte) (int, error) {
+	rw.Body = data
+	return len(data), nil
+}
+
 func setupEnv() error {
 	environment["_HANDLER"], _ = os.LookupEnv("_HANDLER")
 	environment["LAMBDA_TASK_ROOT"], _ = os.LookupEnv("LAMBDA_TASK_ROOT")
@@ -187,6 +199,21 @@ func responseHandler(w http.ResponseWriter, r *http.Request) {
 	return
 }
 
+func mapEvent(h http.Handler) http.Handler {
+	eventType, _ := os.LookupEnv("EVENT")
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		rw := responseWrapper{w, []byte{}}
+		switch eventType {
+		case "API_GATEWAY":
+			apiGateway.Request(r)
+			h.ServeHTTP(&rw, r)
+			apiGateway.Response(w, rw.Body)
+		default:
+			h.ServeHTTP(w, r)
+		}
+	})
+}
+
 func api() error {
 	apiRouter := http.NewServeMux()
 	apiRouter.HandleFunc(awsEndpoint+"/init/error", initError)
@@ -220,7 +247,8 @@ func main() {
 	}
 
 	taskRouter := http.NewServeMux()
-	taskRouter.HandleFunc("/", newTask)
+	taskHandler := http.HandlerFunc(newTask)
+	taskRouter.Handle("/", mapEvent(taskHandler))
 	fmt.Println("Listening...")
 	log.Fatalln(http.ListenAndServe(":8080", taskRouter))
 }

pkg/events/apiGateway/mapper.go

@@ -0,0 +1,53 @@
+package apiGateway
+
+import (
+	"bytes"
+	"encoding/json"
+	"io/ioutil"
+	"log"
+	"net/http"
+	"strings"
+
+	"github.com/aws/aws-lambda-go/events"
+)
+
+func Request(r *http.Request) {
+	event := events.APIGatewayProxyRequest{}
+	body, err := ioutil.ReadAll(r.Body)
+	if err != nil {
+		log.Fatalln(err)
+	}
+
+	event.Body = string(body)
+	event.Headers = make(map[string]string)
+	for k, v := range r.Header {
+		event.Headers[k] = strings.Join(v, ",")
+	}
+	event.HTTPMethod = r.Method
+	event.Path = r.URL.Path
+	event.QueryStringParameters = make(map[string]string)
+	for k, v := range r.URL.Query() {
+		event.QueryStringParameters[k] = strings.Join(v, ",")
+	}
+	event.RequestContext = events.APIGatewayProxyRequestContext{}
+	// event.Resource = ""
+	// event.QueryStringParameters = make(map[string]string)
+	// event.IsBase64Encoded = false
+	js, err := json.Marshal(event)
+	if err != nil {
+		log.Fatalln(err)
+	}
+	r.Body = ioutil.NopCloser(bytes.NewBuffer(js))
+}
+
+func Response(w http.ResponseWriter, data []byte) (int, error) {
+	var js events.APIGatewayProxyResponse
+	if err := json.Unmarshal(data, &js); err != nil {
+		return 0, err
+	}
+	for k, v := range js.Headers {
+		w.Header().Set(k, v)
+	}
+	w.WriteHeader(js.StatusCode)
+	return w.Write([]byte(js.Body))
+}