feat: add cache, outbound HTTP, request protection, tasks and module scaffolding
This commit is contained in:
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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user