243 lines
7.6 KiB
Go
243 lines
7.6 KiB
Go
package middleware
|
|
|
|
import (
|
|
"bytes"
|
|
"crypto/sha256"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
apierrors "git.yuntuo.hk/k2safe/go-omni/errors"
|
|
"git.yuntuo.hk/k2safe/go-omni/httpx"
|
|
)
|
|
|
|
type IdempotencyConfig struct {
|
|
Store IdempotencyStore `json:"-"` // Optional shared store; nil keeps local behavior.
|
|
TTL time.Duration `json:"ttl"`
|
|
MaxEntries int `json:"maxEntries"`
|
|
MaxBytes int `json:"maxBytes"`
|
|
MaxRequestBytes int64 `json:"maxRequestBytes"`
|
|
MaxResponseBytes int `json:"maxResponseBytes"`
|
|
MaxInFlight int `json:"maxInFlight"`
|
|
// Scope must return a server-verified tenant/user identity. Register auth
|
|
// before this middleware. Client-provided identity is not authorization.
|
|
Scope func(*httpx.Context) (string, error) `json:"-"`
|
|
}
|
|
|
|
type replay struct {
|
|
fingerprint [32]byte
|
|
status int
|
|
headers http.Header
|
|
body []byte
|
|
expires time.Time
|
|
running, failed bool
|
|
size int
|
|
}
|
|
type Idempotency struct {
|
|
mu sync.Mutex
|
|
config IdempotencyConfig
|
|
items map[[32]byte]*replay
|
|
bytes int
|
|
slots chan struct{}
|
|
now func() time.Time
|
|
}
|
|
|
|
func NewIdempotency(cfg IdempotencyConfig) (*Idempotency, error) {
|
|
if cfg.Scope == nil || cfg.TTL < 0 || cfg.MaxEntries < 0 || cfg.MaxBytes < 0 || cfg.MaxRequestBytes < 0 || cfg.MaxResponseBytes < 0 || cfg.MaxInFlight < 0 {
|
|
return nil, errors.New("middleware: invalid idempotency configuration or missing scope")
|
|
}
|
|
if cfg.TTL == 0 {
|
|
cfg.TTL = 10 * time.Minute
|
|
}
|
|
if cfg.Store != nil && cfg.TTL > 30*24*time.Hour {
|
|
return nil, errors.New("middleware: shared idempotency TTL exceeds 30d")
|
|
}
|
|
if cfg.MaxEntries == 0 {
|
|
cfg.MaxEntries = 256
|
|
}
|
|
if cfg.MaxBytes == 0 {
|
|
cfg.MaxBytes = 16 << 20
|
|
}
|
|
if cfg.MaxRequestBytes == 0 {
|
|
cfg.MaxRequestBytes = 1 << 20
|
|
}
|
|
if cfg.MaxResponseBytes == 0 {
|
|
cfg.MaxResponseBytes = 64 << 10
|
|
}
|
|
if cfg.MaxInFlight == 0 {
|
|
cfg.MaxInFlight = 32
|
|
}
|
|
if cfg.MaxRequestBytes > 1<<30 {
|
|
return nil, errors.New("middleware: maxRequestBytes exceeds 1 GiB")
|
|
}
|
|
return &Idempotency{config: cfg, items: make(map[[32]byte]*replay), slots: make(chan struct{}, cfg.MaxInFlight), now: time.Now}, nil
|
|
}
|
|
|
|
func (i *Idempotency) Middleware() httpx.Middleware {
|
|
return func(next httpx.Handler) httpx.Handler {
|
|
return func(c *httpx.Context) error {
|
|
r := c.Request()
|
|
if r.Method == "GET" || r.Method == "HEAD" || r.Method == "OPTIONS" {
|
|
return next(c)
|
|
}
|
|
select {
|
|
case i.slots <- struct{}{}:
|
|
defer func() { <-i.slots }()
|
|
default:
|
|
return apierrors.New(503, "IDEMPOTENCY_BUSY", "idempotency capacity reached")
|
|
}
|
|
key := r.Header.Get("Idempotency-Key")
|
|
if len(r.Header.Values("Idempotency-Key")) != 1 || len(key) < 1 || len(key) > 128 || strings.IndexFunc(key, func(r rune) bool { return r < 33 || r > 126 }) >= 0 {
|
|
return apierrors.BadRequest("INVALID_IDEMPOTENCY_KEY", "Idempotency-Key must contain 1..128 printable ASCII characters")
|
|
}
|
|
scope, err := i.config.Scope(c)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if scope == "" || len(scope) > 1024 {
|
|
return apierrors.Unauthorized("IDEMPOTENCY_SCOPE_REQUIRED", "authenticated scope is required")
|
|
}
|
|
if r.ContentLength > i.config.MaxRequestBytes {
|
|
return apierrors.New(413, "BODY_TOO_LARGE", "request body too large")
|
|
}
|
|
var body []byte
|
|
if r.Body != nil {
|
|
body, err = io.ReadAll(io.LimitReader(r.Body, i.config.MaxRequestBytes+1))
|
|
if err != nil {
|
|
return apierrors.BadRequest("INVALID_BODY", "unable to read request body").WithCause(err)
|
|
}
|
|
if int64(len(body)) > i.config.MaxRequestBytes {
|
|
return apierrors.New(413, "BODY_TOO_LARGE", "request body too large")
|
|
}
|
|
}
|
|
if err := c.Context().Err(); err != nil {
|
|
return err
|
|
}
|
|
originalBody := r.Body
|
|
r.Body = io.NopCloser(bytes.NewReader(body))
|
|
defer func() { r.Body = originalBody }()
|
|
parts, _ := json.Marshal([]string{r.Method, r.URL.RequestURI(), r.Header.Get("Content-Type")})
|
|
hash := sha256.New()
|
|
hash.Write(parts)
|
|
hash.Write(body)
|
|
var fingerprint [32]byte
|
|
copy(fingerprint[:], hash.Sum(nil))
|
|
namespace, _ := json.Marshal([]string{scope, key})
|
|
id := sha256.Sum256(namespace)
|
|
if i.config.Store != nil {
|
|
return i.shared(c, next, id, fingerprint)
|
|
}
|
|
i.mu.Lock()
|
|
now := i.now()
|
|
for k, v := range i.items {
|
|
if !v.running && !now.Before(v.expires) {
|
|
delete(i.items, k)
|
|
i.bytes -= v.size
|
|
}
|
|
}
|
|
if saved := i.items[id]; saved != nil {
|
|
if saved.fingerprint != fingerprint {
|
|
i.mu.Unlock()
|
|
return apierrors.New(409, "IDEMPOTENCY_CONFLICT", "key was used with a different request")
|
|
}
|
|
if saved.running || saved.failed {
|
|
i.mu.Unlock()
|
|
return apierrors.New(409, "IDEMPOTENCY_UNAVAILABLE", "request is in progress or its outcome cannot be replayed")
|
|
}
|
|
status, headers, data := saved.status, saved.headers.Clone(), append([]byte(nil), saved.body...)
|
|
i.mu.Unlock()
|
|
headers.Set("Idempotency-Replayed", "true")
|
|
return c.WriteResponse(status, headers, data)
|
|
}
|
|
if len(i.items) >= i.config.MaxEntries {
|
|
i.mu.Unlock()
|
|
return apierrors.New(503, "IDEMPOTENCY_BUSY", "idempotency capacity reached")
|
|
}
|
|
saved := &replay{fingerprint: fingerprint, running: true}
|
|
i.items[id] = saved
|
|
i.mu.Unlock()
|
|
buffer := &boundedResponse{headers: c.Raw().Header().Clone(), limit: i.config.MaxResponseBytes}
|
|
// Always retain an outcome tombstone on panic/error/overflow. Deleting
|
|
// the key could execute a partially committed business operation twice.
|
|
defer func() {
|
|
i.mu.Lock()
|
|
defer i.mu.Unlock()
|
|
saved.running = false
|
|
saved.failed = true
|
|
saved.expires = i.now().Add(i.config.TTL)
|
|
size := buffer.body.Len()
|
|
for k, values := range buffer.headers {
|
|
size += len(k)
|
|
for _, v := range values {
|
|
size += len(v)
|
|
}
|
|
}
|
|
if buffer.complete && buffer.err == nil && buffer.status < 500 && buffer.headers.Get("Set-Cookie") == "" && !strings.Contains(strings.ToLower(buffer.headers.Get("Cache-Control")), "no-store") && size <= i.config.MaxBytes-i.bytes {
|
|
saved.failed = false
|
|
saved.status = buffer.status
|
|
saved.headers = buffer.headers.Clone()
|
|
saved.body = append([]byte(nil), buffer.body.Bytes()...)
|
|
saved.size = size
|
|
i.bytes += size
|
|
}
|
|
}()
|
|
child := c.Clone(c.Context(), buffer)
|
|
err = next(child)
|
|
// Keep an ambiguous outcome tombstone, but let an outer Timeout
|
|
// middleware turn deadline expiry into 504 before committing bytes.
|
|
if ctxErr := child.Context().Err(); ctxErr != nil {
|
|
return ctxErr
|
|
}
|
|
if err != nil {
|
|
child.WriteError(err)
|
|
}
|
|
if buffer.err != nil {
|
|
return apierrors.Internal(buffer.err)
|
|
}
|
|
if buffer.status == 0 {
|
|
buffer.status = 200
|
|
}
|
|
buffer.complete = true
|
|
return c.WriteResponse(buffer.status, buffer.headers, buffer.body.Bytes())
|
|
}
|
|
}
|
|
}
|
|
|
|
// Intentionally lacks Flush/Hijack: only buffered, bounded endpoints are valid.
|
|
type boundedResponse struct {
|
|
headers http.Header
|
|
status, limit int
|
|
body bytes.Buffer
|
|
err error
|
|
complete bool
|
|
}
|
|
|
|
func (b *boundedResponse) Header() http.Header { return b.headers }
|
|
func (b *boundedResponse) WriteHeader(status int) {
|
|
if status < 200 {
|
|
b.err = errors.New("idempotency: informational/upgrade responses are unsupported")
|
|
return
|
|
}
|
|
if b.status == 0 {
|
|
b.status = status
|
|
}
|
|
}
|
|
func (b *boundedResponse) Write(data []byte) (int, error) {
|
|
if b.status == 0 {
|
|
b.status = 200
|
|
}
|
|
if b.err != nil {
|
|
return 0, b.err
|
|
}
|
|
if len(data) > b.limit-b.body.Len() {
|
|
b.err = fmt.Errorf("idempotency: response exceeds %d bytes", b.limit)
|
|
return 0, b.err
|
|
}
|
|
return b.body.Write(data)
|
|
}
|