From eb9d898d814f558bd54f22c77a674b5915af6e56 Mon Sep 17 00:00:00 2001
From: Koji Saiki <saikou9901@gmail.com>
Date: Wed, 3 Jan 2024 10:22:03 +0900
Subject: [PATCH] Try routing refactoring

---
 app/router/router.go      | 41 ++++++++++++++++++++++++++++++++++++---
 app/router/router_test.go | 24 +++++++++++++++++++++++
 2 files changed, 62 insertions(+), 3 deletions(-)

diff --git a/app/router/router.go b/app/router/router.go
index 699b6ab0..91273adc 100644
--- a/app/router/router.go
+++ b/app/router/router.go
@@ -3,6 +3,7 @@ package router
 import (
 	"io"
 	"net/http"
+	"strings"
 
 	log "github.com/sirupsen/logrus"
 
@@ -65,14 +66,15 @@ func health(w http.ResponseWriter, req *http.Request) {
 }
 
 func actionHandler(w http.ResponseWriter, req *http.Request) {
+	action := extractAction(req)
 	log.WithFields(
 		log.Fields{
-			"action": req.FormValue("Action"),
+			"action": action,
 			"url":    req.URL,
 		}).Debug("Handling URL request")
-	fn, ok := routingTable[req.FormValue("Action")]
+	fn, ok := routingTable[action]
 	if !ok {
-		log.Println("Bad Request - Action:", req.FormValue("Action"))
+		log.Println("Bad Request - Action:", action)
 		w.WriteHeader(http.StatusBadRequest)
 		io.WriteString(w, "Bad Request")
 		return
@@ -85,3 +87,36 @@ func pemHandler(w http.ResponseWriter, req *http.Request) {
 	w.WriteHeader(http.StatusOK)
 	w.Write(sns.PemKEY)
 }
+
+type AwsProtocol int
+
+const (
+	AwsJsonProtocol  AwsProtocol = iota
+	AwsQueryProtocol AwsProtocol = iota
+)
+
+// Extract target Action from the request.
+// How contains the Action name is different with aws-query protocol and aws-json protocol.
+func extractAction(req *http.Request) string {
+	protocol := resolveProtocol(req)
+	switch protocol {
+	case AwsJsonProtocol:
+		// Get action from X-Amz-Target header
+		action := req.Header.Get("X-Amz-Target")
+		// Action value will be like as "AmazonSQS.CreateQueue".
+		// After dot should be the action name.
+		return strings.Split(action, ".")[1]
+	case AwsQueryProtocol:
+		return req.FormValue("Action")
+	}
+	return ""
+}
+
+// Determine which protocol is used.
+func resolveProtocol(req *http.Request) AwsProtocol {
+	// Use content-type to determine protocol
+	if req.Header.Get("Content-Type") == "application/x-amz-json-1.0" {
+		return AwsJsonProtocol
+	}
+	return AwsQueryProtocol
+}
diff --git a/app/router/router_test.go b/app/router/router_test.go
index 41288d70..b255073a 100644
--- a/app/router/router_test.go
+++ b/app/router/router_test.go
@@ -1,6 +1,8 @@
 package router
 
 import (
+	"bytes"
+	"encoding/json"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
@@ -91,6 +93,28 @@ func TestIndexServerhandler_POST_GoodRequest_With_URL(t *testing.T) {
 	}
 }
 
+func TestIndexServerhandler_POST_GoodRequest_With_URL_And_Aws_Json_Protocol(t *testing.T) {
+	json, _ := json.Marshal(map[string]string{
+		"QueueName": "local-queue1",
+	})
+	req, err := http.NewRequest("POST", "/100010001000/local-queue1", bytes.NewBuffer(json))
+	if err != nil {
+		t.Fatal(err)
+	}
+	req.Header.Set("X-Amz-Target", "AmazonSQS.CreateQueue")
+	req.Header.Set("Content-Type", "application/x-amz-json-1.0")
+
+	rr := httptest.NewRecorder()
+
+	New().ServeHTTP(rr, req)
+
+	// Check the status code is what we expect.
+	if status := rr.Code; status != http.StatusOK {
+		t.Errorf("handler returned wrong status code: got %v want %v",
+			status, http.StatusOK)
+	}
+}
+
 func TestIndexServerhandler_GET_GoodRequest_Pem_cert(t *testing.T) {
 
 	req, err := http.NewRequest("GET", "/SimpleNotificationService/100010001000.pem", nil)