228 lines
6.8 KiB
Go
228 lines
6.8 KiB
Go
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)
|
|
}
|
|
}
|