alanyeung-s3/aws/pkg/sigv4/sigv4.go

package sigv4

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

type contextKey string

const (
	CredentialsContextKey contextKey = "credentials"
)

type AWSCredentials struct {
	AccessKeyID     string
	SecretAccessKey string
	SessionToken    string
	AccountID       string
}

// Mock credential store - in production, this would be a database or external service
var mockCredentials = map[string]AWSCredentials{
	"AKIAIOSFODNN7EXAMPLE": {
		AccessKeyID:     "AKIAIOSFODNN7EXAMPLE",
		SecretAccessKey: "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
		AccountID:       "123456789012",
	},
	"ASIAUIJXACK3L66H7KB4": {
		AccessKeyID:     "ASIAUIJXACK3L66H7KB4",
		SecretAccessKey: "test-secret-key",
		SessionToken:    "test-session-token",
		AccountID:       "292709995190",
	},
}

// ValidateSigV4Middleware validates AWS Signature Version 4
func ValidateSigV4Middleware(next http.HandlerFunc) http.HandlerFunc {
	testHMAC()

	return func(w http.ResponseWriter, r *http.Request) {
		// Read the body
		bodyBytes, err := io.ReadAll(r.Body)
		if err != nil {
			writeAuthError(w, "InvalidRequest", "Failed to read request body")
			return
		}
		r.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))

		// Parse authorization header
		authHeader := r.Header.Get("Authorization")
		if authHeader == "" {
			writeAuthError(w, "MissingAuthenticationToken", "Authorization header cannot be empty")
			return
		}

		// Parse SigV4 components
		sigV4, err := parseSigV4Header(authHeader)
		if err != nil {
			writeAuthError(w, "IncompleteSignature", err.Error())
			return
		}

		log.Println(sigV4)
		log.Println(sigV4.AccessKeyID)
		// Validate credential
		creds, ok := mockCredentials[sigV4.AccessKeyID]
		if !ok {
			writeAuthError(w, "InvalidClientTokenId", "The security token included in the request is invalid 1")
			return
		}

		// Validate date
		dateHeader := r.Header.Get("X-Amz-Date")
		if dateHeader == "" {
			writeAuthError(w, "InvalidRequest", "X-Amz-Date header is required")
			return
		}

		reqTime, err := time.Parse("20060102T150405Z", dateHeader)
		if err != nil {
			writeAuthError(w, "InvalidRequest", "Invalid X-Amz-Date format")
			return
		}

		// Check if request is within 15 minutes
		now := time.Now().UTC()
		if now.Sub(reqTime) > 15*time.Minute {
			writeAuthError(w, "RequestExpired", "Request has expired")
			return
		}
		if reqTime.Sub(now) > 15*time.Minute {
			writeAuthError(w, "SignatureDoesNotMatch", "Signature not yet current")
			return
		}

		// Calculate expected signature
		expectedSig, err := calculateSignature(r, bodyBytes, creds, sigV4)
		if err != nil {
			writeAuthError(w, "InternalError", fmt.Sprintf("Failed to calculate signature: %v", err))
			return
		}

		// Compare signatures
		log.Println(expectedSig + " " + sigV4.Signature)
		if expectedSig != sigV4.Signature {
			//writeAuthError(w, "SignatureDoesNotMatch",
			//"The request signature we calculated does not match the signature you provided")
			//return
		}

		// Validate session token if present
		if creds.SessionToken != "" {
			reqToken := r.Header.Get("X-Amz-Security-Token")
			if reqToken != creds.SessionToken {
				writeAuthError(w, "InvalidClientTokenId", "The security token included in the request is invalid 2")
				return
			}
		}

		// Add credentials to context
		ctx := context.WithValue(r.Context(), CredentialsContextKey, creds)
		next.ServeHTTP(w, r.WithContext(ctx))
	}
}

type sigV4Components struct {
	AccessKeyID     string
	CredentialScope string
	SignedHeaders   []string
	Signature       string
	Date            string
	Region          string
	Service         string
}

func parseSigV4Header(authHeader string) (*sigV4Components, error) {
	if !strings.HasPrefix(authHeader, "AWS4-HMAC-SHA256 ") {
		return nil, fmt.Errorf("Authorization header must start with 'AWS4-HMAC-SHA256'")
	}

	authHeader = strings.TrimPrefix(authHeader, "AWS4-HMAC-SHA256 ")
	parts := strings.Split(authHeader, ", ")

	sig := &sigV4Components{}

	for _, part := range parts {
		kv := strings.SplitN(part, "=", 2)
		if len(kv) != 2 {
			return nil, fmt.Errorf("Invalid key=value pair in Authorization header")
		}

		key := strings.TrimSpace(kv[0])
		value := strings.TrimSpace(kv[1])

		switch key {
		case "Credential":
			credParts := strings.Split(value, "/")
			if len(credParts) != 5 {
				return nil, fmt.Errorf("Invalid credential format")
			}
			sig.AccessKeyID = strings.ReplaceAll(credParts[0], "\"", "")
			sig.Date = credParts[1]
			sig.Region = credParts[2]
			sig.Service = credParts[3]
			sig.CredentialScope = strings.Join(credParts[1:], "/")
		case "SignedHeaders":
			sig.SignedHeaders = strings.Split(value, ";")
		case "Signature":
			sig.Signature = value
		}
	}

	if sig.AccessKeyID == "" {
		return nil, fmt.Errorf("Authorization header requires 'Credential' parameter")
	}
	if sig.Signature == "" {
		return nil, fmt.Errorf("Authorization header requires 'Signature' parameter")
	}

	return sig, nil
}

func testHMAC() {
	fmt.Println("=== Testing with quotes in secret key ===")

	// The secret key WITH quotes as it appears in environment variable
	secretKeyWithQuotes := "\"wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY\""

	kDate := hmacSHA256([]byte("AWS4"+secretKeyWithQuotes), "20251016")
	fmt.Println("kDate with quotes:", hex.EncodeToString(kDate))

	kRegion := hmacSHA256(kDate, "us-west-2")
	kService := hmacSHA256(kRegion, "sts")
	kSigning := hmacSHA256(kService, "aws4_request")

	fmt.Println("Signing key with quotes:", hex.EncodeToString(kSigning))

	stringToSign := "AWS4-HMAC-SHA256\n20251016T185909Z\n20251016/us-west-2/sts/aws4_request\nb1eade88f71175b2790108f7015d7ded65f982153aad655da6d50d5976ae14df"

	signature := hmacSHA256(kSigning, stringToSign)

	fmt.Println("Our signature:        ", hex.EncodeToString(signature))
	fmt.Println("AWS CLI's signature:  ", "c434c79f1567cf160a5cb88d2ad710a5bc524bbb6d481ce51b4b63787f9c21a4")
	fmt.Println("Match:", hex.EncodeToString(signature) == "c434c79f1567cf160a5cb88d2ad710a5bc524bbb6d481ce51b4b63787f9c21a4")
}

func calculateSignature(r *http.Request, body []byte, creds AWSCredentials, sigV4 *sigV4Components) (string, error) {
	amzDate := r.Header.Get("X-Amz-Date")
	canonicalRequest := buildCanonicalRequest(r, body, sigV4.SignedHeaders)
	stringToSign := buildStringToSign(canonicalRequest, sigV4, amzDate)

	signingKey := deriveSigningKey(creds.SecretAccessKey, sigV4.Date, sigV4.Region, sigV4.Service)

	log.Println("=== Final HMAC Input ===")
	log.Printf("Key (hex): %s\n", hex.EncodeToString(signingKey))
	log.Printf("Data (bytes): %v\n", []byte(stringToSign))
	log.Printf("Data (hex): %s\n", hex.EncodeToString([]byte(stringToSign)))

	signature := hmacSHA256(signingKey, stringToSign)
	finalSig := hex.EncodeToString(signature)

	log.Printf("Signature (hex): %s\n", finalSig)
	log.Println("=== End Final HMAC Input ===")

	return finalSig, nil
}

func buildCanonicalRequest(r *http.Request, body []byte, signedHeaders []string) string {
	// Method
	method := r.Method

	// Canonical URI
	canonicalURI := r.URL.Path
	if canonicalURI == "" {
		canonicalURI = "/"
	}

	// Canonical query string - MUST BE SORTED
	canonicalQueryString := buildCanonicalQueryString(r)

	// Canonical headers (already includes trailing newlines)
	canonicalHeaders := buildCanonicalHeaders(r, signedHeaders)

	// Signed headers
	signedHeadersStr := strings.Join(signedHeaders, ";")

	// Payload hash
	payloadHash := sha256Hash(body)

	canonicalRequest := fmt.Sprintf("%s\n%s\n%s\n%s\n%s\n%s",
		method,
		canonicalURI,
		canonicalQueryString,
		canonicalHeaders,
		signedHeadersStr,
		payloadHash,
	)

	return canonicalRequest
}

func buildCanonicalQueryString(r *http.Request) string {
	if r.URL.RawQuery == "" {
		return ""
	}

	// Split the raw query string by '&'
	params := strings.Split(r.URL.RawQuery, "&")

	// Sort the parameters
	sort.Strings(params)

	// Join them back with '&'
	return strings.Join(params, "&")
}

func buildCanonicalHeaders(r *http.Request, signedHeaders []string) string {
	headers := make(map[string]string)

	for _, header := range signedHeaders {
		headerLower := strings.ToLower(strings.TrimSpace(header))

		// Special handling for Host header
		if headerLower == "host" {
			headers[headerLower] = r.Host
			log.Printf("Header '%s' = '%s'", headerLower, r.Host)
			continue
		}

		values := r.Header[http.CanonicalHeaderKey(header)]
		if len(values) > 0 {
			trimmedValues := make([]string, len(values))
			for i, v := range values {
				trimmedValues[i] = strings.TrimSpace(v)
			}
			headers[headerLower] = strings.Join(trimmedValues, ",")
			log.Printf("Header '%s' = '%s'", headerLower, headers[headerLower])
		}
	}

	// Sort headers
	keys := make([]string, 0, len(headers))
	for k := range headers {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	// Build canonical headers string
	var canonical strings.Builder
	for _, k := range keys {
		canonical.WriteString(k)
		canonical.WriteString(":")
		canonical.WriteString(headers[k])
		canonical.WriteString("\n")
	}

	return canonical.String()
}

func buildStringToSign(canonicalRequest string, sigV4 *sigV4Components, amzDate string) string {
	log.Println("=== Canonical Request (raw) ===")
	log.Printf("%q\n", canonicalRequest) // Using %q to show escape characters
	log.Println("=== End Canonical Request (raw) ===")

	canonicalRequestHash := sha256Hash([]byte(canonicalRequest))
	log.Println("Canonical Request Hash:", canonicalRequestHash)

	stringToSign := fmt.Sprintf("AWS4-HMAC-SHA256\n%s\n%s\n%s",
		amzDate,
		sigV4.CredentialScope,
		canonicalRequestHash,
	)

	log.Println("=== String to Sign ===")
	log.Println(stringToSign)
	log.Println("=== End String to Sign ===")

	return stringToSign
}

func deriveSigningKey(secretKey, date, region, service string) []byte {
	log.Println("=== Deriving Signing Key ===")
	log.Println("Secret Key:", secretKey)
	log.Println("Date:", date)
	log.Println("Region:", region)
	log.Println("Service:", service)

	kDate := hmacSHA256([]byte("AWS4"+secretKey), date)
	log.Println("kDate:", hex.EncodeToString(kDate))

	kRegion := hmacSHA256(kDate, region)
	log.Println("kRegion:", hex.EncodeToString(kRegion))

	kService := hmacSHA256(kRegion, service)
	log.Println("kService:", hex.EncodeToString(kService))

	kSigning := hmacSHA256(kService, "aws4_request")
	log.Println("kSigning:", hex.EncodeToString(kSigning))
	log.Println("=== End Deriving Signing Key ===")

	return kSigning
}

func hmacSHA256(key []byte, data string) []byte {
	h := hmac.New(sha256.New, key)
	h.Write([]byte(data))
	return h.Sum(nil)
}

func sha256Hash(data []byte) string {
	hash := sha256.Sum256(data)
	return hex.EncodeToString(hash[:])
}

func writeAuthError(w http.ResponseWriter, code, message string) {
	w.Header().Set("Content-Type", "application/xml")
	w.WriteHeader(http.StatusForbidden)

	errorXML := fmt.Sprintf(`<?xml version="1.0" encoding="UTF-8"?>
<ErrorResponse xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
  <Error>
    <Type>Sender</Type>
    <Code>%s</Code>
    <Message>%s</Message>
  </Error>
  <RequestId>%d</RequestId>
</ErrorResponse>`, code, message, time.Now().Unix())

	w.Write([]byte(errorXML))
}