go-omni/storage/upload.go

package storage

import (
	"bytes"
	"context"
	"errors"
	"io"
	"mime"
	"net/http"
	"strings"
	"time"
)

var ErrUnsupported = errors.New("storage: capability not supported")
var ErrContentType = errors.New("storage: content type not allowed")
var ErrSizeMismatch = errors.New("storage: declared size mismatch")

// UploadRequest is a short-lived bearer capability; never log its URL. Headers
// must accompany the request. Browsers derive Content-Length from the Blob body.
type UploadRequest struct {
	Method    string      `json:"method"`
	URL       string      `json:"url"`
	Headers   http.Header `json:"headers"`
	ExpiresAt time.Time   `json:"expiresAt"`
}

type UploadPresigner interface {
	PresignPut(context.Context, string, int64, string, time.Duration) (UploadRequest, error)
}

// UploadPolicy contains only transport-neutral constraints. ContentTypes is an
// exact MIME allowlist (parameters are ignored). Empty permits any valid type.
type UploadPolicy struct {
	MaxBytes     int64    `json:"maxBytes"`
	ContentTypes []string `json:"contentTypes"`
}

func (p UploadPolicy) Validate() error {
	if p.MaxBytes <= 0 || p.MaxBytes > 5<<40 {
		return errors.New("storage: maxBytes must be 1..5 TiB")
	}
	for _, value := range p.ContentTypes {
		t, err := MediaType(value)
		if err != nil || strings.Contains(t, "*") {
			return ErrContentType
		}
	}
	return nil
}
func MediaType(value string) (string, error) {
	if value == "" {
		return "application/octet-stream", nil
	}
	if len(value) > 256 || strings.ContainsAny(value, "\r\n") {
		return "", ErrContentType
	}
	t, _, err := mime.ParseMediaType(value)
	if err != nil || !strings.Contains(t, "/") || strings.Contains(t, "*") {
		return "", ErrContentType
	}
	return strings.ToLower(t), nil
}
func (p UploadPolicy) Check(size int64, contentType string) error {
	if err := p.Validate(); err != nil {
		return err
	}
	if size < 0 {
		return ErrSizeMismatch
	}
	if size > p.MaxBytes {
		return ErrTooLarge
	}
	t, err := MediaType(contentType)
	if err != nil {
		return err
	}
	if len(p.ContentTypes) == 0 {
		return nil
	}
	for _, allowed := range p.ContentTypes {
		a, _ := MediaType(allowed)
		if a == t {
			return nil
		}
	}
	return ErrContentType
}
func CheckUpload(ctx context.Context, key string, size, maxBytes int64, contentType string, ttl time.Duration) error {
	if err := Check(ctx, key); err != nil {
		return err
	}
	if err := (UploadPolicy{MaxBytes: maxBytes}).Check(size, contentType); err != nil {
		return err
	}
	if ttl < time.Second || ttl > time.Hour {
		return errors.New("storage: TTL must be 1s..1h")
	}
	return nil
}

// PrepareUpload checks the first 512 bytes against the declared type when an
// allowlist is set, then streams the body with exact-length checking. It does
// not close the caller's reader. MIME sniffing is not antivirus/file validation.
func (p UploadPolicy) PrepareUpload(ctx context.Context, reader io.Reader, size int64, contentType string) (io.Reader, error) {
	if ctx == nil || reader == nil {
		return nil, errors.New("storage: context and reader required")
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if err := p.Check(size, contentType); err != nil {
		return nil, err
	}
	var r io.Reader = ContextReader{Context: ctx, Reader: reader}
	if len(p.ContentTypes) > 0 && size > 0 {
		prefix := make([]byte, min(size, 512))
		if _, err := io.ReadFull(r, prefix); err != nil {
			return nil, errors.Join(ErrSizeMismatch, err)
		}
		detected, _ := MediaType(http.DetectContentType(prefix))
		declared, _ := MediaType(contentType)
		if detected != declared {
			return nil, ErrContentType
		}
		r = io.MultiReader(bytes.NewReader(prefix), r)
	}
	if size == 0 {
		var b [1]byte
		n, err := io.ReadFull(r, b[:])
		if n > 0 {
			return nil, ErrSizeMismatch
		}
		if err != io.EOF {
			return nil, err
		}
		return bytes.NewReader(nil), nil
	}
	return &exactReader{reader: r, remaining: size}, nil
}

type exactReader struct {
	reader    io.Reader
	remaining int64
	done      bool
}

func (r *exactReader) Read(p []byte) (int, error) {
	if len(p) == 0 {
		return 0, nil
	}
	if r.done {
		return 0, io.EOF
	}
	if int64(len(p)) > r.remaining {
		p = p[:r.remaining]
	}
	n, err := r.reader.Read(p)
	r.remaining -= int64(n)
	if r.remaining == 0 {
		r.done = true
		// Probe before handing the last byte to a known-length SDK, which might
		// otherwise stop reading without observing a trailing error.
		var b [1]byte
		extra, probeErr := io.ReadFull(r.reader, b[:])
		if extra > 0 {
			return max(0, n-1), ErrSizeMismatch
		}
		if probeErr != io.EOF {
			return max(0, n-1), probeErr
		}
		if err != nil && err != io.EOF {
			return max(0, n-1), err
		}
		return n, io.EOF
	}
	if err == io.EOF {
		r.done = true
		return n, ErrSizeMismatch
	}
	return n, err
}

// Uploader applies a policy without coupling the framework to business fields.
// Store remains unchanged; optional operations return ErrUnsupported when absent.
type Uploader struct {
	store  Store
	policy UploadPolicy
}

func NewUploader(store Store, policy UploadPolicy) (*Uploader, error) {
	if store == nil {
		return nil, errors.New("storage: store required")
	}
	if err := policy.Validate(); err != nil {
		return nil, err
	}
	policy.ContentTypes = append([]string(nil), policy.ContentTypes...)
	return &Uploader{store: store, policy: policy}, nil
}
func (u *Uploader) Put(ctx context.Context, key string, body io.Reader, size int64, contentType string) (Info, error) {
	if err := Check(ctx, key); err != nil {
		return Info{}, err
	}
	reader, err := u.policy.PrepareUpload(ctx, body, size, contentType)
	if err != nil {
		return Info{}, err
	}
	return u.store.Put(ctx, key, reader, size, contentType)
}
func (u *Uploader) PresignPut(ctx context.Context, key string, size int64, contentType string, ttl time.Duration) (UploadRequest, error) {
	if err := u.policy.Check(size, contentType); err != nil {
		return UploadRequest{}, err
	}
	signer, ok := u.store.(UploadPresigner)
	if !ok {
		return UploadRequest{}, ErrUnsupported
	}
	return signer.PresignPut(ctx, key, size, contentType, ttl)
}