s3 proxy

package main

import (
	"flag"

	"bytes"
	"encoding/xml"
	"fmt"
	"io"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws/credentials"
	v4 "github.com/aws/aws-sdk-go/aws/signer/v4"
)

// Handler is a special handler that re-signs any AWS S3 request and sends it upstream
type Handler struct{}

var (
	accessKey  string
	secretKey  string
	region     string
	s3Endpoint string
	proxy      *httputil.ReverseProxy
)

/**
<?xml version="1.0" encoding="UTF-8"?>
<Error>
  <Code>NoSuchKey</Code>
  <Message>The resource you requested does not exist</Message>
  <Resource>/mybucket/myfoto.jpg</Resource>
  <RequestId>4442587FB7D0A2F9</RequestId>
</Error>
****/

/***

2022-12-19 09:00:17,412 - MainThread - botocore.parsers - DEBUG - Response headers: {'x-amz-request-id': 'SVD07H84ERJKTTWV', 'x-amz-id-2': '3ug19S7p5ucOL+PM+tXqtUmM8uGxxDr1tS0C9M+VAp9vdULO0tTfRs1XmakYcNh/Cgfny9XG6zQ=', 'Content-Type': 'application/xml', 'Transfer-Encoding': 'chunked', 'Date': 'Mon, 19 Dec 2022 17:00:16 GMT', 'Server': 'AmazonS3'}
2022-12-19 09:00:17,412 - MainThread - botocore.parsers - DEBUG - Response body:
b'<?xml version="1.0" encoding="UTF-8"?>\n
<Error>
<Code>InvalidAccessKeyId</Code>
<Message>The AWS Access Key Id you provided does not exist in our records.</Message>
<AWSAccessKeyId>...</AWSAccessKeyId>
<RequestId>SVD07H84ERJKTTWV</RequestId>
<HostId>3ug19S7p5ucOL+PM+tXqtUmM8uGxxDr1tS0C9M+VAp9vdULO0tTfRs1XmakYcNh/Cgfny9XG6zQ=</HostId>
</Error>'
***/

type xmlErrorResponse struct {
	XMLName   xml.Name `xml:"Error"`
	Code      string   `xml:"Code"`
	Message   string   `xml:"Message"`
	RequestId string   `xml:"RequestId"`
}

// func combineURL(base string, paths ...string) string {
// 	p := path.Join(paths...)
// 	return fmt.Sprintf("%s/%s", strings.TrimRight(base, "/"), strings.TrimLeft(p, "/"))
// }

func sendForbidden(response http.ResponseWriter, msg string) {
	resp := &xmlErrorResponse{
		Code:      "InvalidAccessKeyId",
		Message:   "The AWS Access Key ID is invalid: " + msg,
		RequestId: "id",
	}

	x, err := xml.MarshalIndent(resp, "", "	")

	if err != nil {
		http.Error(response, err.Error(), http.StatusInternalServerError)
		return
	}

	response.WriteHeader(http.StatusForbidden)
	response.Header().Set("x-amz-request-id", "application/xml")
	response.Header().Set("x-amz-id-2", "application/xml")
	response.Header().Set("Content-Type", "application/xml")
	response.Write(x)
}

func (h *Handler) assembleUpstreamReq(signer *v4.Signer, req *http.Request, region string) (*http.Request, error) {
	upstreamUrl, err := url.Parse(s3Endpoint)
	if err != nil {
		return nil, fmt.Errorf("failed to parse upstream url %v", err)
	}

	proxyURL := *req.URL
	proxyURL.Scheme = upstreamUrl.Scheme
	proxyURL.Host = upstreamUrl.Host
	proxyURL.RawPath = req.URL.Path
	proxyReq, err := http.NewRequest(req.Method, proxyURL.String(), req.Body)
	if err != nil {
		return nil, err
	}
	if val, ok := req.Header["Content-Type"]; ok {
		proxyReq.Header["Content-Type"] = val
	}
	if val, ok := req.Header["Content-Md5"]; ok {
		proxyReq.Header["Content-Md5"] = val
	}

	// Sign the upstream request
	if err := h.sign(signer, proxyReq, region); err != nil {
		return nil, err
	}

	// Add origin headers after request is signed (no overwrite)
	copyHeaderWithoutOverwrite(proxyReq.Header, req.Header)

	return proxyReq, nil
}

// Do validates the incoming request and create a new request for an upstream server
func (h *Handler) buildUpstreamRequest(req *http.Request) (*http.Request, error) {

	signer := v4.NewSigner(credentials.NewStaticCredentials(accessKey, secretKey, ""),
		func(s *v4.Signer) { s.DisableURIPathEscaping = true })

	// Assemble a new upstream request
	proxyReq, err := h.assembleUpstreamReq(signer, req, region)
	if err != nil {
		return nil, err
	}

	// Disable Go's "Transfer-Encoding: chunked" madness
	proxyReq.ContentLength = req.ContentLength

	proxyReqDump, _ := httputil.DumpRequest(proxyReq, false)
	fmt.Printf("Proxying request: %v\n", string(proxyReqDump))

	return proxyReq, nil
}

func (h *Handler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
	initialReqDump, _ := httputil.DumpRequest(req, false)
	fmt.Printf("Initial request dump: %v\n", string(initialReqDump))

	authHeader := req.Header["Authorization"]

	/** expected Authorization header
	Authorization: AWS4-HMAC-SHA256
	Credential=AKIAIOSFODNN7EXAMPLE/20130524/us-east-1/s3/aws4_request,
	SignedHeaders=host;range;x-amz-date,
	Signature=fe5f80f77d5fa3beca038a248ff027d0445342fe2855ddc963176630326f1024
	***/

	tokens := strings.Split(authHeader[0], " ")
	rawToken := strings.Split(strings.Split(tokens[1], "=")[1], "/")[0]

	// VALIDATE rawToken, extract user info, verify issuer etc

	// Get the filepath (if any) to make policy decisions
	u, _ := url.Parse(req.URL.Path)

	urlArr := strings.SplitN(u.Path, "/", 3)

	bucket := urlArr[1]

	fmt.Printf("urlArr %v:%v:%v\n", urlArr, len(urlArr), bucket)

	filePath := "/"
	if len(urlArr) == 3 {
		filePath = filePath + urlArr[2]
	}

	// MAKE POLICY DECISION and sendForbidden()

	fmt.Printf("method %v rawToken %v FILEPATH bucket %v filePath %v path %v\n", req.Method, rawToken, bucket, filePath, u.Path)

	proxyReq, err := h.buildUpstreamRequest(req)
	if err != nil {
		fmt.Printf("unable to proxy request: %v\n", err)
		resp.WriteHeader(http.StatusBadRequest)
		return
	}

	url, _ := url.Parse(s3Endpoint)
	proxy = httputil.NewSingleHostReverseProxy(url)
	proxy.FlushInterval = 1
	proxy.ServeHTTP(resp, proxyReq)
}

func (h *Handler) sign(signer *v4.Signer, req *http.Request, region string) error {
	body := bytes.NewReader([]byte{})
	if req.Body != nil {
		b, err := io.ReadAll(req.Body)
		if err != nil {
			return err
		}
		body = bytes.NewReader(b)
	}

	_, err := signer.Sign(req, body, "s3", region, time.Now())
	return err
}

func copyHeaderWithoutOverwrite(dst http.Header, src http.Header) {
	for k, v := range src {
		if _, ok := dst[k]; !ok {
			for _, vv := range v {
				dst.Add(k, vv)
			}
		}
	}
}

func main() {

	flag.StringVar(&s3Endpoint, "s3Endpoint", "https://s3.us-east-1.amazonaws.com/", "s3 endpoint")
	flag.StringVar(&accessKey, "accessKey", "", "Access Key")
	flag.StringVar(&secretKey, "secretKey", "", "Secret Key")
	flag.StringVar(&region, "region", "us-east-1", "s3 authorizer configuration file")
	flag.Parse()

	if region == "" {
		region = "us-east-1"
	}

	var wrappedHandler http.Handler = &Handler{}

	fmt.Printf("s3 proxy listening \n")
	http.ListenAndServe("127.0.0.1:30000", wrappedHandler)
}