Memoizing SigV4

awssigner.go

package main

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"log"
	"net/http"
	"sort"
	"strings"
	"time"
)

const (
	Debug = false
)

type AWSSigner struct {
	Host      string
	Region    string
	Service   string
	AccessKey string
	SecretKey string
}

func NewAWSSigner(host, region, service, accessKey, secretKey string) *AWSSigner {
	return &AWSSigner{
		Host:      host,
		Region:    region,
		Service:   service,
		AccessKey: accessKey,
		SecretKey: secretKey,
	}
}

func (s *AWSSigner) SendPOSTRequest(path string, body []byte) (*http.Response, error) {
	req, err := http.NewRequest("POST", "https://"+s.Host+path, bytes.NewReader(body))

	if err != nil {
		return nil, err
	}
	req.Header.Set("Host", s.Host)
	req.Header.Set("Content-Type", "application/json")

	if Debug == true {
		fmt.Println(req)
		fmt.Println("req.ContentLength", req.ContentLength)
	}

	s.signRequest(req, body)

	// HTTP Optimizations
	t := http.DefaultTransport.(*http.Transport).Clone()
	t.MaxIdleConns = 100
	t.MaxConnsPerHost = 100
	t.MaxIdleConnsPerHost = 100

	client := &http.Client{
		Timeout:   time.Second * 10,
		Transport: t,
	}
	return client.Do(req)
}

func (s *AWSSigner) signRequest(req *http.Request, body []byte) {
	t := time.Now().UTC()
	canonicalRequest := s.buildCanonicalRequest(req, body)
	stringToSign := s.buildStringToSign(t, canonicalRequest)
	signature := s.calculateSignature(t, stringToSign)

	if Debug == true {
		fmt.Println("Canonical request\n====================\n" + canonicalRequest + "\n====================\n\n")
		fmt.Println("String to sign\n====================\n" + stringToSign + "\n====================\n\n")
		fmt.Println("Signature\n====================\n" + signature + "\n====================\n\n")
	}

	authHeader := fmt.Sprintf("AWS4-HMAC-SHA256 Credential=%s/%s/%s/%s/aws4_request, SignedHeaders=%s, Signature=%s",
		s.AccessKey, t.Format("20060102"), s.Region, s.Service, s.signedHeaders(req), signature)

	req.Header.Set("Authorization", authHeader)
	req.Header.Set("X-Amz-Date", t.Format("20060102T150405Z"))
}

func (s *AWSSigner) buildCanonicalRequest(req *http.Request, body []byte) string {
	hashedBody := sha256.Sum256(body)
	canonicalBody := hex.EncodeToString(hashedBody[:])

	return strings.Join([]string{
		req.Method,
		req.URL.Path,
		req.URL.RawQuery,
		s.canonicalHeaders(req) + "\n",
		s.signedHeaders(req),
		canonicalBody,
	}, "\n")
}

func (s *AWSSigner) canonicalHeaders(req *http.Request) string {
	var headers []string
	for k, v := range req.Header {
		header := strings.ToLower(k) + ":" + strings.Join(v, ",")
		headers = append(headers, header)
	}

	sort.Strings(headers)
	return strings.Join(headers, "\n")
}

func (s *AWSSigner) signedHeaders(req *http.Request) string {
	var headers []string
	for k := range req.Header {
		headers = append(headers, strings.ToLower(k))
	}

	sort.Strings(headers)
	return strings.Join(headers, ";")
}

func (s *AWSSigner) buildStringToSign(t time.Time, canonicalRequest string) string {
	hashedRequest := sha256.Sum256([]byte(canonicalRequest))
	hashedRequestHex := hex.EncodeToString(hashedRequest[:])

	return strings.Join([]string{
		"AWS4-HMAC-SHA256",
		t.Format("20060102T150405Z"),
		s.credentialScope(t),
		hashedRequestHex,
	}, "\n")
}

func (s *AWSSigner) credentialScope(t time.Time) string {
	return strings.Join([]string{
		t.Format("20060102"),
		s.Region,
		s.Service,
		"aws4_request",
	}, "/")
}

type signingKeyCacheKey struct {
	date    string
	region  string
	service string
}

func (s *AWSSigner) calculateSignature(t time.Time, stringToSign string) string {
	cacheKey := signingKeyCacheKey{
		date:    t.Format("20060102"),
		region:  s.Region,
		service: s.Service,
	}

	kSigning := s.getSigningKey(cacheKey)

	signature := s.hmacSHA256(kSigning, []byte(stringToSign))
	return hex.EncodeToString(signature)
}

var signingKeyCache = make(map[signingKeyCacheKey][]byte)

func (s *AWSSigner) getSigningKey(cacheKey signingKeyCacheKey) []byte {
	if kSigning, ok := signingKeyCache[cacheKey]; ok {
		println("CACHE HIT")
		return kSigning
	}
	println("CACHE MISS")
	kSecret := []byte("AWS4" + s.SecretKey)
	kDate := s.hmacSHA256(kSecret, []byte(cacheKey.date))
	kRegion := s.hmacSHA256(kDate, []byte(cacheKey.region))
	kService := s.hmacSHA256(kRegion, []byte(cacheKey.service))
	kSigning := s.hmacSHA256(kService, []byte("aws4_request"))

	signingKeyCache[cacheKey] = kSigning
	return kSigning
}

func (s *AWSSigner) hmacSHA256(key, data []byte) []byte {
	hash := hmac.New(sha256.New, key)
	hash.Write(data)
	return hash.Sum(nil)
}

func do_request(host string, uri string, service string, body []byte) (string, error) {

	signer := NewAWSSigner(host, AWSRegion, service, AWSAccessKey, AWSSecretKey)

	resp, err := signer.SendPOSTRequest(uri, body)
	if err != nil {
		log.Fatal(err)
	}
	defer resp.Body.Close()

	responseBody, err := io.ReadAll(resp.Body)
	if err != nil {
		log.Fatal(err)
	}

	// fmt.Println("Response:", string(responseBody))
	return string(responseBody), nil
}