feat: add cache, outbound HTTP, request protection, tasks and module scaffolding
This commit is contained in:
235
middleware/idempotency.go
Normal file
235
middleware/idempotency.go
Normal file
@@ -0,0 +1,235 @@
|
||||
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 {
|
||||
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.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)
|
||||
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)
|
||||
}
|
||||
227
middleware/protection_test.go
Normal file
227
middleware/protection_test.go
Normal file
@@ -0,0 +1,227 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.yuntuo.hk/k2safe/go-omni/httpx"
|
||||
)
|
||||
|
||||
func invoke(h httpx.Handler, method, body, scope, key string) *httptest.ResponseRecorder {
|
||||
r := httptest.NewRequest(method, "/records?version=1", strings.NewReader(body))
|
||||
r.Header.Set("Content-Type", "application/json")
|
||||
if key != "" {
|
||||
r.Header.Set("Idempotency-Key", key)
|
||||
}
|
||||
w := httptest.NewRecorder()
|
||||
c := httpx.NewContext(r, w, nil)
|
||||
c.Set("verifiedIdentity", scope)
|
||||
if err := h(c); err != nil {
|
||||
c.WriteError(err)
|
||||
}
|
||||
return w
|
||||
}
|
||||
|
||||
func identity(c *httpx.Context) (string, error) {
|
||||
v, _ := c.Get("verifiedIdentity")
|
||||
s, _ := v.(string)
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func TestRateLimitRefillCapacityAndIdleExpiry(t *testing.T) {
|
||||
l, err := NewRateLimiter(RateLimitConfig{Rate: 2, Burst: 2, MaxKeys: 1, IdleTTL: time.Second, Key: func(c *httpx.Context) string { s, _ := identity(c); return s }})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
now := time.Unix(100, 0)
|
||||
l.now = func() time.Time { return now }
|
||||
h := l.Middleware()(func(c *httpx.Context) error { return c.NoContent() })
|
||||
for _, status := range []int{204, 204, 429} {
|
||||
w := invoke(h, "GET", "", "a", "")
|
||||
if w.Code != status {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
if status == 429 && w.Header().Get("Retry-After") != "1" {
|
||||
t.Fatal(w.Header())
|
||||
}
|
||||
}
|
||||
if w := invoke(h, "GET", "", "b", ""); w.Code != 503 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
now = now.Add(500 * time.Millisecond)
|
||||
if w := invoke(h, "GET", "", "a", ""); w.Code != 204 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
now = now.Add(time.Second)
|
||||
if w := invoke(h, "GET", "", "b", ""); w.Code != 204 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
if _, err := NewRateLimiter(RateLimitConfig{Rate: 1, Burst: 2, IdleTTL: time.Second}); err == nil {
|
||||
t.Fatal("early bucket expiry accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimitConcurrentBurst(t *testing.T) {
|
||||
l, _ := NewRateLimiter(RateLimitConfig{Rate: 1, Burst: 7})
|
||||
l.now = func() time.Time { return time.Unix(100, 0) }
|
||||
var allowed atomic.Int32
|
||||
var wg sync.WaitGroup
|
||||
for n := 0; n < 100; n++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
delay, full := l.allow("same")
|
||||
if delay == 0 && !full {
|
||||
allowed.Add(1)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
if allowed.Load() != 7 {
|
||||
t.Fatal(allowed.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestTimeoutCooperativeAndCommittedResponse(t *testing.T) {
|
||||
for _, committed := range []bool{false, true} {
|
||||
h := Timeout(5 * time.Millisecond)(func(c *httpx.Context) error {
|
||||
if committed {
|
||||
_ = c.Accepted("accepted")
|
||||
}
|
||||
<-c.Context().Done()
|
||||
return c.Context().Err()
|
||||
})
|
||||
w := invoke(h, "GET", "", "", "")
|
||||
want := 504
|
||||
if committed {
|
||||
want = 202
|
||||
}
|
||||
if w.Code != want {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestIdempotencyReplayScopeConflictAndExpiry(t *testing.T) {
|
||||
i, err := NewIdempotency(IdempotencyConfig{Scope: identity, TTL: time.Minute})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
now := time.Unix(100, 0)
|
||||
i.now = func() time.Time { return now }
|
||||
calls := 0
|
||||
h := i.Middleware()(func(c *httpx.Context) error { calls++; return c.Created(map[string]int{"calls": calls}) })
|
||||
first := invoke(h, "POST", `{"name":"one"}`, "alice", "key1")
|
||||
replay := invoke(h, "POST", `{"name":"one"}`, "alice", "key1")
|
||||
if first.Code != 201 || replay.Code != 201 || first.Body.String() != replay.Body.String() || replay.Header().Get("Idempotency-Replayed") != "true" || calls != 1 {
|
||||
t.Fatal(first, replay, calls)
|
||||
}
|
||||
if w := invoke(h, "POST", `{"name":"two"}`, "alice", "key1"); w.Code != 409 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
if w := invoke(h, "POST", `{"name":"one"}`, "bob", "key1"); w.Code != 201 || calls != 2 {
|
||||
t.Fatal(w.Code, calls)
|
||||
}
|
||||
if w := invoke(h, "POST", `{}`, "", "key1"); w.Code != 401 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
if w := invoke(h, "POST", `{}`, "alice", ""); w.Code != 400 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
now = now.Add(time.Minute)
|
||||
if w := invoke(h, "POST", `{"name":"one"}`, "alice", "key1"); w.Code != 201 || calls != 3 {
|
||||
t.Fatal(w.Code, calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIdempotencyActiveClaimNeverExpires(t *testing.T) {
|
||||
i, _ := NewIdempotency(IdempotencyConfig{Scope: identity, TTL: time.Nanosecond, MaxInFlight: 2})
|
||||
started, release := make(chan struct{}), make(chan struct{})
|
||||
var calls atomic.Int32
|
||||
h := i.Middleware()(func(c *httpx.Context) error { calls.Add(1); close(started); <-release; return c.Created("saved") })
|
||||
done := make(chan *httptest.ResponseRecorder, 1)
|
||||
go func() { done <- invoke(h, "POST", "{}", "alice", "key") }()
|
||||
<-started
|
||||
for n := 0; n < 30; n++ {
|
||||
if w := invoke(h, "POST", "{}", "alice", "key"); w.Code != 409 {
|
||||
t.Errorf("duplicate status %d", w.Code)
|
||||
}
|
||||
}
|
||||
close(release)
|
||||
if w := <-done; w.Code != 201 || calls.Load() != 1 {
|
||||
t.Fatal(w.Code, calls.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestIdempotencyBoundsAndAmbiguousOutcomes(t *testing.T) {
|
||||
for _, mode := range []string{"panic", "error", "overflow", "cookie", "noStore", "memory"} {
|
||||
t.Run(mode, func(t *testing.T) {
|
||||
cfg := IdempotencyConfig{Scope: identity, MaxRequestBytes: 8, MaxResponseBytes: 128, MaxEntries: 1}
|
||||
if mode == "memory" {
|
||||
cfg.MaxBytes = 1
|
||||
}
|
||||
i, _ := NewIdempotency(cfg)
|
||||
calls := 0
|
||||
h := Recovery()(i.Middleware()(func(c *httpx.Context) error {
|
||||
calls++
|
||||
switch mode {
|
||||
case "panic":
|
||||
panic("failure")
|
||||
case "error":
|
||||
return errors.New("failed")
|
||||
case "overflow":
|
||||
return c.OK(strings.Repeat("x", 200))
|
||||
case "cookie":
|
||||
c.SetResponseHeader("Set-Cookie", "session=secret")
|
||||
case "noStore":
|
||||
c.SetResponseHeader("Cache-Control", "no-store")
|
||||
}
|
||||
return c.Created("ok")
|
||||
}))
|
||||
if w := invoke(h, "POST", "123456789", "alice", "key"); w.Code != 413 || calls != 0 {
|
||||
t.Fatal(w.Code, calls)
|
||||
}
|
||||
first := invoke(h, "POST", "{}", "alice", "key")
|
||||
if mode == "panic" || mode == "error" || mode == "overflow" {
|
||||
if first.Code != 500 {
|
||||
t.Fatal(first.Code)
|
||||
}
|
||||
} else if first.Code != 201 {
|
||||
t.Fatal(first.Code)
|
||||
}
|
||||
if w := invoke(h, "POST", "{}", "alice", "key"); w.Code != 409 || calls != 1 {
|
||||
t.Fatal(w.Code, calls)
|
||||
}
|
||||
if w := invoke(h, "POST", "{}", "alice", "newKey"); w.Code != 503 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestIdempotencyInFlightLimitAndTimeoutComposition(t *testing.T) {
|
||||
i, _ := NewIdempotency(IdempotencyConfig{Scope: identity, MaxInFlight: 1})
|
||||
started, release := make(chan struct{}), make(chan struct{})
|
||||
h := i.Middleware()(func(c *httpx.Context) error { close(started); <-release; return c.NoContent() })
|
||||
done := make(chan struct{})
|
||||
go func() { defer close(done); invoke(h, "POST", "{}", "alice", "one") }()
|
||||
<-started
|
||||
if w := invoke(h, "POST", "{}", "alice", "two"); w.Code != 503 {
|
||||
t.Error(w.Code)
|
||||
}
|
||||
close(release)
|
||||
<-done
|
||||
timed := Timeout(5 * time.Millisecond)(i.Middleware()(func(c *httpx.Context) error { <-c.Context().Done(); return context.DeadlineExceeded }))
|
||||
if w := invoke(timed, "POST", "{}", "alice", "timeout"); w.Code != 504 {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
if w := invoke(timed, "POST", "{}", "alice", "timeout"); w.Code != 409 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
}
|
||||
109
middleware/rateLimit.go
Normal file
109
middleware/rateLimit.go
Normal file
@@ -0,0 +1,109 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"container/list"
|
||||
"errors"
|
||||
"math"
|
||||
"strconv"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
apierrors "git.yuntuo.hk/k2safe/go-omni/errors"
|
||||
"git.yuntuo.hk/k2safe/go-omni/httpx"
|
||||
)
|
||||
|
||||
type RateLimitConfig struct {
|
||||
Rate float64 `json:"rate"` // Tokens per second, per key.
|
||||
Burst int `json:"burst"`
|
||||
MaxKeys int `json:"maxKeys"`
|
||||
IdleTTL time.Duration `json:"idleTtl"`
|
||||
Key func(*httpx.Context) string `json:"-"` // Default: trusted ClientIP.
|
||||
}
|
||||
|
||||
type bucket struct {
|
||||
key string
|
||||
tokens float64
|
||||
updated time.Time
|
||||
}
|
||||
type RateLimiter struct {
|
||||
mu sync.Mutex
|
||||
config RateLimitConfig
|
||||
items map[string]*list.Element
|
||||
lru *list.List
|
||||
now func() time.Time
|
||||
}
|
||||
|
||||
func NewRateLimiter(cfg RateLimitConfig) (*RateLimiter, error) {
|
||||
if math.IsNaN(cfg.Rate) || math.IsInf(cfg.Rate, 0) || cfg.Rate <= 0 || cfg.Burst < 1 || cfg.MaxKeys < 0 || cfg.IdleTTL < 0 {
|
||||
return nil, errors.New("middleware: invalid rate limit configuration")
|
||||
}
|
||||
refill := float64(cfg.Burst) / cfg.Rate
|
||||
if refill > (24 * time.Hour).Seconds() {
|
||||
return nil, errors.New("middleware: bucket refill time must not exceed 24h")
|
||||
}
|
||||
minimum := time.Duration(math.Ceil(refill * float64(time.Second)))
|
||||
if cfg.IdleTTL == 0 {
|
||||
cfg.IdleTTL = max(5*time.Minute, minimum)
|
||||
}
|
||||
if cfg.IdleTTL < minimum {
|
||||
return nil, errors.New("middleware: idleTtl must allow a full bucket refill")
|
||||
}
|
||||
if cfg.MaxKeys == 0 {
|
||||
cfg.MaxKeys = 10000
|
||||
}
|
||||
if cfg.Key == nil {
|
||||
cfg.Key = func(c *httpx.Context) string { return c.ClientIP() }
|
||||
}
|
||||
return &RateLimiter{config: cfg, items: make(map[string]*list.Element), lru: list.New(), now: time.Now}, nil
|
||||
}
|
||||
|
||||
func (l *RateLimiter) Middleware() httpx.Middleware {
|
||||
return func(next httpx.Handler) httpx.Handler {
|
||||
return func(c *httpx.Context) error {
|
||||
key := l.config.Key(c)
|
||||
if key == "" || len(key) > 1024 {
|
||||
return apierrors.BadRequest("INVALID_RATE_KEY", "invalid rate limit key")
|
||||
}
|
||||
delay, full := l.allow(key)
|
||||
if full {
|
||||
return apierrors.New(503, "RATE_LIMIT_CAPACITY", "rate limiter capacity reached")
|
||||
}
|
||||
if delay > 0 {
|
||||
c.SetResponseHeader("Retry-After", strconv.Itoa(max(1, int(math.Ceil(delay.Seconds())))))
|
||||
return apierrors.New(429, "RATE_LIMITED", "too many requests")
|
||||
}
|
||||
return next(c)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (l *RateLimiter) allow(key string) (time.Duration, bool) {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
now := l.now()
|
||||
e := l.items[key]
|
||||
if e == nil {
|
||||
for l.lru.Len() > 0 {
|
||||
old := l.lru.Back()
|
||||
if now.Sub(old.Value.(*bucket).updated) < l.config.IdleTTL {
|
||||
break
|
||||
}
|
||||
delete(l.items, old.Value.(*bucket).key)
|
||||
l.lru.Remove(old)
|
||||
}
|
||||
if len(l.items) >= l.config.MaxKeys {
|
||||
return 0, true
|
||||
}
|
||||
e = l.lru.PushFront(&bucket{key: key, tokens: float64(l.config.Burst), updated: now})
|
||||
l.items[key] = e
|
||||
}
|
||||
b := e.Value.(*bucket)
|
||||
b.tokens = min(float64(l.config.Burst), b.tokens+max(0, now.Sub(b.updated).Seconds())*l.config.Rate)
|
||||
b.updated = now
|
||||
l.lru.MoveToFront(e)
|
||||
if b.tokens >= 1 {
|
||||
b.tokens--
|
||||
return 0, false
|
||||
}
|
||||
return time.Duration(math.Ceil((1 - b.tokens) / l.config.Rate * float64(time.Second))), false
|
||||
}
|
||||
35
middleware/timeout.go
Normal file
35
middleware/timeout.go
Normal file
@@ -0,0 +1,35 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
apierrors "git.yuntuo.hk/k2safe/go-omni/errors"
|
||||
"git.yuntuo.hk/k2safe/go-omni/httpx"
|
||||
)
|
||||
|
||||
// Timeout applies a cooperative deadline to handlers and downstream calls.
|
||||
// It does not spawn a goroutine or race a handler's response. Business code must
|
||||
// observe Context().Done(); an already committed response cannot be replaced.
|
||||
func Timeout(duration time.Duration) httpx.Middleware {
|
||||
if duration <= 0 {
|
||||
panic("middleware: timeout must be positive")
|
||||
}
|
||||
return func(next httpx.Handler) httpx.Handler {
|
||||
return func(c *httpx.Context) error {
|
||||
parent := c.Context()
|
||||
ctx, cancel := context.WithTimeout(parent, duration)
|
||||
defer cancel()
|
||||
c.WithContext(ctx)
|
||||
defer c.WithContext(parent)
|
||||
if ctx.Err() != nil {
|
||||
return apierrors.New(504, "REQUEST_TIMEOUT", "request deadline exceeded")
|
||||
}
|
||||
err := next(c)
|
||||
if ctx.Err() == context.DeadlineExceeded && !c.Responded() {
|
||||
return apierrors.New(504, "REQUEST_TIMEOUT", "request deadline exceeded").WithCause(err)
|
||||
}
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user