go-omni/internal/storagefixture/multipart.go

package storagefixture

import (
	"crypto/md5"
	"encoding/xml"
	"fmt"
	"io"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"
)

// MultipartHandler checks protocol sequencing, not provider authentication.
// The fixture is memory bounded and should only be mounted on httptest servers.
func MultipartHandler() http.Handler {
	type part struct {
		Number int    `xml:"PartNumber"`
		ETag   string `xml:"ETag"`
		Size   int64  `xml:"Size"`
		data   []byte
	}
	type upload struct {
		key, contentType string
		parts            map[int]part
	}
	var mu sync.Mutex
	uploads := map[string]*upload{}
	objects := map[string]value{}
	sequence := 0
	simple := Handler()
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		q := r.URL.Query()
		id := q.Get("uploadId")
		if _, start := q["uploads"]; !start && id == "" {
			mu.Lock()
			v, ok := objects[r.URL.Path]
			if ok && r.Method == "DELETE" {
				delete(objects, r.URL.Path)
				mu.Unlock()
				w.WriteHeader(204)
				return
			}
			mu.Unlock()
			if ok && (r.Method == "GET" || r.Method == "HEAD") {
				w.Header().Set("Content-Length", strconv.Itoa(len(v.data)))
				w.Header().Set("Content-Type", v.contentType)
				w.Header().Set("ETag", v.etag)
				w.Header().Set("Last-Modified", time.Now().UTC().Format(http.TimeFormat))
				if r.Method == "GET" {
					w.Write(v.data)
				}
				return
			}
			simple.ServeHTTP(w, r)
			return
		}
		mu.Lock()
		defer mu.Unlock()
		fail := func(code string, status int) {
			w.Header().Set("Content-Type", "application/xml")
			w.WriteHeader(status)
			fmt.Fprintf(w, "<Error><Code>%s</Code><Message>fixture rejection</Message></Error>", code)
		}
		if r.Header.Get("Authorization") == "" && q.Get("X-Amz-Signature") == "" && q.Get("x-oss-signature") == "" {
			fail("AccessDenied", 403)
			return
		}
		if _, ok := q["uploads"]; ok && r.Method == "POST" {
			if len(uploads) >= 8 {
				fail("TooManyUploads", 503)
				return
			}
			sequence++
			id = strconv.Itoa(sequence)
			uploads[id] = &upload{key: r.URL.Path, contentType: r.Header.Get("Content-Type"), parts: map[int]part{}}
			fmt.Fprintf(w, "<InitiateMultipartUploadResult><Bucket>testbucket</Bucket><Key>file</Key><UploadId>%s</UploadId></InitiateMultipartUploadResult>", id)
			return
		}
		u, ok := uploads[id]
		if !ok || u.key != r.URL.Path {
			fail("NoSuchUpload", 404)
			return
		}
		switch r.Method {
		case "PUT":
			number, _ := strconv.Atoi(q.Get("partNumber"))
			if number < 1 || number > 3 {
				fail("InvalidPart", 400)
				return
			}
			data, err := io.ReadAll(io.LimitReader(r.Body, (6<<20)+1))
			if err != nil || len(data) > 6<<20 {
				fail("EntityTooLarge", 413)
				return
			}
			p := part{Number: number, ETag: fmt.Sprintf("\"%x\"", md5.Sum(data)), Size: int64(len(data)), data: data}
			u.parts[number] = p
			w.Header().Set("ETag", p.ETag)
		case "GET":
			marker, _ := strconv.Atoi(q.Get("part-number-marker"))
			limit, _ := strconv.Atoi(q.Get("max-parts"))
			if limit < 1 {
				limit = 1000
			}
			numbers := []int{}
			for n := range u.parts {
				if n > marker {
					numbers = append(numbers, n)
				}
			}
			sort.Ints(numbers)
			var result struct {
				XMLName   xml.Name `xml:"ListPartsResult"`
				Truncated bool     `xml:"IsTruncated"`
				Next      int      `xml:"NextPartNumberMarker"`
				Parts     []part   `xml:"Part"`
			}
			if len(numbers) > limit {
				result.Truncated = true
				numbers = numbers[:limit]
				result.Next = numbers[len(numbers)-1]
			}
			for _, n := range numbers {
				result.Parts = append(result.Parts, u.parts[n])
			}
			xml.NewEncoder(w).Encode(result)
		case "POST":
			var request struct {
				Parts []part `xml:"Part"`
			}
			if xml.NewDecoder(io.LimitReader(r.Body, 65537)).Decode(&request) != nil || len(request.Parts) == 0 {
				fail("InvalidPart", 400)
				return
			}
			var data []byte
			for _, p := range request.Parts {
				actual, ok := u.parts[p.Number]
				if !ok || strings.Trim(p.ETag, "\"") != strings.Trim(actual.ETag, "\"") {
					fail("InvalidPart", 400)
					return
				}
				if len(data)+len(actual.data) > 7<<20 {
					fail("EntityTooLarge", 413)
					return
				}
				data = append(data, actual.data...)
			}
			objects[u.key] = value{data: data, contentType: u.contentType, etag: fmt.Sprintf("\"%x\"", md5.Sum(data))}
			delete(uploads, id)
			fmt.Fprint(w, "<CompleteMultipartUploadResult><Bucket>testbucket</Bucket><Key>file</Key><ETag>\"completed\"</ETag></CompleteMultipartUploadResult>")
		case "DELETE":
			delete(uploads, id)
			w.WriteHeader(204)
		default:
			fail("InvalidRequest", 400)
		}
	})
}