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) }