mirror of
https://github.com/lukaszraczylo/graphql-monitoring-proxy.git
synced 2026-11-04 08:04:10 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e3066ad7f0 | ||
|
|
e138b44d00 | ||
|
|
3a6403c30b | ||
|
|
cff11f2b25 | ||
|
|
4c5d01aef3 | ||
|
|
6e7e2ecfae | ||
|
|
3e187790f2 | ||
|
|
ea30694dc9 | ||
|
|
c128a0f4d9 |
@@ -821,7 +821,7 @@ Enable rate limiting based on user roles using the `ROLE_RATE_LIMIT` environment
|
||||
- `ROLE_RATE_LIMIT`: Enable role-based rate limiting (default: `false`)
|
||||
|
||||
**Features:**
|
||||
- **Burst control**: Optional burst limits for handling traffic spikes
|
||||
- **Burst control**: Optional `burst` caps requests in a short window (one tenth of the interval). `req` counts requests per `interval`. A `burst` at or above `req` has no effect.
|
||||
|
||||
`ratelimit.json` is loaded once at startup; changing the file requires a restart.
|
||||
|
||||
|
||||
@@ -143,8 +143,13 @@ func periodicallyReloadBannedUsers(ctx context.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
func checkIfUserIsBanned(c fiber.Ctx, userID string) bool {
|
||||
func isUserBanned(userID string) bool {
|
||||
_, found := bannedUsersIDs.Load(userID)
|
||||
return found
|
||||
}
|
||||
|
||||
func checkIfUserIsBanned(c fiber.Ctx, userID string) bool {
|
||||
found := isUserBanned(userID)
|
||||
|
||||
cfg.Logger.Debug(&libpack_logger.LogMessage{
|
||||
Message: "Checking if user is banned",
|
||||
|
||||
+1
-1
@@ -118,7 +118,7 @@ func (bhm *BackendHealthManager) WaitForBackendReady(timeout time.Duration) erro
|
||||
}
|
||||
|
||||
// Exponential backoff (1.5x growth, capped at maxDelay), no jitter
|
||||
time.Sleep(currentDelay)
|
||||
time.Sleep(min(currentDelay, time.Until(deadline)))
|
||||
currentDelay = time.Duration(float64(currentDelay) * 1.5)
|
||||
if currentDelay > maxDelay {
|
||||
currentDelay = maxDelay
|
||||
|
||||
@@ -0,0 +1,32 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"net"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
libpack_logger "github.com/lukaszraczylo/graphql-monitoring-proxy/logging"
|
||||
"github.com/valyala/fasthttp"
|
||||
)
|
||||
|
||||
// Pins theorems/Theorems/ReadyBackoff.lean: the startup wait never sleeps past its timeout.
|
||||
func TestWaitForBackendReady_HonoursTimeout(t *testing.T) {
|
||||
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
addr := ln.Addr().String()
|
||||
_ = ln.Close() // closed port: every probe is refused immediately
|
||||
|
||||
bhm := NewBackendHealthManager(&fasthttp.Client{}, "http://"+addr, "", libpack_logger.New())
|
||||
defer bhm.cancel()
|
||||
|
||||
const timeout = 3 * time.Second
|
||||
start := time.Now()
|
||||
if err := bhm.WaitForBackendReady(timeout); err == nil {
|
||||
t.Fatal("want error for unreachable backend")
|
||||
}
|
||||
if elapsed := time.Since(start); elapsed > timeout+500*time.Millisecond {
|
||||
t.Fatalf("waited %v, want <= %v", elapsed, timeout+500*time.Millisecond)
|
||||
}
|
||||
}
|
||||
Vendored
+2
-1
@@ -158,7 +158,8 @@ func (c *Cache) Set(key string, value []byte, ttl time.Duration) {
|
||||
c.evictToFreeMemory(memoryToFree)
|
||||
} else if atomic.LoadInt64(&c.entryCount) >= c.maxCacheSize {
|
||||
// Fall back to count-based eviction for backward compatibility
|
||||
c.evictOldest(int(c.maxCacheSize / 10)) // Evict 10% of entries
|
||||
// 10% of entries, at least one: max/10 is 0 for caps below 10.
|
||||
c.evictOldest(int(max(1, c.maxCacheSize/10)))
|
||||
}
|
||||
|
||||
// A non-positive TTL is clamped to defaultTTL so this backend expires
|
||||
|
||||
Vendored
+23
@@ -0,0 +1,23 @@
|
||||
package libpack_cache_memory
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
// Pins theorems/Theorems/MemCacheCap.lean: every cap >= 1 bounds the entry count.
|
||||
func TestSet_SmallEntryCapIsEnforced(t *testing.T) {
|
||||
for _, maxEntries := range []int64{1, 2, 5, 9, 10, 25} {
|
||||
t.Run(fmt.Sprintf("cap_%d", maxEntries), func(t *testing.T) {
|
||||
c := NewWithSize(time.Minute, DefaultMaxMemorySize, maxEntries)
|
||||
defer c.Shutdown()
|
||||
for i := 0; i < 100; i++ {
|
||||
c.Set(fmt.Sprintf("k%d", i), []byte("v"), time.Minute)
|
||||
assert.LessOrEqual(t, c.CountQueries(), maxEntries)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Vendored
+9
-1
@@ -83,7 +83,15 @@ type RedisClientConfig struct {
|
||||
// entry.
|
||||
const defaultTTL = 60 * time.Second
|
||||
|
||||
// DefaultKeyPrefix namespaces cache keys when no prefix is configured, so Clear
|
||||
// and CountQueries never touch keys of other applications on a shared Redis.
|
||||
const DefaultKeyPrefix = "gmp:"
|
||||
|
||||
func New(redisClientConfig *RedisClientConfig) (*RedisConfig, error) {
|
||||
prefix := redisClientConfig.Prefix
|
||||
if prefix == "" {
|
||||
prefix = DefaultKeyPrefix
|
||||
}
|
||||
c := &RedisConfig{
|
||||
client: redis.NewClient(&redis.Options{
|
||||
Addr: redisClientConfig.RedisServer,
|
||||
@@ -91,7 +99,7 @@ func New(redisClientConfig *RedisClientConfig) (*RedisConfig, error) {
|
||||
DB: redisClientConfig.RedisDB,
|
||||
}),
|
||||
ctx: context.Background(),
|
||||
prefix: redisClientConfig.Prefix,
|
||||
prefix: prefix,
|
||||
builderPool: &sync.Pool{
|
||||
New: func() any {
|
||||
return &strings.Builder{}
|
||||
|
||||
+30
@@ -0,0 +1,30 @@
|
||||
package libpack_cache_redis
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// Pins theorems/Theorems/RedisClear.lean: with no configured prefix, Clear must still spare foreign keys.
|
||||
func TestClear_DefaultPrefixSparesForeignKeys(t *testing.T) {
|
||||
t.Parallel()
|
||||
s, err := miniredis.Run()
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(s.Close)
|
||||
|
||||
require.NoError(t, s.Set("other-app:session", "keep"))
|
||||
|
||||
c, err := New(&RedisClientConfig{RedisServer: s.Addr()})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, c.Set("q1", []byte("v"), 0))
|
||||
|
||||
require.NoError(t, c.Clear())
|
||||
|
||||
assert.True(t, s.Exists("other-app:session"), "foreign key was deleted")
|
||||
_, found, err := c.Get("q1")
|
||||
require.NoError(t, err)
|
||||
assert.False(t, found, "own key must be cleared")
|
||||
}
|
||||
Vendored
+4
-4
@@ -131,11 +131,11 @@ func (suite *RedisConfigSuite) TestGetKeys() {
|
||||
err = suite.redisConfig.Set("otherkey", []byte("othervalue"), ttl)
|
||||
assert.NoError(suite.T(), err)
|
||||
|
||||
keys, _ := suite.redisConfig.client.Keys(suite.redisConfig.ctx, "testkey*").Result()
|
||||
expectedKeys := []string{"testkey1", "testkey2"}
|
||||
keys, _ := suite.redisConfig.client.Keys(suite.redisConfig.ctx, DefaultKeyPrefix+"testkey*").Result()
|
||||
expectedKeys := []string{DefaultKeyPrefix + "testkey1", DefaultKeyPrefix + "testkey2"}
|
||||
assert.ElementsMatch(suite.T(), expectedKeys, keys)
|
||||
|
||||
suite.redisConfig.client.Del(suite.redisConfig.ctx, "testkey1", "testkey2", "otherkey")
|
||||
suite.redisConfig.client.Del(suite.redisConfig.ctx, DefaultKeyPrefix+"testkey1", DefaultKeyPrefix+"testkey2", DefaultKeyPrefix+"otherkey")
|
||||
}
|
||||
|
||||
func (suite *RedisConfigSuite) TestGetKeysCount() {
|
||||
@@ -154,5 +154,5 @@ func (suite *RedisConfigSuite) TestGetKeysCount() {
|
||||
assert.NoError(suite.T(), err)
|
||||
assert.Equal(suite.T(), int64(3), count3)
|
||||
|
||||
suite.redisConfig.client.Del(suite.redisConfig.ctx, "testkey1", "testkey2", "otherkey")
|
||||
suite.redisConfig.client.Del(suite.redisConfig.ctx, DefaultKeyPrefix+"testkey1", DefaultKeyPrefix+"testkey2", DefaultKeyPrefix+"otherkey")
|
||||
}
|
||||
|
||||
@@ -12,6 +12,9 @@ import (
|
||||
const (
|
||||
initialDelay = 60 * time.Second
|
||||
cleanupInterval = 1 * time.Hour
|
||||
|
||||
// minRetentionDays guards against a cutoff at or after NOW(), which would delete all event history.
|
||||
minRetentionDays = 1
|
||||
)
|
||||
|
||||
// Use parameterized queries to prevent SQL injection
|
||||
@@ -46,6 +49,10 @@ func enableHasuraEventCleaner(ctx context.Context) error {
|
||||
logger := cfg.Logger
|
||||
cfgMutex.RUnlock()
|
||||
|
||||
if clearOlderThan < minRetentionDays {
|
||||
return fmt.Errorf("HASURA_EVENT_CLEANER_OLDER_THAN must be at least %d day, got %d: event cleaner not started", minRetentionDays, clearOlderThan)
|
||||
}
|
||||
|
||||
logger.Info(&libpack_logger.LogMessage{
|
||||
Message: "Event cleaner enabled",
|
||||
Pairs: map[string]any{"interval_in_days": clearOlderThan},
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
libpack_logger "github.com/lukaszraczylo/graphql-monitoring-proxy/logging"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
// Pins theorems/Theorems/EventCleaner.lean: a retention below one day would delete all event history.
|
||||
func TestEnableHasuraEventCleaner_RejectsNonPositiveRetention(t *testing.T) {
|
||||
for _, days := range []int{0, -1, -30} {
|
||||
cfgMutex.Lock()
|
||||
cfg = &config{Logger: libpack_logger.New()}
|
||||
cfg.HasuraEventCleaner.Enable = true
|
||||
cfg.HasuraEventCleaner.EventMetadataDb = "postgres://fake:[email protected]:1/fake"
|
||||
cfg.HasuraEventCleaner.ClearOlderThan = days
|
||||
cfgMutex.Unlock()
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
err := enableHasuraEventCleaner(ctx)
|
||||
cancel()
|
||||
assert.Error(t, err, "days=%d", days)
|
||||
}
|
||||
}
|
||||
@@ -5,7 +5,7 @@ go 1.26.0
|
||||
toolchain go1.27.3
|
||||
|
||||
require (
|
||||
github.com/MicahParks/keyfunc/v3 v3.8.2
|
||||
github.com/MicahParks/keyfunc/v3 v3.9.1
|
||||
github.com/VictoriaMetrics/metrics v1.44.1
|
||||
github.com/alicebob/miniredis/v2 v2.40.0
|
||||
github.com/avast/retry-go/v5 v5.0.0
|
||||
@@ -20,7 +20,6 @@ require (
|
||||
github.com/graphql-go/graphql v0.8.1
|
||||
github.com/jackc/pgx/v5 v5.12.0
|
||||
github.com/lukaszraczylo/ask v0.0.0-20240916204100-6e9ef53a62d9
|
||||
github.com/lukaszraczylo/go-ratecounter v0.1.13
|
||||
github.com/lukaszraczylo/go-simple-graphql v1.2.89
|
||||
github.com/lukaszraczylo/oss-telemetry v0.2.5
|
||||
github.com/redis/go-redis/v9 v9.23.0
|
||||
@@ -36,7 +35,7 @@ require (
|
||||
)
|
||||
|
||||
require (
|
||||
github.com/MicahParks/jwkset v0.11.3 // indirect
|
||||
github.com/MicahParks/jwkset v0.11.4 // indirect
|
||||
github.com/avast/retry-go/v4 v4.7.0 // indirect
|
||||
github.com/cenkalti/backoff/v5 v5.0.3 // indirect
|
||||
github.com/cespare/xxhash/v2 v2.3.0 // indirect
|
||||
@@ -66,7 +65,7 @@ require (
|
||||
go.opentelemetry.io/otel/metric v1.47.0 // indirect
|
||||
go.opentelemetry.io/proto/otlp v1.11.1 // indirect
|
||||
go.uber.org/atomic v1.12.0 // indirect
|
||||
go.yaml.in/yaml/v3 v3.0.5 // indirect
|
||||
go.yaml.in/yaml/v3 v3.0.6 // indirect
|
||||
golang.org/x/crypto v0.58.0 // indirect
|
||||
golang.org/x/net v0.62.0 // indirect
|
||||
golang.org/x/sync v0.24.0 // indirect
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
github.com/MicahParks/jwkset v0.11.3 h1:Phli4RdTDdIdLXZpuO7abkwZyzIk0RDTUPVVBHPRdkQ=
|
||||
github.com/MicahParks/jwkset v0.11.3/go.mod h1:U2oRhRaLgDCLjtpGL2GseNKGmZtLs/3O7p+OZaL5vo0=
|
||||
github.com/MicahParks/keyfunc/v3 v3.8.2 h1:eydEwk/pBAVrDIpmFfB/gkCcrp++xQ7YYXirrI2zlWE=
|
||||
github.com/MicahParks/keyfunc/v3 v3.8.2/go.mod h1:T4snFPe26GwMg45bBAdM5P6qWQyLxZHLwBhxR/9PnCs=
|
||||
github.com/MicahParks/jwkset v0.11.4 h1:sm6pqlEpa1YDu82Hoxho/O+Y291lPcksNImVIsV74gg=
|
||||
github.com/MicahParks/jwkset v0.11.4/go.mod h1:U2oRhRaLgDCLjtpGL2GseNKGmZtLs/3O7p+OZaL5vo0=
|
||||
github.com/MicahParks/keyfunc/v3 v3.9.1 h1:10dlh2kS3kRtfj5x2bokCxLhTpmIpf3WtsvBdJ2TfOo=
|
||||
github.com/MicahParks/keyfunc/v3 v3.9.1/go.mod h1:h8DMhdulFa3yjG1q74FqaehNDekB18HrVCXPl9Raa4I=
|
||||
github.com/VictoriaMetrics/metrics v1.44.1 h1:OMyGrMkhI3hnq+IpJ1gZK9VhXcsOjUX1CoOFv//bJ0I=
|
||||
github.com/VictoriaMetrics/metrics v1.44.1/go.mod h1:xDM82ULLYCYdFRgQ2JBxi8Uf1+8En1So9YUwlGTOqTc=
|
||||
github.com/alicebob/miniredis/v2 v2.40.0 h1:6IgTQbvm2JiMYZ8yur/acNJzVjtSOWCelxXbk/LUBXc=
|
||||
@@ -72,8 +72,6 @@ github.com/klauspost/cpuid/v2 v2.4.0 h1:S6Hrbc7+ywsr0r+RLapfGBHfyefhCTwEh3A0tV91
|
||||
github.com/klauspost/cpuid/v2 v2.4.0/go.mod h1:19jmZ9mjzoF//ddRSUsv0zfBTJWh3QJh9FNxZTMrGxU=
|
||||
github.com/lukaszraczylo/ask v0.0.0-20240916204100-6e9ef53a62d9 h1:pL8B9mjv6RPUfKYYGm/uJ7QL6Ndf+z+OEl0qJE6KmEc=
|
||||
github.com/lukaszraczylo/ask v0.0.0-20240916204100-6e9ef53a62d9/go.mod h1:M+UVdyqZs++xtEPrascaVmZdOMhCnxjZ2SgH+xHpR0c=
|
||||
github.com/lukaszraczylo/go-ratecounter v0.1.13 h1:p1r53oyMb9ZnMbqub6IF4CWhURFiz7k967I0GbUyQMA=
|
||||
github.com/lukaszraczylo/go-ratecounter v0.1.13/go.mod h1:TqXEOCtFJStk1i0tkipprv1kiDHGon1MVUisjSTBSKM=
|
||||
github.com/lukaszraczylo/go-simple-graphql v1.2.89 h1:Xbu1Ny+a0lT2Sr2SaSC8mcHmGQDwGD4TJKk4DDd+PwA=
|
||||
github.com/lukaszraczylo/go-simple-graphql v1.2.89/go.mod h1:PxQYblQDZISmYYj8sNfazAWxAOh1rhAtU208y+uPV8s=
|
||||
github.com/lukaszraczylo/oss-telemetry v0.2.5 h1:W6Qgwp7IY4yBRO361MmnEAmM59motdQpbAFAx8v6fJM=
|
||||
@@ -144,8 +142,8 @@ go.uber.org/automaxprocs v1.6.0 h1:O3y2/QNTOdbF+e/dpXNNW7Rx2hZ4sTIPyybbxyNqTUs=
|
||||
go.uber.org/automaxprocs v1.6.0/go.mod h1:ifeIMSnPZuznNm6jmdzmU3/bfk01Fe2fotchwEFJ8r8=
|
||||
go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
|
||||
go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
|
||||
go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw=
|
||||
go.yaml.in/yaml/v3 v3.0.5/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg=
|
||||
go.yaml.in/yaml/v3 v3.0.6 h1:zcDfvjXEqTQmpjniUBLc91xYJMB2c4QOFxEF/LCEckM=
|
||||
go.yaml.in/yaml/v3 v3.0.6/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg=
|
||||
golang.org/x/crypto v0.58.0 h1:COkYLr4k7nDI4r6QtMML2AfzeIswWHAJthLS6K/M/54=
|
||||
golang.org/x/crypto v0.58.0/go.mod h1:Xh+kl5A+M0gMiWDytH2Hhr5Z7WbNCpv/eB/EFgfSk+w=
|
||||
golang.org/x/net v0.62.0 h1:18ELalfM5+XNRwzM2nYqXsZQ9ByS1uhj2O/U7+5jqRY=
|
||||
|
||||
+42
-113
@@ -8,17 +8,17 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/goccy/go-json"
|
||||
goratecounter "github.com/lukaszraczylo/go-ratecounter"
|
||||
libpack_logger "github.com/lukaszraczylo/graphql-monitoring-proxy/logging"
|
||||
)
|
||||
|
||||
// RateLimitConfig holds the rate limit configuration for a role
|
||||
type RateLimitConfig struct {
|
||||
RateCounterTicker *goratecounter.RateCounter
|
||||
// sustained enforces Req over the full Interval.
|
||||
sustained *attemptWindow
|
||||
// burstCounter enforces Burst over a short trailing window, independent
|
||||
// of RateCounterTicker's sustained-rate window. Nil when Burst <= 0
|
||||
// of the sustained window. Nil when Burst <= 0
|
||||
// (burst disabled). See burstWindow and checkRateLimit for semantics.
|
||||
burstCounter *burstWindow
|
||||
burstCounter *attemptWindow
|
||||
Endpoints []string `json:"endpoints,omitempty"`
|
||||
Interval time.Duration `json:"interval"`
|
||||
Req int `json:"req"`
|
||||
@@ -107,51 +107,6 @@ func burstWindowDuration(interval time.Duration) time.Duration {
|
||||
return w
|
||||
}
|
||||
|
||||
// burstWindow is a minimal, clock-injectable, trailing-window request
|
||||
// counter used to enforce RateLimitConfig.Burst independently of the
|
||||
// sustained-rate RateCounterTicker.
|
||||
//
|
||||
// [email protected] cannot provide this second, shorter window
|
||||
// cleanly: the window length lives on *RateCounter itself, not per named
|
||||
// Counter (helpers.go: GetRate -> c.parent.interval.Seconds()), so two
|
||||
// different window lengths need two separate RateCounter instances. Each
|
||||
// instance drives its own background goroutine off wall-clock time.Now with
|
||||
// no seam to inject a fake clock (helpers.go: start(), cleanUpOldValues(),
|
||||
// addValue() all call time.Now() directly). That would force any test of
|
||||
// window rollover to sleep on real time. burstWindow is implemented here
|
||||
// instead, with a swappable clock, so tests can advance time deterministically.
|
||||
type burstWindow struct {
|
||||
mu sync.Mutex
|
||||
now func() time.Time
|
||||
window time.Duration
|
||||
hits []time.Time
|
||||
}
|
||||
|
||||
// newBurstWindow returns a burstWindow that tracks hits over the trailing
|
||||
// duration window, using the real wall clock.
|
||||
func newBurstWindow(window time.Duration) *burstWindow {
|
||||
return &burstWindow{now: time.Now, window: window}
|
||||
}
|
||||
|
||||
// recordAndCount records a hit at the current time, drops hits older than
|
||||
// the trailing window, and returns the resulting hit count within the
|
||||
// window. Safe for concurrent use.
|
||||
func (b *burstWindow) recordAndCount() int {
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
now := b.now()
|
||||
cutoff := now.Add(-b.window)
|
||||
live := b.hits[:0]
|
||||
for _, t := range b.hits {
|
||||
if t.After(cutoff) {
|
||||
live = append(live, t)
|
||||
}
|
||||
}
|
||||
live = append(live, now)
|
||||
b.hits = live
|
||||
return len(b.hits)
|
||||
}
|
||||
|
||||
// Variable to hold the current load config function - allows for testing
|
||||
var loadConfigFunc = loadConfigFromPath
|
||||
|
||||
@@ -237,14 +192,12 @@ func loadConfigFromPath(path string) error {
|
||||
|
||||
newRateLimits := make(map[string]RateLimitConfig, len(config.RateLimit))
|
||||
for key, value := range config.RateLimit {
|
||||
value.RateCounterTicker = goratecounter.NewRateCounter().WithConfig(goratecounter.RateCounterConfig{
|
||||
Interval: value.Interval,
|
||||
})
|
||||
value.sustained = newAttemptWindow(value.Interval, value.Req)
|
||||
|
||||
// Burst disabled (<=0) leaves burstCounter nil, so checkRateLimit
|
||||
// performs the Req-only check exactly as before burst existed.
|
||||
if value.Burst > 0 {
|
||||
value.burstCounter = newBurstWindow(burstWindowDuration(value.Interval))
|
||||
value.burstCounter = newAttemptWindow(burstWindowDuration(value.Interval), value.Burst)
|
||||
}
|
||||
|
||||
if cfg.LogLevel == "DEBUG" {
|
||||
@@ -279,7 +232,7 @@ func rateLimitedRequest(userID, userRole string) bool {
|
||||
// Try to get config from atomic value first for better performance
|
||||
if configInterface := rateLimitConfigAtomic.Load(); configInterface != nil {
|
||||
if config, ok := configInterface.(map[string]RateLimitConfig); ok {
|
||||
if roleConfig, exists := config[userRole]; exists && roleConfig.RateCounterTicker != nil {
|
||||
if roleConfig, exists := config[userRole]; exists && roleConfig.sustained != nil {
|
||||
return checkRateLimit(userID, userRole, roleConfig, "")
|
||||
}
|
||||
}
|
||||
@@ -290,7 +243,7 @@ func rateLimitedRequest(userID, userRole string) bool {
|
||||
roleConfig, ok := rateLimits[userRole]
|
||||
rateLimitMu.RUnlock()
|
||||
|
||||
if !ok || roleConfig.RateCounterTicker == nil {
|
||||
if !ok || roleConfig.sustained == nil {
|
||||
cfg.Logger.Warning(&libpack_logger.LogMessage{
|
||||
Message: "Rate limit role not found or ticker not initialized - defaulting to deny",
|
||||
Pairs: map[string]any{"user_role": userRole},
|
||||
@@ -304,72 +257,48 @@ func rateLimitedRequest(userID, userRole string) bool {
|
||||
|
||||
// checkRateLimit performs the actual rate limit check.
|
||||
//
|
||||
// Two independent gates apply, both must pass:
|
||||
// Two gates apply, both must pass:
|
||||
//
|
||||
// 1. Sustained: tickerRate (requests measured over the full Interval) must
|
||||
// stay at or below Req. This is the only gate when Burst <= 0, and its
|
||||
// comparison is unchanged from before the burst fix.
|
||||
// 2. Burst (only when roleConfig.burstCounter != nil, i.e. Burst > 0): the
|
||||
// raw request count within a short trailing window (see
|
||||
// burstWindowDuration), much shorter than Interval, must stay at or
|
||||
// below Burst.
|
||||
// 1. Sustained: the request count over the full Interval must stay at or
|
||||
// below Req (Req is per Interval, not per second).
|
||||
// 2. Burst (only when Burst > 0): the request count within a short trailing
|
||||
// window (see burstWindowDuration) must stay at or below Burst.
|
||||
//
|
||||
// Because the burst window is short relative to Interval, a client can
|
||||
// legally land up to Burst requests in a tight cluster - more than a
|
||||
// steady Req-paced stream would deliver in that same short span - as long
|
||||
// as the sustained gate still holds over the full Interval. That is what
|
||||
// makes burst > req a real spike allowance instead of a no-op. Burst is a
|
||||
// short-window ceiling layered ON TOP of the sustained Req rate, not a
|
||||
// replacement for it. burst < req is an unusual config: the burst gate
|
||||
// becomes the tighter, binding constraint whenever requests cluster in
|
||||
// time, but the sustained cap for evenly-spread traffic is still Req
|
||||
// alone. The pre-fix code compared the SAME sustained-window tickerRate
|
||||
// against both Burst and Req, so when burst < req its Burst check tripped
|
||||
// first and Burst was effectively the sustained cap too. This gate's
|
||||
// Req-only sustained check can therefore allow MORE through than the
|
||||
// pre-fix code did for evenly-spread traffic under a burst < req config.
|
||||
// In both cases the burst gate never allows more than Burst requests
|
||||
// within its own window (fail-safe).
|
||||
// The burst window lies inside the Interval, so its count never exceeds the
|
||||
// sustained count. Burst >= Req therefore never rejects anything the sustained
|
||||
// gate admits; Burst only takes effect when it is below Req.
|
||||
func checkRateLimit(userID, userRole string, roleConfig RateLimitConfig, endpoint string) bool {
|
||||
roleConfig.RateCounterTicker.Incr(1)
|
||||
tickerRate := roleConfig.RateCounterTicker.GetRate()
|
||||
debug := cfg.Logger.IsLevelEnabled(libpack_logger.LEVEL_DEBUG)
|
||||
|
||||
logDetails := map[string]any{
|
||||
"user_role": userRole,
|
||||
"user_id": userID,
|
||||
"rate": tickerRate,
|
||||
"config_rate": roleConfig.Req,
|
||||
"interval": roleConfig.Interval,
|
||||
"endpoint": endpoint,
|
||||
}
|
||||
|
||||
cfg.Logger.Debug(&libpack_logger.LogMessage{
|
||||
Message: "Rate limit ticker",
|
||||
Pairs: map[string]any{"log_details": logDetails},
|
||||
})
|
||||
|
||||
if tickerRate > float64(roleConfig.Req) {
|
||||
cfg.Logger.Debug(&libpack_logger.LogMessage{
|
||||
Message: "Rate limit exceeded",
|
||||
Pairs: map[string]any{"log_details": logDetails},
|
||||
})
|
||||
if roleConfig.sustained.recordAndExceeds() {
|
||||
if debug {
|
||||
cfg.Logger.Debug(&libpack_logger.LogMessage{
|
||||
Message: "Rate limit exceeded",
|
||||
Pairs: map[string]any{
|
||||
"user_role": userRole,
|
||||
"user_id": userID,
|
||||
"config_rate": roleConfig.Req,
|
||||
"interval": roleConfig.Interval,
|
||||
"endpoint": endpoint,
|
||||
},
|
||||
})
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
if roleConfig.burstCounter != nil {
|
||||
burstCount := roleConfig.burstCounter.recordAndCount()
|
||||
burstLogDetails := map[string]any{
|
||||
"user_role": userRole,
|
||||
"user_id": userID,
|
||||
"burst_count": burstCount,
|
||||
"config_burst": roleConfig.Burst,
|
||||
"endpoint": endpoint,
|
||||
}
|
||||
if burstCount > roleConfig.Burst {
|
||||
cfg.Logger.Debug(&libpack_logger.LogMessage{
|
||||
Message: "Burst limit exceeded",
|
||||
Pairs: map[string]any{"log_details": burstLogDetails},
|
||||
})
|
||||
if roleConfig.burstCounter.recordAndExceeds() {
|
||||
if debug {
|
||||
cfg.Logger.Debug(&libpack_logger.LogMessage{
|
||||
Message: "Burst limit exceeded",
|
||||
Pairs: map[string]any{
|
||||
"user_role": userRole,
|
||||
"user_id": userID,
|
||||
"config_burst": roleConfig.Burst,
|
||||
"endpoint": endpoint,
|
||||
},
|
||||
})
|
||||
}
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
+21
-63
@@ -10,7 +10,6 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/goccy/go-json"
|
||||
goratecounter "github.com/lukaszraczylo/go-ratecounter"
|
||||
libpack_logger "github.com/lukaszraczylo/graphql-monitoring-proxy/logging"
|
||||
)
|
||||
|
||||
@@ -65,8 +64,8 @@ func (suite *Tests) Test_loadRatelimitConfig() {
|
||||
suite.Contains(rateLimits, "user")
|
||||
suite.Equal(100, rateLimits["admin"].Req)
|
||||
suite.Equal(10, rateLimits["user"].Req)
|
||||
suite.NotNil(rateLimits["admin"].RateCounterTicker)
|
||||
suite.NotNil(rateLimits["user"].RateCounterTicker)
|
||||
suite.NotNil(rateLimits["admin"].sustained)
|
||||
suite.NotNil(rateLimits["user"].sustained)
|
||||
})
|
||||
|
||||
// Test loading config from non-existent path
|
||||
@@ -154,23 +153,17 @@ func (suite *Tests) Test_rateLimitedRequest() {
|
||||
rateLimits = make(map[string]RateLimitConfig)
|
||||
|
||||
// Admin role with high limit
|
||||
adminCounter := goratecounter.NewRateCounter().WithConfig(goratecounter.RateCounterConfig{
|
||||
Interval: 1 * time.Second,
|
||||
})
|
||||
rateLimits["admin"] = RateLimitConfig{
|
||||
RateCounterTicker: adminCounter,
|
||||
Interval: 1 * time.Second,
|
||||
Req: 100,
|
||||
sustained: newAttemptWindow(1*time.Second, 100),
|
||||
Interval: 1 * time.Second,
|
||||
Req: 100,
|
||||
}
|
||||
|
||||
// User role with low limit
|
||||
userCounter := goratecounter.NewRateCounter().WithConfig(goratecounter.RateCounterConfig{
|
||||
Interval: 1 * time.Second,
|
||||
})
|
||||
rateLimits["user"] = RateLimitConfig{
|
||||
RateCounterTicker: userCounter,
|
||||
Interval: 1 * time.Second,
|
||||
Req: 2, // Set very low for testing
|
||||
sustained: newAttemptWindow(1*time.Second, 2),
|
||||
Interval: 1 * time.Second,
|
||||
Req: 2, // Set very low for testing
|
||||
}
|
||||
rateLimitMu.Unlock()
|
||||
|
||||
@@ -327,19 +320,17 @@ func (f *fakeClock) Advance(d time.Duration) {
|
||||
}
|
||||
|
||||
// newBurstTestRoleConfig mirrors the RateLimitConfig construction performed
|
||||
// by loadConfigFromPath: a real RateCounterTicker for the sustained check,
|
||||
// by loadConfigFromPath: a real sustained window for the Req check,
|
||||
// and (only when burst > 0) a burstCounter sized by burstWindowDuration.
|
||||
func newBurstTestRoleConfig(interval time.Duration, req, burst int) RateLimitConfig {
|
||||
rc := RateLimitConfig{
|
||||
RateCounterTicker: goratecounter.NewRateCounter().WithConfig(goratecounter.RateCounterConfig{
|
||||
Interval: interval,
|
||||
}),
|
||||
Interval: interval,
|
||||
Req: req,
|
||||
Burst: burst,
|
||||
sustained: newAttemptWindow(interval, req),
|
||||
Interval: interval,
|
||||
Req: req,
|
||||
Burst: burst,
|
||||
}
|
||||
if burst > 0 {
|
||||
rc.burstCounter = newBurstWindow(burstWindowDuration(interval))
|
||||
rc.burstCounter = newAttemptWindow(burstWindowDuration(interval), burst)
|
||||
}
|
||||
return rc
|
||||
}
|
||||
@@ -372,59 +363,26 @@ func Test_checkRateLimit_BurstSemantics(t *testing.T) {
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("b_burst_above_req_allows_short_spike_sustained_still_capped", func(t *testing.T) {
|
||||
t.Run("b_burst_above_req_is_noop_sustained_caps_at_req", func(t *testing.T) {
|
||||
const (
|
||||
interval = 10 * time.Second
|
||||
req = 3
|
||||
burst = 6 // burst > req
|
||||
burst = 6 // burst >= req: dominated by the sustained gate
|
||||
)
|
||||
rc := newBurstTestRoleConfig(interval, req, burst)
|
||||
fc := newFakeClock(time.Now())
|
||||
rc.burstCounter.now = fc.Now
|
||||
|
||||
// A rapid cluster of `burst` requests must all be allowed, even
|
||||
// though burst (6) > req (3): the sustained gate only bounds the
|
||||
// total over the full 10s Interval (tickerRate <= 3 needs count <=
|
||||
// 30), so a small cluster like this never trips it.
|
||||
for i := 1; i <= burst; i++ {
|
||||
for i := 1; i <= req; i++ {
|
||||
if !checkRateLimit("u", "role-b", rc, "") {
|
||||
t.Fatalf("request %d in first burst: got denied, want allowed (<= Burst=%d)", i, burst)
|
||||
t.Fatalf("request %d: got denied, want allowed (<= Req=%d per Interval)", i, req)
|
||||
}
|
||||
}
|
||||
|
||||
// The (burst+1)th request in the SAME short window must be denied
|
||||
// by the burst gate specifically - not by the sustained gate,
|
||||
// which alone would still allow it (tickerRate = 7/10 = 0.7 <= 3).
|
||||
// This is the proof burst > req is not a no-op: it is the gate
|
||||
// that actually fires here.
|
||||
// Clearing the burst window must not readmit: Req is per Interval.
|
||||
fc.Advance(2 * burstWindowDuration(interval))
|
||||
if checkRateLimit("u", "role-b", rc, "") {
|
||||
t.Fatalf("request %d: got allowed, want denied by burst gate (> Burst=%d)", burst+1, burst)
|
||||
}
|
||||
|
||||
// Roll the burst window forward (fake clock only - the real
|
||||
// sustained ticker has not aged, since no real time passed) and
|
||||
// repeat clusters of `burst` requests. Each cluster alone passes
|
||||
// burst, but the cumulative sustained count keeps climbing until
|
||||
// it exceeds req*interval.Seconds() = 30, at which point the
|
||||
// sustained gate must deny regardless of how the burst window
|
||||
// resets. This proves "sustained still capped at Req".
|
||||
sustainedTripped := false
|
||||
for batch := 0; batch < 10 && !sustainedTripped; batch++ {
|
||||
fc.Advance(2 * burstWindowDuration(interval)) // clears the burst window
|
||||
for i := 1; i <= burst; i++ {
|
||||
allowed := checkRateLimit("u", "role-b", rc, "")
|
||||
rate := rc.RateCounterTicker.GetRate()
|
||||
if rate > float64(req) {
|
||||
if allowed {
|
||||
t.Fatalf("batch %d req %d: allowed with sustained rate %.2f > Req=%d", batch, i, rate, req)
|
||||
}
|
||||
sustainedTripped = true
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
if !sustainedTripped {
|
||||
t.Fatalf("sustained gate never tripped after repeated bursts; expected it to cap total requests at Req over the Interval")
|
||||
t.Fatalf("request %d: got allowed, want denied by the sustained gate (Req=%d per Interval)", req+1, req)
|
||||
}
|
||||
})
|
||||
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
libpack_logger "github.com/lukaszraczylo/graphql-monitoring-proxy/logging"
|
||||
)
|
||||
|
||||
func setupRateLimitCfg() {
|
||||
cfg = &config{}
|
||||
parseConfig()
|
||||
cfg.Logger = libpack_logger.New()
|
||||
}
|
||||
|
||||
// Pins theorems/Theorems/RateLimitUnits.lean: Req counts requests per Interval, not per second.
|
||||
func TestCheckRateLimit_ReqIsPerInterval(t *testing.T) {
|
||||
const req = 3
|
||||
setupRateLimitCfg()
|
||||
rc := newBurstTestRoleConfig(time.Minute, req, 0)
|
||||
for i := 1; i <= req; i++ {
|
||||
if !checkRateLimit("u", "r", rc, "") {
|
||||
t.Fatalf("request %d denied, want allowed (Req=%d per minute)", i, req)
|
||||
}
|
||||
}
|
||||
if checkRateLimit("u", "r", rc, "") {
|
||||
t.Fatalf("request %d allowed, want denied: Req=%d per minute", req+1, req)
|
||||
}
|
||||
}
|
||||
|
||||
// With Burst >= Req the burst gate never rejects what the sustained gate admits.
|
||||
func TestCheckRateLimit_BurstAboveReqIsNoOp(t *testing.T) {
|
||||
const req, burst = 3, 6
|
||||
setupRateLimitCfg()
|
||||
rc := newBurstTestRoleConfig(time.Minute, req, burst)
|
||||
for i := 1; i <= req; i++ {
|
||||
if !checkRateLimit("u", "r", rc, "") {
|
||||
t.Fatalf("request %d denied", i)
|
||||
}
|
||||
}
|
||||
if checkRateLimit("u", "r", rc, "") {
|
||||
t.Fatalf("request %d allowed, want denied by the sustained gate", req+1)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// attemptWindow decides "more than limit attempts in the trailing window" while
|
||||
// storing only the newest limit+1 attempt times. The attempt number limit+1
|
||||
// counted back from now lies inside the window exactly when the full count
|
||||
// exceeds limit (see theorems/Theorems/SlidingWindow.lean), so memory and time
|
||||
// per attempt stay constant however hard a client floods.
|
||||
type attemptWindow struct {
|
||||
now func() time.Time
|
||||
base time.Time // monotonic origin for stored offsets
|
||||
ring []int64 // offsets from base, grows to capacity then wraps
|
||||
mu sync.Mutex
|
||||
window time.Duration
|
||||
head int // index of the oldest stored attempt once the ring is full
|
||||
size int // limit+1: attempts kept; the ring grows lazily up to this
|
||||
}
|
||||
|
||||
func newAttemptWindow(window time.Duration, limit int) *attemptWindow {
|
||||
return &attemptWindow{
|
||||
now: time.Now,
|
||||
base: time.Now(),
|
||||
window: window,
|
||||
size: max(limit, 0) + 1,
|
||||
}
|
||||
}
|
||||
|
||||
// recordAndExceeds records an attempt now and reports whether the attempts in
|
||||
// the trailing window, this one included, outnumber the limit.
|
||||
func (w *attemptWindow) recordAndExceeds() bool {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
|
||||
t := int64(w.now().Sub(w.base))
|
||||
if len(w.ring) < w.size {
|
||||
w.ring = append(w.ring, t)
|
||||
} else {
|
||||
w.ring[w.head] = t
|
||||
w.head = (w.head + 1) % len(w.ring)
|
||||
}
|
||||
if len(w.ring) < w.size {
|
||||
return false
|
||||
}
|
||||
return w.ring[w.head] > t-int64(w.window)
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"math/rand"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// naiveExceeds is the exact reference: count every attempt in the trailing window.
|
||||
func naiveExceeds(attempts []time.Time, now time.Time, window time.Duration, limit int) bool {
|
||||
cutoff := now.Add(-window)
|
||||
n := 0
|
||||
for _, t := range attempts {
|
||||
if t.After(cutoff) {
|
||||
n++
|
||||
}
|
||||
}
|
||||
return n > limit
|
||||
}
|
||||
|
||||
// Pins theorems/Theorems/SlidingWindow.lean: storing the newest limit+1 attempts decides like the full count.
|
||||
func TestAttemptWindow_MatchesExactCount(t *testing.T) {
|
||||
for _, limit := range []int{0, 1, 3, 10} {
|
||||
rng := rand.New(rand.NewSource(int64(limit) + 1))
|
||||
window := 10 * time.Second
|
||||
fc := newFakeClock(time.Unix(1_000_000, 0))
|
||||
w := newAttemptWindow(window, limit)
|
||||
w.now = fc.Now
|
||||
w.base = fc.Now()
|
||||
|
||||
var attempts []time.Time
|
||||
for i := 0; i < 5000; i++ {
|
||||
fc.Advance(time.Duration(rng.Intn(3000)) * time.Millisecond)
|
||||
attempts = append(attempts, fc.Now())
|
||||
want := naiveExceeds(attempts, fc.Now(), window, limit)
|
||||
if got := w.recordAndExceeds(); got != want {
|
||||
t.Fatalf("limit=%d step=%d: got %v, want %v", limit, i, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAttemptWindow_StorageIsBounded(t *testing.T) {
|
||||
const limit = 10
|
||||
fc := newFakeClock(time.Unix(1_000_000, 0))
|
||||
w := newAttemptWindow(time.Hour, limit)
|
||||
w.now = fc.Now
|
||||
w.base = fc.Now()
|
||||
for i := 0; i < 100_000; i++ {
|
||||
w.recordAndExceeds()
|
||||
fc.Advance(time.Millisecond)
|
||||
}
|
||||
if got := len(w.ring); got > limit+1 {
|
||||
t.Fatalf("stored %d timestamps, want <= %d", got, limit+1)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkAttemptWindow(b *testing.B) {
|
||||
w := newAttemptWindow(time.Hour, 1000)
|
||||
b.ReportAllocs()
|
||||
for i := 0; i < b.N; i++ {
|
||||
w.recordAndExceeds()
|
||||
}
|
||||
}
|
||||
@@ -442,6 +442,11 @@ func processGraphQLRequest(c fiber.Ctx) error {
|
||||
// never an error; the request stays anonymous, exactly like before
|
||||
// verification existed.
|
||||
func extractUserInfo(c fiber.Ctx) (string, string, error) {
|
||||
return extractUserInfoFromHeaders(func(key string) string { return c.Get(key) })
|
||||
}
|
||||
|
||||
// extractUserInfoFromHeaders resolves user id and role from request headers via get.
|
||||
func extractUserInfoFromHeaders(get func(string) string) (string, string, error) {
|
||||
extractedUserID := defaultValue
|
||||
extractedRoleName := defaultValue
|
||||
|
||||
@@ -450,7 +455,7 @@ func extractUserInfo(c fiber.Ctx) (string, string, error) {
|
||||
// Authorization header must still be checked even with no claim path
|
||||
// set, otherwise JWT_VERIFY_SIGNATURE=true alone would never call
|
||||
// Verify and the forged-token rejection below would never trigger.
|
||||
if authorization := c.Get("Authorization"); authorization != "" &&
|
||||
if authorization := get("Authorization"); authorization != "" &&
|
||||
(len(cfg.Client.JWTUserClaimPath) > 0 || len(cfg.Client.JWTRoleClaimPath) > 0 || cfg.Client.JWTVerifier != nil) {
|
||||
usr, role, err := extractClaimsFromJWTHeader(authorization)
|
||||
if err != nil {
|
||||
@@ -461,7 +466,7 @@ func extractUserInfo(c fiber.Ctx) (string, string, error) {
|
||||
|
||||
// Override role from header if configured
|
||||
if cfg.Client.RoleFromHeader != "" {
|
||||
if role := c.Get(cfg.Client.RoleFromHeader); role != "" {
|
||||
if role := get(cfg.Client.RoleFromHeader); role != "" {
|
||||
extractedRoleName = role
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,6 +2,7 @@ package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"sync"
|
||||
@@ -231,6 +232,22 @@ func (wsp *WebSocketProxy) handleConnection(ctx context.Context, clientConn *web
|
||||
// Try to extract headers from connection_init payload (for GraphQL WebSocket protocols)
|
||||
enrichedHeaders := wsp.extractAuthFromPayload(message, headers)
|
||||
|
||||
// The upgrade request was gated on its own headers; the backend acts as the
|
||||
// identity in enrichedHeaders, so gate that one too.
|
||||
if authErr := authorizePayloadIdentity(headers, enrichedHeaders); authErr != nil {
|
||||
wsp.errors.Add(1)
|
||||
if wsp.logger != nil {
|
||||
wsp.logger.Warning(&libpack_logger.LogMessage{
|
||||
Message: "WebSocket connection_init identity rejected",
|
||||
Pairs: map[string]any{"connection_id": connectionID, "reason": authErr.Error()},
|
||||
})
|
||||
}
|
||||
_ = clientConn.WriteMessage(websocket.CloseMessage,
|
||||
websocket.FormatCloseMessage(websocket.ClosePolicyViolation, authErr.Error()))
|
||||
_ = clientConn.Close() // Best-effort cleanup
|
||||
return
|
||||
}
|
||||
|
||||
// Connect to backend WebSocket with enriched headers
|
||||
backendConn, err := wsp.dialBackend(ctx, enrichedHeaders)
|
||||
if err != nil {
|
||||
@@ -586,6 +603,36 @@ func (wsp *WebSocketProxy) extractAuthFromPayload(message []byte, originalHeader
|
||||
return enrichedHeaders
|
||||
}
|
||||
|
||||
var (
|
||||
errWSUnauthorized = errors.New("invalid or expired authorization token")
|
||||
errWSBanned = errors.New("user is banned")
|
||||
errWSRateLimited = errors.New("rate limit exceeded")
|
||||
)
|
||||
|
||||
// authorizePayloadIdentity applies the ban and rate-limit checks of the upgrade
|
||||
// path to the identity that connection_init supplied. It does nothing when the
|
||||
// payload left Authorization and the role header unchanged, because the upgrade
|
||||
// request already passed those checks (and counted against the rate limit).
|
||||
func authorizePayloadIdentity(upgrade, enriched http.Header) error {
|
||||
roleHeader := cfg.Client.RoleFromHeader
|
||||
if upgrade.Get("Authorization") == enriched.Get("Authorization") &&
|
||||
(roleHeader == "" || upgrade.Get(roleHeader) == enriched.Get(roleHeader)) {
|
||||
return nil
|
||||
}
|
||||
|
||||
userID, role, err := extractUserInfoFromHeaders(enriched.Get)
|
||||
if err != nil {
|
||||
return errWSUnauthorized
|
||||
}
|
||||
if isUserBanned(userID) {
|
||||
return errWSBanned
|
||||
}
|
||||
if cfg.Client.RoleRateLimit && !rateLimitedRequest(userID, role) {
|
||||
return errWSRateLimited
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// dialBackend establishes a WebSocket connection to the backend
|
||||
func (wsp *WebSocketProxy) dialBackend(ctx context.Context, headers http.Header) (*gorillaws.Conn, error) {
|
||||
// Convert http:// to ws:// or https:// to wss://
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"net/http"
|
||||
"testing"
|
||||
|
||||
libpack_logger "github.com/lukaszraczylo/graphql-monitoring-proxy/logging"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func unsignedJWT(sub, role string) string {
|
||||
enc := base64.RawURLEncoding.EncodeToString
|
||||
return "Bearer " + enc([]byte(`{"alg":"none"}`)) + "." +
|
||||
enc([]byte(`{"sub":"`+sub+`","role":"`+role+`"}`)) + "."
|
||||
}
|
||||
|
||||
// Pins theorems/Theorems/WsIdentityGate.lean: the identity the backend acts as is the one that is gated.
|
||||
func TestAuthorizePayloadIdentity(t *testing.T) {
|
||||
cfg = &config{Logger: libpack_logger.New()}
|
||||
cfg.Client.JWTUserClaimPath = "sub"
|
||||
cfg.Client.JWTRoleClaimPath = "role"
|
||||
bannedUsersIDs.Store("banned-user", "test")
|
||||
defer bannedUsersIDs.Delete("banned-user")
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
upgrade http.Header
|
||||
enriched http.Header
|
||||
wantErr error
|
||||
}{
|
||||
{"payload adds banned identity", http.Header{}, http.Header{"Authorization": {unsignedJWT("banned-user", "user")}}, errWSBanned},
|
||||
{"payload replaces clean identity with banned", http.Header{"Authorization": {unsignedJWT("ok-user", "user")}}, http.Header{"Authorization": {unsignedJWT("banned-user", "user")}}, errWSBanned},
|
||||
{"payload identity clean", http.Header{}, http.Header{"Authorization": {unsignedJWT("ok-user", "user")}}, nil},
|
||||
{"payload does not change identity", http.Header{"Authorization": {unsignedJWT("ok-user", "user")}}, http.Header{"Authorization": {unsignedJWT("ok-user", "user")}}, nil},
|
||||
{"no authorization anywhere", http.Header{}, http.Header{}, nil},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
assert.ErrorIs(t, authorizePayloadIdentity(tt.upgrade, tt.enriched), tt.wantErr)
|
||||
})
|
||||
}
|
||||
}
|
||||
+3
-6
@@ -11,7 +11,6 @@ import (
|
||||
|
||||
fiber "github.com/gofiber/fiber/v3"
|
||||
gorillaws "github.com/gorilla/websocket"
|
||||
goratecounter "github.com/lukaszraczylo/go-ratecounter"
|
||||
libpack_logger "github.com/lukaszraczylo/graphql-monitoring-proxy/logging"
|
||||
libpack_monitoring "github.com/lukaszraczylo/graphql-monitoring-proxy/monitoring"
|
||||
"github.com/stretchr/testify/assert"
|
||||
@@ -652,11 +651,9 @@ func TestHandleWebSocketOrDefault_RateLimitedUserRejected(t *testing.T) {
|
||||
rateLimitMu.Lock()
|
||||
rateLimits = map[string]RateLimitConfig{
|
||||
defaultValue: {
|
||||
RateCounterTicker: goratecounter.NewRateCounter().WithConfig(goratecounter.RateCounterConfig{
|
||||
Interval: 1 * time.Second,
|
||||
}),
|
||||
Interval: 1 * time.Second,
|
||||
Req: 0,
|
||||
sustained: newAttemptWindow(1*time.Second, 0),
|
||||
Interval: 1 * time.Second,
|
||||
Req: 0,
|
||||
},
|
||||
}
|
||||
rateLimitConfigAtomic.Store(rateLimits)
|
||||
|
||||
Reference in New Issue
Block a user