go-omni/storage/signed.go

package storage

import (
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"errors"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"time"
)

type Signer struct {
	store        Store
	base         *url.URL
	secret       []byte
	now          func() time.Time
	uploadPolicy *UploadPolicy
}

func NewSigner(store Store, baseURL string, secret []byte) (*Signer, error) {
	base, err := url.Parse(baseURL)
	if err != nil {
		return nil, err
	}
	if store == nil || len(secret) < 32 || base.Host == "" || (base.Scheme != "https" && base.Scheme != "http") || base.RawQuery != "" || base.Fragment != "" || base.User != nil {
		return nil, errors.New("storage: store, HTTP base URL and >=32-byte secret required")
	}
	return &Signer{store: store, base: base, secret: append([]byte(nil), secret...), now: time.Now}, nil
}
func (s *Signer) signature(key, expiry string) string {
	mac := hmac.New(sha256.New, s.secret)
	io.WriteString(mac, "GET\n"+key+"\n"+expiry)
	return base64.RawURLEncoding.EncodeToString(mac.Sum(nil))
}
func (s *Signer) PresignGet(ctx context.Context, key string, ttl time.Duration) (*url.URL, error) {
	if err := Check(ctx, key); err != nil {
		return nil, err
	}
	if ttl < time.Second || ttl > time.Hour {
		return nil, errors.New("storage: TTL must be 1s..1h")
	}
	expiry := strconv.FormatInt(s.now().Add(ttl).Unix(), 10)
	u := *s.base
	u.RawQuery = url.Values{"key": {key}, "expires": {expiry}, "signature": {s.signature(key, expiry)}}.Encode()
	return &u, nil
}

func (s *Signer) readSignature(q url.Values) string {
	values := make(url.Values, len(q))
	for key, value := range q {
		if key != "signature" {
			values[key] = value
		}
	}
	mac := hmac.New(sha256.New, s.secret)
	io.WriteString(mac, "GET-OPTIONS-v2\n"+values.Encode())
	return base64.RawURLEncoding.EncodeToString(mac.Sum(nil))
}

func (s *Signer) PresignGetWithOptions(ctx context.Context, key string, ttl time.Duration, options ReadOptions) (ReadLink, error) {
	if err := Check(ctx, key); err != nil {
		return ReadLink{}, err
	}
	if ttl < time.Second || ttl > time.Hour {
		return ReadLink{}, errors.New("storage: TTL must be 1s..1h")
	}
	headers, err := options.ResponseHeaders()
	if err != nil {
		return ReadLink{}, err
	}
	if options.Disposition == "" {
		options.Disposition = DispositionAttachment
	}
	expires := s.now().Add(ttl).UTC().Truncate(time.Second)
	q := url.Values{"v": {"2"}, "key": {key}, "expires": {strconv.FormatInt(expires.Unix(), 10)}, "disposition": {options.Disposition}, "filename": {options.FileName}, "contentType": {headers.Get("Content-Type")}}
	q.Set("signature", s.readSignature(q))
	u := *s.base
	u.RawQuery = q.Encode()
	return ReadLink{URL: u.String(), ExpiresAt: expires}, nil
}

// ServeHTTP serves signed GET/HEAD requests, with optional preview headers.
// Seekable object bodies support Range and conditional reads; other stores keep
// streaming full responses. No full-object buffer is used. Issue links only after
// authorization; possession of an unexpired URL grants access to its object.
func (s *Signer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method == http.MethodPut && s.uploadPolicy != nil {
		s.serveUpload(w, r)
		return
	}
	w = &readResponseWriter{ResponseWriter: w, head: r.Method == http.MethodHead}
	if r.Method != "GET" && r.Method != "HEAD" {
		w.Header().Set("Allow", "GET, HEAD")
		http.Error(w, "method not allowed", 405)
		return
	}
	q, err := url.ParseQuery(r.URL.RawQuery)
	keys := []string{"key", "expires", "signature"}
	if q.Get("v") == "2" {
		keys = append(keys, "v", "disposition", "filename", "contentType")
	}
	valid := err == nil && len(q) == len(keys)
	for _, key := range keys {
		valid = valid && len(q[key]) == 1
	}
	if !valid {
		http.Error(w, "invalid link", 403)
		return
	}
	key, expiry := q.Get("key"), q.Get("expires")
	expires, err := strconv.ParseInt(expiry, 10, 64)
	now := s.now().Unix()
	expected := s.signature(key, expiry)
	if q.Get("v") == "2" {
		expected = s.readSignature(q)
	}
	if err != nil || expires <= now || expires > now+3600 || Check(r.Context(), key) != nil || !hmac.Equal([]byte(q.Get("signature")), []byte(expected)) {
		http.Error(w, "invalid or expired link", 403)
		return
	}
	options := ReadOptions{}
	if q.Get("v") == "2" {
		options = ReadOptions{Disposition: q.Get("disposition"), FileName: q.Get("filename"), ContentType: q.Get("contentType")}
	}
	headers, err := options.ResponseHeaders()
	if err != nil {
		http.Error(w, "invalid link", 403)
		return
	}
	object, err := s.store.Get(r.Context(), key)
	if err != nil {
		status := 500
		if errors.Is(err, ErrNotFound) {
			status = 404
		}
		http.Error(w, "object unavailable", status)
		return
	}
	if object.Body == nil {
		http.Error(w, "object unavailable", 500)
		return
	}
	defer object.Body.Close()
	if object.Size < 0 {
		http.Error(w, "object unavailable", 500)
		return
	}
	contentType := object.ContentType
	if options.ContentType != "" {
		contentType = headers.Get("Content-Type")
	}
	if contentType == "" {
		contentType = "application/octet-stream"
	}
	w.Header().Set("Content-Type", contentType)
	w.Header().Set("Content-Disposition", headers.Get("Content-Disposition"))
	if body, ok := object.Body.(io.ReadSeeker); ok {
		// Range describes a GET representation. HEAD reports full-object headers.
		if r.Method == http.MethodHead {
			r = r.Clone(r.Context())
			r.Header.Del("Range")
		}
		http.ServeContent(w, r, "", object.ModifiedAt, body)
		return
	}
	w.Header().Set("Content-Length", strconv.FormatInt(object.Size, 10))
	w.WriteHeader(http.StatusOK)
	if r.Method == "GET" {
		_, _ = io.Copy(w, object.Body)
	}
}

// ServeContent may strip cache headers on errors in newer Go versions. Restore
// the signed-link policy at the actual write, including 304/416 and HEAD errors.
type readResponseWriter struct {
	http.ResponseWriter
	head, wrote bool
}

func (w *readResponseWriter) WriteHeader(status int) {
	if w.wrote {
		return
	}
	w.wrote = true
	w.Header().Set("Cache-Control", "private, no-store")
	w.Header().Set("X-Content-Type-Options", "nosniff")
	w.ResponseWriter.WriteHeader(status)
}
func (w *readResponseWriter) Write(p []byte) (int, error) {
	if !w.wrote {
		w.WriteHeader(http.StatusOK)
	}
	if w.head {
		return len(p), nil
	}
	return w.ResponseWriter.Write(p)
}
func (w *readResponseWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }

var _ Presigner = (*Signer)(nil)
var _ ReadPresigner = (*Signer)(nil)