mirror of
https://github.com/lukaszraczylo/graphql-monitoring-proxy.git
synced 2026-11-04 08:04:10 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
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)
|
||||
}
|
||||
}
|
||||
+10
-26
@@ -304,35 +304,19 @@ 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()
|
||||
tickerRate := roleConfig.RateCounterTicker.Get()
|
||||
|
||||
logDetails := map[string]any{
|
||||
"user_role": userRole,
|
||||
@@ -348,7 +332,7 @@ func checkRateLimit(userID, userRole string, roleConfig RateLimitConfig, endpoin
|
||||
Pairs: map[string]any{"log_details": logDetails},
|
||||
})
|
||||
|
||||
if tickerRate > float64(roleConfig.Req) {
|
||||
if tickerRate > int64(roleConfig.Req) {
|
||||
cfg.Logger.Debug(&libpack_logger.LogMessage{
|
||||
Message: "Rate limit exceeded",
|
||||
Pairs: map[string]any{"log_details": logDetails},
|
||||
|
||||
+7
-40
@@ -372,59 +372,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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user