Compare commits

...
14 Commits
Author SHA1 Message Date
github-actions[bot]andGitHub e3066ad7f0 chore(deps): update dependencies
Automated dependency update. All tests passed.
2026-11-04 03:05:05 +00:00
github-actions[bot]andGitHub e138b44d00 chore(deps): update module go.yaml.in/yaml/v3 to v3.0.6
Automated dependency update. All tests passed.
2026-11-03 03:05:12 +00:00
lukaszraczylo 3a6403c30b fix(ratelimit): bound memory and CPU per request
The sustained gate stored one entry for every attempt, denied ones too, and summed all entries on each request. A flood grew memory with rate times interval and made each request slower. Both gates now keep only the newest limit plus one attempt times. A Lean proof shows this decides the same as the full count. One check dropped from 29,502 ns and 11 allocations to 43 ns and none. The go-ratecounter dependency is gone.
2026-11-03 00:55:35 +00:00
lukaszraczylo cff11f2b25 fix(redis): namespace cache keys when no prefix is set
With an empty prefix, Clear deleted every key in the Redis database, including keys of other applications. Cache keys now start with gmp: by default. Existing cache entries become unreachable after the upgrade and expire by TTL.
2026-11-03 00:49:02 +00:00
lukaszraczylo 4c5d01aef3 fix(ws): gate the identity sent in connection_init
The upgrade request passed the ban and rate-limit checks on its own headers. The backend then acted as the identity in the connection_init payload. A banned user could connect without credentials and send their token in the payload. The proxy now runs the same checks on the payload identity.
2026-11-03 00:34:18 +00:00
lukaszraczylo 6e7e2ecfae fix(events): refuse a retention below one day
A retention of zero or less moved the delete cutoff to now or later. The cleaner then deleted all event history. The cleaner now refuses to start and returns an error.
2026-11-03 00:31:53 +00:00
lukaszraczylo 3e187790f2 fix(health): stop startup wait at its timeout
The startup probe slept a full backoff step even when the deadline was closer. A 31 second timeout ended after about 41 seconds. The probe now sleeps no longer than the time left.
2026-11-03 00:23:56 +00:00
lukaszraczylo ea30694dc9 fix(cache): enforce entry cap below ten in memory cache
The count-based eviction removed a tenth of the cap. For a cap under ten, that was zero entries, so the cache grew without limit. The cache now evicts at least one entry when it is full.
2026-11-03 00:19:35 +00:00
lukaszraczylo c128a0f4d9 fix(ratelimit): count Req per Interval, not per second
The sustained gate compared requests per second with Req, which is a count per Interval. An interval of one minute with req 10 allowed 600 requests. The gate now compares the count in the interval with Req. One-second intervals behave as before. A burst at or above req has no effect, because the sustained gate always binds first.
2026-11-03 00:07:57 +00:00
lukaszraczylo e2af5b079f fix(coalesce): copy the primary response body for waiters
Waiters read the primary response buffer while fasthttp could reuse it for another request. A waiter could then send the wrong body. The primary now gives waiters their own copy.
2026-11-03 00:02:15 +00:00
lukaszraczylo 10d3946f1e fix(rps): stop dropping requests recorded during a sample
The sampler read the counter and then reset it in two steps. Requests recorded between the steps were lost. The sampler now swaps the counter to zero in one atomic step.
2026-11-03 00:00:24 +00:00
lukaszraczylo ab63f924a9 fix(retry): carry fractional tokens across refill ticks
The refill loop dropped the fractional part of each tick. At 15 tokens per second and 100 ms ticks, it added 10 tokens per second instead of 15. The clock now advances only by the time the added tokens pay for.
2026-11-02 23:57:46 +00:00
github-actions[bot]andGitHub a0a8283b52 chore(deps): update dependencies
Automated dependency update. All tests passed.
2026-11-02 03:04:09 +00:00
github-actions[bot]andGitHub cf4f2785e2 chore(deps): update dependencies
Automated dependency update. All tests passed.
2026-11-01 03:04:09 +00:00
29 changed files with 652 additions and 247 deletions
+3
View File
@@ -5,3 +5,6 @@ dist/
coverage.out
CLAUDE.md
graphql-monitoring-proxy
# Lean proofs: local only, never committed
theorems/
+1 -1
View File
@@ -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.
+6 -1
View File
@@ -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
View File
@@ -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
+32
View File
@@ -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)
}
}
+2 -1
View File
@@ -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
+23
View File
@@ -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)
}
})
}
}
+9 -1
View File
@@ -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
View File
@@ -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")
}
+4 -4
View File
@@ -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")
}
+27
View File
@@ -0,0 +1,27 @@
package main
import (
"testing"
"github.com/gofiber/fiber/v3"
"github.com/stretchr/testify/assert"
"github.com/valyala/fasthttp"
)
// Pins theorems/Theorems/CoalesceAlias.lean: waiters never alias the primary's response buffer.
func TestSnapshotCoalescedResponse_OwnsBody(t *testing.T) {
app := fiber.New()
fctx := &fasthttp.RequestCtx{}
c := app.AcquireCtx(fctx)
defer app.ReleaseCtx(c)
c.Response().SetStatusCode(201)
c.Response().SetBody([]byte(`{"user":"alice"}`))
snap := snapshotCoalescedResponse(c)
c.Response().SetBody([]byte(`{"user":"bobby"}`)) // buffer reused by a later request
assert.Equal(t, 201, snap.StatusCode)
assert.Equal(t, `{"user":"alice"}`, string(snap.Body))
}
+7
View File
@@ -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},
+26
View File
@@ -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)
}
}
+7 -8
View File
@@ -5,12 +5,12 @@ 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
github.com/goccy/go-json v0.11.2
github.com/gofiber/contrib/v3/websocket v1.2.9
github.com/goccy/go-json v0.12.0
github.com/gofiber/contrib/v3/websocket v1.2.10
github.com/gofiber/fiber/v3 v3.6.1
github.com/gofrs/flock v0.13.1
github.com/golang-jwt/jwt/v5 v5.3.1
@@ -18,9 +18,8 @@ require (
github.com/gookit/goutil v0.8.6
github.com/gorilla/websocket v1.5.3
github.com/graphql-go/graphql v0.8.1
github.com/jackc/pgx/v5 v5.11.0
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,14 +35,14 @@ 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
github.com/fasthttp/websocket v1.5.12 // indirect
github.com/go-logr/logr v1.4.4 // indirect
github.com/go-logr/stdr v1.2.2 // indirect
github.com/gofiber/schema v1.9.1 // indirect
github.com/gofiber/schema v1.9.2 // indirect
github.com/gofiber/utils/v2 v2.6.3 // indirect
github.com/grpc-ecosystem/grpc-gateway/v2 v2.33.0 // indirect
github.com/jackc/pgpassfile v1.0.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
+14 -16
View File
@@ -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=
@@ -28,16 +28,16 @@ github.com/go-logr/logr v1.4.4 h1:tG4xh9yMsRCAiodLVTxyrkzSZ9+o0L1Kg/+cPVcbP/8=
github.com/go-logr/logr v1.4.4/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
github.com/goccy/go-json v0.11.2 h1:jdZv93Tt4ioR8yW1CoNsvSxrcZlCXAUU1aZXN7gpXUA=
github.com/goccy/go-json v0.11.2/go.mod h1:3NdmfEkZlB7YI5UFw/qdFKq8XN1aiWR0YyRPWZNQltY=
github.com/goccy/go-json v0.12.0 h1:pD+oWouAiz4yruwXIt/iT+lOMeyhr1aTtxzaMP/RgVI=
github.com/goccy/go-json v0.12.0/go.mod h1:3NdmfEkZlB7YI5UFw/qdFKq8XN1aiWR0YyRPWZNQltY=
github.com/goccy/go-reflect v1.2.0 h1:O0T8rZCuNmGXewnATuKYnkL0xm6o8UNOJZd/gOkb9ms=
github.com/goccy/go-reflect v1.2.0/go.mod h1:n0oYZn8VcV2CkWTxi8B9QjkCoq6GTtCEdfmR66YhFtE=
github.com/gofiber/contrib/v3/websocket v1.2.9 h1:DgZgxWK5Yua3g5+I4TmoA0RRx6MILIaqooSGKMgjWZE=
github.com/gofiber/contrib/v3/websocket v1.2.9/go.mod h1:+unMmdAcLfIywqYDaH8eTynn503n7Lf4Cr4kXM5ghO0=
github.com/gofiber/contrib/v3/websocket v1.2.10 h1:tUyRcHSH9GUsh6WO03Xn+Kjvt3850W6raF6WPngdJLE=
github.com/gofiber/contrib/v3/websocket v1.2.10/go.mod h1:O6MOrNHCVdCrrstdLrYmRIL0eJy2N87Wa/y2dszuxvQ=
github.com/gofiber/fiber/v3 v3.6.1 h1:wYjbgqpJm6VErIVLSGzGpX7lLZyTNgElsj8nQ1i+iFo=
github.com/gofiber/fiber/v3 v3.6.1/go.mod h1:+nxjBal3gw0EPf3gqK5EmtxaxTL01+zbSTagFletlDI=
github.com/gofiber/schema v1.9.1 h1:QjicFdTDMlkzCOoZAlNWqdkR2YHhgH8/Q/MpcbHurjs=
github.com/gofiber/schema v1.9.1/go.mod h1:219nYAgzIEBWl/aHCyhQtFoXBvfLjsC6+FI3DUN1+Ck=
github.com/gofiber/schema v1.9.2 h1:Gzw2PF1K5DpWgG2L51/9t9dkmxw4kqK6auKMvgD9mIU=
github.com/gofiber/schema v1.9.2/go.mod h1:219nYAgzIEBWl/aHCyhQtFoXBvfLjsC6+FI3DUN1+Ck=
github.com/gofiber/utils/v2 v2.6.3 h1:VlnOB89rFnonlXoYOkESOkXtIkAw4CBBwpBHu+/a/YQ=
github.com/gofiber/utils/v2 v2.6.3/go.mod h1:Mq3OXwyBPxORb3jPafy6Swki4/TJJW9TDPHeerXkXpg=
github.com/gofrs/flock v0.13.1 h1:jjREztyBeSKBZYAC+mgc1laB+xsgy4kYMf3FbKF2UBo=
@@ -62,8 +62,8 @@ github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsI
github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
github.com/jackc/pgx/v5 v5.11.0 h1:IzBBtyK9AHqf98cctWFifYSci2hgQR/cd56wB4p+ogg=
github.com/jackc/pgx/v5 v5.11.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4=
github.com/jackc/pgx/v5 v5.12.0 h1:hWwax8YiGT0YiN50U1oIDn5kgR7KC2tG3rmavwpPkxY=
github.com/jackc/pgx/v5 v5.12.0/go.mod h1:u2Fi6no03HCSBYE+ZVdLxbwdu3xyi2vU+K/YNJioSds=
github.com/jackc/puddle/v2 v2.2.3 h1:f8HoREISb3l5eO9coPgK7U3hKLhWGokOFVcThl5jrao=
github.com/jackc/puddle/v2 v2.2.3/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
github.com/klauspost/compress v1.20.3 h1:zWNoFqq/lSuYQRim3dUXlXeA0IqwuZGicdG5ZTfY5j4=
@@ -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=
+12 -5
View File
@@ -740,11 +740,7 @@ func performProxyRequest(c fiber.Ctx, proxyURL string) error {
}, proxyErr
}
return &CoalescedResponse{
Body: c.Response().Body(),
StatusCode: c.Response().StatusCode(),
// Headers intentionally left nil; not populated or read anywhere.
}, nil
return snapshotCoalescedResponse(c), nil
})
// Check for error from rc.Do (though it typically returns nil)
@@ -773,6 +769,17 @@ func performProxyRequest(c fiber.Ctx, proxyURL string) error {
return performProxyRequestCore(c, proxyURL, cacheKey)
}
// snapshotCoalescedResponse captures the primary's response for waiters. It
// copies the body because fasthttp reuses the buffer once the primary handler
// returns, while waiters can still be reading.
func snapshotCoalescedResponse(c fiber.Ctx) *CoalescedResponse {
return &CoalescedResponse{
Body: bytes.Clone(c.Response().Body()),
StatusCode: c.Response().StatusCode(),
// Headers intentionally left nil; not populated or read anywhere.
}
}
// performProxyRequestCore executes the proxy request with retries and circuit breaker
// This is the core implementation used by both direct calls and coalesced requests
func performProxyRequestCore(c fiber.Ctx, proxyURL string, cacheKey string) error {
+42 -113
View File
@@ -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
View File
@@ -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)
}
})
+44
View File
@@ -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)
}
}
+49
View File
@@ -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)
}
+64
View File
@@ -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()
}
}
+27 -19
View File
@@ -148,31 +148,39 @@ func (rb *RetryBudget) refill() {
now := time.Now().UnixNano()
last := rb.lastRefill.Load()
// Calculate elapsed time in seconds
elapsed := float64(now-last) / float64(time.Second)
tokensToAdd, newLast := refillStep(tokensPerSecond, last, now)
if tokensToAdd <= 0 {
return
}
if rb.lastRefill.CompareAndSwap(last, newLast) {
for {
current := rb.currentTokens.Load()
newValue := current + tokensToAdd
if newValue > maxTokens {
newValue = maxTokens
}
// Calculate tokens to add
tokensToAdd := int64(elapsed * tokensPerSecond)
if tokensToAdd > 0 {
// Update last refill time
if rb.lastRefill.CompareAndSwap(last, now) {
// Add tokens, capped at maxTokens
for {
current := rb.currentTokens.Load()
newValue := current + tokensToAdd
if newValue > maxTokens {
newValue = maxTokens
}
if rb.currentTokens.CompareAndSwap(current, newValue) {
break
}
if rb.currentTokens.CompareAndSwap(current, newValue) {
break
}
}
}
}
// refillStep returns the whole tokens earned between last and now (ns) and the
// clock to store. The clock advances only by the time those tokens pay for, so
// the fractional remainder carries into the next tick.
func refillStep(tokensPerSecond float64, last, now int64) (int64, int64) {
if tokensPerSecond <= 0 || now <= last {
return 0, last
}
tokens := int64(float64(now-last) / float64(time.Second) * tokensPerSecond)
if tokens <= 0 {
return 0, last
}
return tokens, last + int64(float64(tokens)/tokensPerSecond*float64(time.Second))
}
// GetStats returns current statistics
func (rb *RetryBudget) GetStats() map[string]any {
totalAttempts := rb.totalAttempts.Load()
+47
View File
@@ -0,0 +1,47 @@
package main
import (
"testing"
"time"
"github.com/stretchr/testify/assert"
)
// Pins theorems/Theorems/RetryRefill.lean: total_added(15 tps, 100ms ticks, 10s) = 150.
func TestRefillStep_KeepsFractionalTokens(t *testing.T) {
tests := []struct {
name string
tps float64
want int64
}{
{"15 tps", 15, 150},
{"25 tps", 25, 250},
{"5 tps (sub-token per tick)", 5, 50},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
tick := 100 * time.Millisecond
var last, total int64
now := int64(tick)
for i := 0; i < 100; i++ {
added, newLast := refillStep(tt.tps, last, now)
total += added
last = newLast
now += int64(tick)
}
assert.Equal(t, tt.want, total)
})
}
}
func TestRefillStep_NoElapsedNoChange(t *testing.T) {
added, last := refillStep(10, 1000, 1000)
assert.Zero(t, added)
assert.Equal(t, int64(1000), last)
}
func TestRefillStep_ZeroRateNeverAdds(t *testing.T) {
added, last := refillStep(0, 0, int64(time.Hour))
assert.Zero(t, added)
assert.Zero(t, last)
}
+11 -6
View File
@@ -57,17 +57,23 @@ func (r *RPSTracker) Shutdown() {
// sample calculates RPS since last sample
func (r *RPSTracker) sample() {
now := time.Now()
nowNano := now.UnixNano()
r.sampleCount()
}
// sampleCount performs one sample and returns the requests it consumed.
func (r *RPSTracker) sampleCount() int64 {
nowNano := time.Now().UnixNano()
currentCount := r.lastCount.Load()
lastSampleNano := r.lastSampleTime.Load()
if lastSampleNano == 0 {
r.lastSampleTime.Store(nowNano)
return
return 0
}
// Swap, not Load+Store(0): requests recorded between the two would be lost.
currentCount := r.lastCount.Swap(0)
elapsed := float64(nowNano-lastSampleNano) / float64(time.Second)
if elapsed > 0 {
rps := float64(currentCount) / elapsed
@@ -75,9 +81,8 @@ func (r *RPSTracker) sample() {
r.currentRPS.Store(uint64(rps * 100))
}
// Reset for next sample
r.lastCount.Store(0)
r.lastSampleTime.Store(nowNano)
return currentCount
}
// GetCurrentRPS returns the current requests per second
+43
View File
@@ -0,0 +1,43 @@
package main
import (
"context"
"sync"
"testing"
"github.com/stretchr/testify/assert"
)
// Pins theorems/Theorems/RpsSample.lean: no recorded request is dropped by sample().
func TestRPSTracker_SampleKeepsConcurrentRequests(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
tr := NewRPSTracker(ctx)
tr.Shutdown() // stop the background sampler; sample() is driven by the test
const writers, perWriter = 8, 20000
var wg sync.WaitGroup
for i := 0; i < writers; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for j := 0; j < perWriter; j++ {
tr.RecordRequest()
}
}()
}
var sampled int64
done := make(chan struct{})
go func() { wg.Wait(); close(done) }()
loop:
for {
select {
case <-done:
break loop
default:
sampled += tr.sampleCount()
}
}
sampled += tr.sampleCount()
assert.Equal(t, int64(writers*perWriter), sampled)
}
+7 -2
View File
@@ -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
}
}
+47
View File
@@ -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://
+43
View File
@@ -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
View File
@@ -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)