226 lines
7.7 KiB
Go
226 lines
7.7 KiB
Go
package redisadmission
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/binary"
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
|
|
"proxy-pool/internal/platform/admission"
|
|
)
|
|
|
|
type scriptCall struct {
|
|
keys []string
|
|
args []any
|
|
}
|
|
|
|
type scriptClient struct {
|
|
redis.Scripter
|
|
result any
|
|
err error
|
|
calls []scriptCall
|
|
}
|
|
|
|
func (client *scriptClient) EvalSha(_ context.Context, _ string, keys []string, args ...any) *redis.Cmd {
|
|
client.calls = append(client.calls, scriptCall{keys: append([]string(nil), keys...), args: append([]any(nil), args...)})
|
|
return redis.NewCmdResult(client.result, client.err)
|
|
}
|
|
|
|
func TestNewRejectsInvalidDependenciesAndOptions(t *testing.T) {
|
|
t.Parallel()
|
|
validClient := &scriptClient{result: "ok"}
|
|
valid := Options{Namespace: "listener-a", Window: time.Minute, Global: 10, PerKey: 2}
|
|
var typedNilClient *redis.Client
|
|
|
|
tests := []struct {
|
|
name string
|
|
client redis.Scripter
|
|
options Options
|
|
}{
|
|
{name: "nil client", options: valid},
|
|
{name: "typed nil client", client: typedNilClient, options: valid},
|
|
{name: "empty namespace", client: validClient, options: withNamespace(valid, "")},
|
|
{name: "leading namespace whitespace", client: validClient, options: withNamespace(valid, " listener-a")},
|
|
{name: "trailing namespace whitespace", client: validClient, options: withNamespace(valid, "listener-a ")},
|
|
{name: "cluster tag in namespace", client: validClient, options: withNamespace(valid, "listener{a}")},
|
|
{name: "separator in namespace", client: validClient, options: withNamespace(valid, "listener:a")},
|
|
{name: "zero window", client: validClient, options: withWindow(valid, 0)},
|
|
{name: "sub-millisecond window", client: validClient, options: withWindow(valid, time.Microsecond)},
|
|
{name: "fractional millisecond window", client: validClient, options: withWindow(valid, time.Millisecond+time.Microsecond)},
|
|
{name: "negative global", client: validClient, options: withGlobal(valid, -1)},
|
|
{name: "negative per-key", client: validClient, options: withPerKey(valid, -1)},
|
|
{name: "zero limits", client: validClient, options: withLimits(valid, 0, 0)},
|
|
{name: "global exceeds Lua integer", client: validClient, options: withGlobal(valid, MaximumLuaInteger+1)},
|
|
{name: "per-key exceeds Lua integer", client: validClient, options: withPerKey(valid, MaximumLuaInteger+1)},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
if adapter, err := New(tt.client, tt.options); !errors.Is(err, ErrInvalidOptions) || adapter != nil {
|
|
t.Fatalf("New() = (%v, %v), want nil adapter and ErrInvalidOptions", adapter, err)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAdmitMapsScriptStatusAndPreservesInputErrors(t *testing.T) {
|
|
t.Parallel()
|
|
tests := []struct {
|
|
name string
|
|
status any
|
|
want error
|
|
}{
|
|
{name: "accepted", status: "ok"},
|
|
{name: "global limit", status: "global", want: admission.ErrGlobalLimit},
|
|
{name: "per-key limit", status: "per_key", want: admission.ErrPerKeyLimit},
|
|
{name: "unexpected status", status: "unknown", want: admission.ErrUnavailable},
|
|
{name: "unexpected reply type", status: int64(1), want: admission.ErrUnavailable},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
client := &scriptClient{result: tt.status}
|
|
adapter := mustAdapter(t, client, Options{Namespace: "listener-a", Window: time.Minute, Global: 10, PerKey: 2})
|
|
err := adapter.Admit(context.Background(), "client-a")
|
|
if !errors.Is(err, tt.want) {
|
|
t.Fatalf("Admit() error = %v, want %v", err, tt.want)
|
|
}
|
|
})
|
|
}
|
|
|
|
adapter := mustAdapter(t, &scriptClient{result: "ok"}, Options{Namespace: "listener-a", Window: time.Minute, Global: 1})
|
|
if err := adapter.Admit(nil, "client-a"); !errors.Is(err, admission.ErrInvalidIdentity) {
|
|
t.Fatalf("Admit(nil context) error = %v, want ErrInvalidIdentity", err)
|
|
}
|
|
if err := adapter.Admit(context.Background(), ""); !errors.Is(err, admission.ErrInvalidIdentity) {
|
|
t.Fatalf("Admit(empty identity) error = %v, want ErrInvalidIdentity", err)
|
|
}
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
if err := adapter.Admit(ctx, "client-a"); !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("Admit(canceled context) error = %v, want context.Canceled", err)
|
|
}
|
|
var nilAdapter *Adapter
|
|
if err := nilAdapter.Admit(context.Background(), "client-a"); !errors.Is(err, admission.ErrUnavailable) {
|
|
t.Fatalf("nil Adapter.Admit() error = %v, want ErrUnavailable", err)
|
|
}
|
|
}
|
|
|
|
func TestAdmitUsesHashedIdentityAndExactIntegerArguments(t *testing.T) {
|
|
t.Parallel()
|
|
client := &scriptClient{result: "ok"}
|
|
adapter := mustAdapter(t, client, Options{
|
|
Namespace: "listener-a", Window: 1500 * time.Millisecond,
|
|
Global: MaximumLuaInteger, PerKey: MaximumLuaInteger - 1,
|
|
})
|
|
const identity = "private-client@example.test"
|
|
if err := adapter.Admit(context.Background(), identity); err != nil {
|
|
t.Fatalf("Admit(): %v", err)
|
|
}
|
|
if len(client.calls) != 1 {
|
|
t.Fatalf("script calls = %d, want 1", len(client.calls))
|
|
}
|
|
call := client.calls[0]
|
|
if len(call.keys) != 1 || call.keys[0] != "pp:{admission}:listener-a:window" {
|
|
t.Fatalf("script keys = %q", call.keys)
|
|
}
|
|
wantArgs := []string{"1500", fmt.Sprint(MaximumLuaInteger), fmt.Sprint(MaximumLuaInteger - 1), clientField(identity)}
|
|
if got := stringify(call.args); fmt.Sprint(got) != fmt.Sprint(wantArgs) {
|
|
t.Fatalf("script args = %q, want %q", got, wantArgs)
|
|
}
|
|
serialized := fmt.Sprint(call.keys, call.args)
|
|
if strings.Contains(serialized, identity) {
|
|
t.Fatalf("Redis input leaks raw identity: %s", serialized)
|
|
}
|
|
if got := clientField(identity); len(got) != len("client:")+sha256.Size*2 || !strings.HasPrefix(got, "client:") {
|
|
t.Fatalf("client field = %q", got)
|
|
}
|
|
}
|
|
|
|
func TestAdmitFailsClosedWithoutLeakingIdentity(t *testing.T) {
|
|
t.Parallel()
|
|
const identity = "sensitive-client"
|
|
client := &scriptClient{err: errors.New("backend failure: " + identity)}
|
|
adapter := mustAdapter(t, client, Options{Namespace: "listener-a", Window: time.Minute, Global: 1})
|
|
err := adapter.Admit(context.Background(), identity)
|
|
if !errors.Is(err, admission.ErrUnavailable) {
|
|
t.Fatalf("Admit() error = %v, want ErrUnavailable", err)
|
|
}
|
|
if strings.Contains(err.Error(), identity) {
|
|
t.Fatalf("Admit() error leaks identity: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestAdmitFailsClosedForUninitializedAdapter(t *testing.T) {
|
|
t.Parallel()
|
|
var typedNilClient *redis.Client
|
|
tests := []*Adapter{
|
|
{},
|
|
{client: typedNilClient, key: "pp:{admission}:test:window", windowMillis: 60_000, global: 1},
|
|
}
|
|
for _, adapter := range tests {
|
|
if err := adapter.Admit(context.Background(), "client-a"); !errors.Is(err, admission.ErrUnavailable) {
|
|
t.Fatalf("Admit() error = %v, want ErrUnavailable", err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func clientField(identity string) string {
|
|
var size [8]byte
|
|
binary.BigEndian.PutUint64(size[:], uint64(len(identity)))
|
|
digest := sha256.New()
|
|
_, _ = digest.Write(size[:])
|
|
_, _ = digest.Write([]byte(identity))
|
|
return fmt.Sprintf("client:%x", digest.Sum(nil))
|
|
}
|
|
|
|
func stringify(values []any) []string {
|
|
result := make([]string, len(values))
|
|
for index, value := range values {
|
|
result[index] = fmt.Sprint(value)
|
|
}
|
|
return result
|
|
}
|
|
|
|
func mustAdapter(t *testing.T, client redis.Scripter, options Options) *Adapter {
|
|
t.Helper()
|
|
adapter, err := New(client, options)
|
|
if err != nil {
|
|
t.Fatalf("New(): %v", err)
|
|
}
|
|
return adapter
|
|
}
|
|
|
|
func withNamespace(options Options, namespace string) Options {
|
|
options.Namespace = namespace
|
|
return options
|
|
}
|
|
|
|
func withWindow(options Options, window time.Duration) Options {
|
|
options.Window = window
|
|
return options
|
|
}
|
|
|
|
func withGlobal(options Options, limit int64) Options {
|
|
options.Global = limit
|
|
return options
|
|
}
|
|
|
|
func withPerKey(options Options, limit int64) Options {
|
|
options.PerKey = limit
|
|
return options
|
|
}
|
|
|
|
func withLimits(options Options, global, perKey int64) Options {
|
|
options.Global = global
|
|
options.PerKey = perKey
|
|
return options
|
|
}
|