From 365993b479896efbe0404b36b55b3a6e2933fa16 Mon Sep 17 00:00:00 2001
From: Chris Broglie <cbroglie@zynga.com>
Date: Mon, 9 Feb 2015 22:15:48 -0800
Subject: [PATCH] Add sign handler which adds Authorization header

This fixes #94
---
 aws/signer/v4/v4util/v4util.go      | 34 +++++++++++++++
 aws/signer/v4/v4util/v4util_test.go | 65 +++++++++++++++++++++++++++++
 2 files changed, 99 insertions(+)
 create mode 100644 aws/signer/v4/v4util/v4util.go
 create mode 100644 aws/signer/v4/v4util/v4util_test.go

diff --git a/aws/signer/v4/v4util/v4util.go b/aws/signer/v4/v4util/v4util.go
new file mode 100644
index 00000000000..e89dd7afc97
--- /dev/null
+++ b/aws/signer/v4/v4util/v4util.go
@@ -0,0 +1,34 @@
+package v4util
+
+import (
+	"fmt"
+
+	"github.com/awslabs/aws-sdk-go/aws"
+)
+
+// Add a Authorization header to requests using the v4 signature format.
+//
+// It is required that the request has previously been signed using the query
+// string. This is useful for non-production services like DynamoDB Local which
+// do not properly support signing using the query string.
+func SignWithHeader(req *aws.Request) {
+	headers := map[string]string{
+		"X-Amz-Algorithm":     "",
+		"X-Amz-Credential":    "",
+		"X-Amz-SignedHeaders": "",
+		"X-Amz-Signature":     "",
+	}
+	query := req.HTTPRequest.URL.Query()
+	for header, _ := range headers {
+		v := query.Get(header)
+		if v == "" {
+			req.Error = fmt.Errorf("'%s' was not found in the query string", header)
+			return
+		}
+		headers[header] = v
+	}
+	authorization := fmt.Sprintf("%s Credential=%s, SignedHeaders=%s, Signature=%s",
+		headers["X-Amz-Algorithm"], headers["X-Amz-Credential"],
+		headers["X-Amz-SignedHeaders"], headers["X-Amz-Signature"])
+	req.HTTPRequest.Header.Set("Authorization", authorization)
+}
diff --git a/aws/signer/v4/v4util/v4util_test.go b/aws/signer/v4/v4util/v4util_test.go
new file mode 100644
index 00000000000..2445038a510
--- /dev/null
+++ b/aws/signer/v4/v4util/v4util_test.go
@@ -0,0 +1,65 @@
+package v4util
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/awslabs/aws-sdk-go/aws"
+	"github.com/awslabs/aws-sdk-go/service/dynamodb"
+)
+
+func TestSignWithHeader(t *testing.T) {
+	server := setupServer()
+	defer server.Close()
+
+	creds := aws.Creds("DUMMY_KEY", "DUMMY_SECRET", "")
+	config := &dynamodb.DynamoDBConfig{
+		Config: &aws.Config{
+			Credentials: creds,
+			Endpoint:    server.URL,
+		},
+	}
+	db := dynamodb.New(config)
+	db.Handlers.Sign.PushBack(SignWithHeader)
+
+	_, err := db.ListTables(&dynamodb.ListTablesInput{})
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestSignWithHeaderFailed(t *testing.T) {
+	server := setupServer()
+	defer server.Close()
+
+	creds := aws.Creds("DUMMY_KEY", "DUMMY_SECRET", "")
+	config := &dynamodb.DynamoDBConfig{
+		Config: &aws.Config{
+			Credentials: creds,
+			Endpoint:    server.URL,
+		},
+	}
+	db := dynamodb.New(config)
+	db.Handlers.Sign.Init() // remove the v4.Sign handler
+	db.Handlers.Sign.PushBack(SignWithHeader)
+
+	_, err := db.ListTables(&dynamodb.ListTablesInput{})
+	if err == nil {
+		t.Fatalf("Expected failure")
+	}
+}
+
+// setupServer creates test server which simply validates the presence of the
+// Authorization header.
+func setupServer() *httptest.Server {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		authorization := r.Header.Get("Authorization")
+		if authorization == "" {
+			w.WriteHeader(http.StatusBadRequest)
+			return
+		}
+		w.WriteHeader(http.StatusOK)
+	}))
+	return server
+}
\ No newline at end of file