go-omni/storage/multipart.go

package storage

import (
	"context"
	"errors"
	"io"
	"strings"
	"time"
	"unicode"
)

// MultipartUpload is provider state, not authorization. Applications must bind
// it to the authenticated caller in their own durable session store.
type MultipartUpload struct {
	Key      string `json:"key"`
	UploadID string `json:"uploadId"`
}
type Part struct {
	Number int    `json:"number"`
	ETag   string `json:"etag"`
	Size   int64  `json:"size"`
}
type PartsPage struct {
	Parts      []Part `json:"parts"`
	NextMarker int    `json:"nextMarker"`
}
type MultipartStore interface {
	CreateMultipart(context.Context, string, string) (MultipartUpload, error)
	UploadPart(context.Context, MultipartUpload, int, io.Reader, int64) (Part, error)
	ListParts(context.Context, MultipartUpload, int, int) (PartsPage, error)
	CompleteMultipart(context.Context, MultipartUpload, []Part) (Info, error)
	AbortMultipart(context.Context, MultipartUpload) error
}
type PartPresigner interface {
	PresignPart(context.Context, MultipartUpload, int, int64, time.Duration) (UploadRequest, error)
}

func CheckMultipart(ctx context.Context, upload MultipartUpload) error {
	if err := Check(ctx, upload.Key); err != nil {
		return err
	}
	if upload.UploadID == "" || len(upload.UploadID) > 2048 || strings.IndexFunc(upload.UploadID, unicode.IsControl) >= 0 {
		return errors.New("storage: invalid upload ID")
	}
	return nil
}
func CheckPart(ctx context.Context, upload MultipartUpload, number int, size, maxBytes int64) error {
	if err := CheckMultipart(ctx, upload); err != nil {
		return err
	}
	if number < 1 || number > 10000 || size <= 0 {
		return errors.New("storage: part number must be 1..10000 and size positive")
	}
	if size > maxBytes || size > 5<<30 {
		return ErrTooLarge
	}
	return nil
}
func CheckPartsPage(ctx context.Context, upload MultipartUpload, after, limit int) error {
	if err := CheckMultipart(ctx, upload); err != nil {
		return err
	}
	if after < 0 || after > 10000 || limit < 1 || limit > 1000 {
		return errors.New("storage: invalid part pagination")
	}
	return nil
}

// VerifyParts checks provider-reported sizes/ETags before completing. Caller
// sizes are not trusted. Do not upload/replace parts concurrently with completion.
// A common 5 MiB minimum (except last) keeps S3/OSS interchangeable.
func VerifyParts(ctx context.Context, store MultipartStore, upload MultipartUpload, parts []Part, maxBytes int64) (int64, error) {
	if err := CheckMultipart(ctx, upload); err != nil {
		return 0, err
	}
	if len(parts) == 0 || len(parts) > 10000 || maxBytes <= 0 {
		return 0, errors.New("storage: invalid completion")
	}
	for i, p := range parts {
		if p.Number != i+1 || p.ETag == "" || len(p.ETag) > 256 || strings.ContainsAny(p.ETag, "\r\n") {
			return 0, errors.New("storage: ordered contiguous parts and ETags required")
		}
	}
	var total int64
	marker, seen := 0, 0
	for {
		page, err := store.ListParts(ctx, upload, marker, 1000)
		if err != nil {
			return 0, err
		}
		if len(page.Parts) > 1000 {
			return 0, errors.New("storage: invalid provider page")
		}
		for _, actual := range page.Parts {
			if seen >= len(parts) || actual.Number != parts[seen].Number || strings.Trim(actual.ETag, "\"") != strings.Trim(parts[seen].ETag, "\"") || actual.Size <= 0 {
				return 0, errors.New("storage: parts changed or incomplete")
			}
			if actual.Size > 5<<30 || actual.Size > maxBytes-total {
				return 0, ErrTooLarge
			}
			if seen < len(parts)-1 && actual.Size < 5<<20 {
				return 0, errors.New("storage: non-final part smaller than 5 MiB")
			}
			total += actual.Size
			seen++
		}
		if page.NextMarker == 0 {
			break
		}
		if len(page.Parts) == 0 || page.NextMarker <= marker || page.NextMarker > 10000 || page.NextMarker != page.Parts[len(page.Parts)-1].Number {
			return 0, errors.New("storage: invalid provider pagination")
		}
		marker = page.NextMarker
	}
	if seen != len(parts) {
		return 0, errors.New("storage: missing parts")
	}
	return total, nil
}
func (u *Uploader) CreateMultipart(ctx context.Context, key, contentType string) (MultipartUpload, error) {
	if err := u.policy.Check(0, contentType); err != nil {
		return MultipartUpload{}, err
	}
	s, ok := u.store.(MultipartStore)
	if !ok {
		return MultipartUpload{}, ErrUnsupported
	}
	return s.CreateMultipart(ctx, key, contentType)
}
func (u *Uploader) UploadPart(ctx context.Context, upload MultipartUpload, number int, body io.Reader, size int64) (Part, error) {
	if err := CheckPart(ctx, upload, number, size, u.policy.MaxBytes); err != nil {
		return Part{}, err
	}
	s, ok := u.store.(MultipartStore)
	if !ok {
		return Part{}, ErrUnsupported
	}
	return s.UploadPart(ctx, upload, number, body, size)
}
func (u *Uploader) PresignPart(ctx context.Context, upload MultipartUpload, number int, size int64, ttl time.Duration) (UploadRequest, error) {
	if err := CheckPart(ctx, upload, number, size, u.policy.MaxBytes); err != nil {
		return UploadRequest{}, err
	}
	s, ok := u.store.(PartPresigner)
	if !ok {
		return UploadRequest{}, ErrUnsupported
	}
	return s.PresignPart(ctx, upload, number, size, ttl)
}
func (u *Uploader) ListParts(ctx context.Context, upload MultipartUpload, after, limit int) (PartsPage, error) {
	s, ok := u.store.(MultipartStore)
	if !ok {
		return PartsPage{}, ErrUnsupported
	}
	return s.ListParts(ctx, upload, after, limit)
}
func (u *Uploader) CompleteMultipart(ctx context.Context, upload MultipartUpload, parts []Part) (Info, error) {
	s, ok := u.store.(MultipartStore)
	if !ok {
		return Info{}, ErrUnsupported
	}
	if _, err := VerifyParts(ctx, s, upload, parts, u.policy.MaxBytes); err != nil {
		return Info{}, err
	}
	return s.CompleteMultipart(ctx, upload, parts)
}
func (u *Uploader) AbortMultipart(ctx context.Context, upload MultipartUpload) error {
	s, ok := u.store.(MultipartStore)
	if !ok {
		return ErrUnsupported
	}
	return s.AbortMultipart(ctx, upload)
}