frostfs-s3-gw/playback/multipart.go

package playback

import (
	"context"
	"encoding/xml"
	"errors"
	"fmt"
	"net/http"
)

type MultipartUpload struct {
	XMLName  xml.Name `xml:"http://s3.amazonaws.com/doc/2006-03-01/ InitiateMultipartUploadResult" json:"-"`
	Bucket   string   `json:"bucket" xml:"Bucket"`
	Key      string   `json:"key" xml:"Key"`
	UploadID string   `json:"uploadId" xml:"UploadId"`
}

func WithMultiparts(ctx context.Context) context.Context {
	return context.WithValue(ctx, multipartKey{}, map[string]MultipartUpload{})
}

func SetMultipartUpload(ctx context.Context, oldUploadID string, upload MultipartUpload) error {
	mparts, ok := ctx.Value(multipartKey{}).(map[string]MultipartUpload)
	if !ok {
		return errors.New("multiparts not set")
	}
	mparts[oldUploadID] = upload

	return nil
}

func GetMultipart(ctx context.Context, oldUploadID string) (MultipartUpload, error) {
	mparts, ok := ctx.Value(multipartKey{}).(map[string]MultipartUpload)
	if !ok {
		return MultipartUpload{}, errors.New("no multipart map set")
	}

	return mparts[oldUploadID], nil
}

func HandleResponse(ctx context.Context, r *http.Request, resp []byte, logResponse []byte) error {
	var mpart, mpartOld MultipartUpload
	if r.Method == "POST" && r.URL.Query().Has("uploads") {
		err1 := xml.Unmarshal(resp, &mpart)
		err2 := xml.Unmarshal(logResponse, &mpartOld)
		if err1 != nil || err2 != nil {
			return errors.New("xml unmarshal error")
		}
		if mpartOld.UploadID != "" {
			if err := SetMultipartUpload(ctx, mpartOld.UploadID, mpart); err != nil {
				return err
			}
		}
	}

	return nil
}

func SwapUploadID(ctx context.Context, r *http.Request) error {
	var uploadID string
	query := r.URL.Query()

	if query.Has("uploadId") {
		uploadID = query.Get("uploadId")
		mpart, err := GetMultipart(ctx, uploadID)
		if err != nil {
			return fmt.Errorf("failed to get multipart upload: %w", err)
		}
		query.Set("uploadId", mpart.UploadID)
		r.URL.RawQuery = query.Encode()
	}

	return nil
}