fix(server): gate WebSocket on ban and rate limit, map errors to status

The WebSocket route skipped the ban and rate-limit checks that the HTTP
path enforced, and returned errors set no custom status, so circuit-open
became HTTP 500. The route now runs the same ban and rate-limit checks
before the upgrade, and a custom error handler maps errors to the
correct status, including 503 for circuit-open. The API and admin
listeners honour the BIND_ADDRESS option.

admin_dashboard_test.go

@@ -161,7 +161,6 @@ func TestAdminDashboard_GetCircuitBreakerStatus(t *testing.T) {
 	cfg = &config{
 		Logger: logger,
 		CircuitBreaker: struct {
-			EndpointConfigs       map[string]*EndpointCBConfig
 			ExcludedStatusCodes   []int
 			MaxFailures           int
 			FailureRatio          float64

api.go

@@ -105,7 +105,7 @@ func enableApi(ctx context.Context) error {
 	// Start server in a goroutine and handle shutdown
 	errCh := make(chan error, 1)
 	go func() {
-		if err := apiserver.Listen(fmt.Sprintf(":%d", cfg.Server.ApiPort), fiber.ListenConfig{DisableStartupMessage: true}); err != nil {
+		if err := apiserver.Listen(fmt.Sprintf("%s:%d", cfg.Server.BindAddress, cfg.Server.ApiPort), fiber.ListenConfig{DisableStartupMessage: true}); err != nil {
 			errCh <- err
 		}
 	}()

server.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"strconv"
 	"sync"
@@ -104,6 +105,7 @@ func StartHTTPProxy() error {
 		WriteTimeout: time.Duration(cfg.Client.ClientTimeout) * time.Second,
 		JSONEncoder:  json.Marshal,
 		JSONDecoder:  json.Unmarshal,
+		ErrorHandler: proxyErrorHandler,
 	}
 
 	server := fiber.New(serverConfig)
@@ -145,15 +147,7 @@ func StartHTTPProxy() error {
 
 	// WebSocket support - must be registered before catch-all routes
 	if cfg.WebSocket.Enable {
-		server.Get("/v1/graphql", func(c fiber.Ctx) error {
-			if IsWebSocketRequest(c) {
-				wsp := GetWebSocketProxy()
-				if wsp != nil {
-					return wsp.HandleWebSocket(c)
-				}
-			}
-			return proxyTheRequestToDefault(c)
-		})
+		server.Get("/v1/graphql", handleWebSocketOrDefault)
 	}
 
 	server.Post("/*", processGraphQLRequest)
@@ -164,7 +158,7 @@ func StartHTTPProxy() error {
 		Pairs:   map[string]any{"port": cfg.Server.PortGraphQL},
 	})
 
-	if err := server.Listen(fmt.Sprintf(":%d", cfg.Server.PortGraphQL), fiber.ListenConfig{DisableStartupMessage: true}); err != nil {
+	if err := server.Listen(fmt.Sprintf("%s:%d", cfg.Server.BindAddress, cfg.Server.PortGraphQL), fiber.ListenConfig{DisableStartupMessage: true}); err != nil {
 		return fmt.Errorf("failed to start HTTP proxy server on port %d: %w",
 			cfg.Server.PortGraphQL, err)
 	}
@@ -172,11 +166,63 @@ func StartHTTPProxy() error {
 	return nil
 }
 
+// proxyErrorHandler maps a handler-returned error to the HTTP status code the
+// client should see and writes a plain-text body, replacing fiber's
+// DefaultErrorHandler (which has no notion of the proxy's error types).
+//
+// A *fiber.Error (returned by fiber itself, e.g. route/body-parsing errors,
+// or explicitly by a handler via fiber.NewError) keeps its own Code, exactly
+// like DefaultErrorHandler does. Any other error - notably a circuit-open
+// sentinel or a *ProxyError - is mapped via StatusCodeForError, so
+// circuit-open surfaces as 503 instead of the generic 500 fiber would
+// otherwise return.
+func proxyErrorHandler(c fiber.Ctx, err error) error {
+	code := fiber.StatusInternalServerError
+
+	var fiberErr *fiber.Error
+	if errors.As(err, &fiberErr) {
+		code = fiberErr.Code
+	} else {
+		code = StatusCodeForError(err)
+	}
+
+	c.Set(fiber.HeaderContentType, fiber.MIMETextPlainCharsetUTF8)
+	return c.Status(code).SendString(err.Error())
+}
+
 // proxyTheRequestToDefault proxies the request to the default GraphQL endpoint.
 func proxyTheRequestToDefault(c fiber.Ctx) error {
 	return proxyTheRequest(c, cfg.Server.HostGraphQL)
 }
 
+// handleWebSocketOrDefault is the handler registered on GET /v1/graphql when
+// WebSocket support is enabled. For an actual upgrade request it applies the
+// SAME ban and rate-limit checks processGraphQLRequest runs for the HTTP
+// path, before proxying the connection - previously a banned or
+// rate-limited user could bypass both simply by using the WebSocket upgrade
+// path instead of a normal POST. Only a user who passes both checks reaches
+// wsp.HandleWebSocket; behaviour for allowed users, and for non-upgrade
+// requests (proxied via proxyTheRequestToDefault), is unchanged.
+func handleWebSocketOrDefault(c fiber.Ctx) error {
+	if IsWebSocketRequest(c) {
+		wsp := GetWebSocketProxy()
+		if wsp != nil {
+			extractedUserID, extractedRoleName := extractUserInfo(c)
+
+			if checkIfUserIsBanned(c, extractedUserID) {
+				return c.Status(fiber.StatusForbidden).SendString("User is banned")
+			}
+
+			if cfg.Client.RoleRateLimit && !rateLimitedRequest(extractedUserID, extractedRoleName) {
+				return c.Status(fiber.StatusTooManyRequests).SendString("Rate limit exceeded, try again later")
+			}
+
+			return wsp.HandleWebSocket(c)
+		}
+	}
+	return proxyTheRequestToDefault(c)
+}
+
 // AddRequestUUID adds a unique request UUID to the context.
 func AddRequestUUID(c fiber.Ctx) error {
 	c.Locals("request_uuid", uuid.NewString())
@@ -428,7 +474,7 @@ func handleCaching(c fiber.Ctx, parsedResult *parseGraphQLQueryResult, userID, u
 		// No caching, just proxy the request
 		if err := proxyTheRequest(c, parsedResult.activeEndpoint); err != nil {
 			cfg.Monitoring.Increment(libpack_monitoring.MetricsFailed, nil)
-			return false, c.Status(fiber.StatusInternalServerError).SendString("Can't proxy the request - try again later")
+			return false, c.Status(StatusCodeForError(err)).SendString("Can't proxy the request - try again later")
 		}
 		return false, nil
 	}
@@ -465,7 +511,7 @@ func proxyAndCacheTheRequest(c fiber.Ctx, queryCacheHash string, cacheTime int,
 			Pairs:   map[string]any{"error": err.Error()},
 		})
 		cfg.Monitoring.Increment(libpack_monitoring.MetricsFailed, nil)
-		return c.Status(fiber.StatusInternalServerError).SendString("Can't proxy the request - try again later")
+		return c.Status(StatusCodeForError(err)).SendString("Can't proxy the request - try again later")
 	}
 
 	libpack_cache.CacheStoreWithTTL(queryCacheHash, c.Response().Body(), time.Duration(cacheTime)*time.Second)

server_handlers_test.go

@@ -3,6 +3,7 @@ package main
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"net"
 	"net/http"
@@ -13,6 +14,7 @@ import (
 
 	"github.com/gofiber/fiber/v3"
 	libpack_cache "github.com/lukaszraczylo/graphql-monitoring-proxy/cache"
+	"github.com/sony/gobreaker"
 	"github.com/valyala/fasthttp"
 )
 
@@ -702,3 +704,165 @@ func (suite *Tests) Test_shutdownHTTPProxy() {
 		suite.Fail("proxy server did not shut down; in-flight requests would be cut off")
 	}
 }
+
+// ---------------------------------------------------------------------------
+// proxyErrorHandler (C4 - circuit-open must surface as 503, not the generic
+// 500 fiber's DefaultErrorHandler would return; *fiber.Error keeps its own
+// code; any other error still maps to 500)
+// ---------------------------------------------------------------------------
+
+func TestProxyErrorHandler_MapsErrorsToStatusCodes(t *testing.T) {
+	tests := []struct {
+		name       string
+		err        error
+		wantStatus int
+	}{
+		{
+			name:       "circuit open sentinel yields 503",
+			err:        ErrCircuitOpen,
+			wantStatus: fiber.StatusServiceUnavailable,
+		},
+		{
+			name:       "gobreaker open state yields 503",
+			err:        gobreaker.ErrOpenState,
+			wantStatus: fiber.StatusServiceUnavailable,
+		},
+		{
+			name:       "wrapped circuit open sentinel yields 503",
+			err:        fmt.Errorf("proxy attempt failed: %w", ErrCircuitOpen),
+			wantStatus: fiber.StatusServiceUnavailable,
+		},
+		{
+			name:       "generic error yields 500",
+			err:        errors.New("boom"),
+			wantStatus: fiber.StatusInternalServerError,
+		},
+		{
+			name:       "proxy error keeps its own status code",
+			err:        NewProxyError(ErrCodeBadGateway, "bad gateway", fiber.StatusBadGateway, false),
+			wantStatus: fiber.StatusBadGateway,
+		},
+		{
+			name:       "fiber.Error keeps its own code",
+			err:        fiber.NewError(fiber.StatusTeapot, "teapot"),
+			wantStatus: fiber.StatusTeapot,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			app := fiber.New(fiber.Config{ErrorHandler: proxyErrorHandler})
+			app.Get("/err", func(c fiber.Ctx) error {
+				return tt.err
+			})
+
+			req := httptest.NewRequest("GET", "/err", nil)
+			resp, err := app.Test(req, fiber.TestConfig{Timeout: 5000 * time.Millisecond})
+			if err != nil {
+				t.Fatalf("app.Test: %v", err)
+			}
+			defer func() { _ = resp.Body.Close() }()
+
+			if resp.StatusCode != tt.wantStatus {
+				t.Errorf("want status %d, got %d", tt.wantStatus, resp.StatusCode)
+			}
+		})
+	}
+}
+
+// TestStartHTTPProxy_ErrorHandlerConfigured is a narrow regression check that
+// StartHTTPProxy actually wires proxyErrorHandler into the fiber app (the
+// bug this fixes was a missing ErrorHandler in that Config), rather than
+// only testing proxyErrorHandler in isolation above.
+func TestStartHTTPProxy_ErrorHandlerConfigured(t *testing.T) {
+	parseConfig()
+	_ = StartMonitoringServer()
+
+	l, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("net.Listen: %v", err)
+	}
+	port := l.Addr().(*net.TCPAddr).Port
+	_ = l.Close()
+
+	cfgMutex.Lock()
+	origPort := cfg.Server.PortGraphQL
+	origTimeout := cfg.Client.ClientTimeout
+	origWS := cfg.WebSocket.Enable
+	origAdmin := cfg.AdminDashboard.Enable
+	origHost := cfg.Server.HostGraphQL
+	origHostRO := cfg.Server.HostGraphQLReadOnly
+	cfg.Server.PortGraphQL = port
+	cfg.Client.ClientTimeout = 5
+	cfg.WebSocket.Enable = false
+	cfg.AdminDashboard.Enable = false
+	// Point at a host that always refuses connections so proxyTheRequest
+	// fails and processGraphQLRequest's proxy-error path (mapped via
+	// StatusCodeForError, not a bare 500) gets exercised.
+	cfg.Server.HostGraphQL = "http://127.0.0.1:1"
+	cfg.Server.HostGraphQLReadOnly = "http://127.0.0.1:1"
+	cfgMutex.Unlock()
+
+	t.Cleanup(func() {
+		cfgMutex.Lock()
+		cfg.Server.PortGraphQL = origPort
+		cfg.Client.ClientTimeout = origTimeout
+		cfg.WebSocket.Enable = origWS
+		cfg.AdminDashboard.Enable = origAdmin
+		cfg.Server.HostGraphQL = origHost
+		cfg.Server.HostGraphQLReadOnly = origHostRO
+		cfgMutex.Unlock()
+	})
+
+	errCh := make(chan error, 1)
+	go func() {
+		errCh <- StartHTTPProxy()
+	}()
+
+	deadline := time.Now().Add(3 * time.Second)
+	var serving bool
+	for time.Now().Before(deadline) {
+		httpProxyMu.RLock()
+		serving = httpProxyApp != nil
+		httpProxyMu.RUnlock()
+		if serving {
+			break
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+	if !serving {
+		t.Fatalf("server did not start serving on port %d within 3s", port)
+	}
+
+	httpResp, err := http.Post(
+		fmt.Sprintf("http://127.0.0.1:%d/v1/graphql", port),
+		"application/json",
+		strings.NewReader(`{"query":"query { __typename }"}`),
+	)
+	if err != nil {
+		t.Fatalf("POST /v1/graphql: %v", err)
+	}
+	_ = httpResp.Body.Close()
+
+	// An unreachable backend is a generic proxy failure, not circuit-open, so
+	// this must stay 500 - proving the ErrorHandler wiring did not change the
+	// status for genuinely-500 errors.
+	if httpResp.StatusCode != fiber.StatusInternalServerError {
+		t.Errorf("want %d for unreachable backend, got %d", fiber.StatusInternalServerError, httpResp.StatusCode)
+	}
+
+	shutdownCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+	if err := shutdownHTTPProxy(shutdownCtx); err != nil {
+		t.Fatalf("shutdownHTTPProxy: %v", err)
+	}
+
+	select {
+	case err := <-errCh:
+		if err != nil {
+			t.Errorf("StartHTTPProxy returned error after shutdown: %v", err)
+		}
+	case <-time.After(3 * time.Second):
+		t.Fatal("StartHTTPProxy did not return after shutdownHTTPProxy")
+	}
+}

websocket_test.go

@@ -11,6 +11,7 @@ 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"
@@ -578,3 +579,134 @@ func TestWebSocketProxy_HandleConnection_PingGoroutineJoinsBeforeReturn(t *testi
 		}, 2*time.Second, time.Millisecond, "iteration %d: handleConnection did not tear down", i)
 	}
 }
+
+// ---------------------------------------------------------------------------
+// handleWebSocketOrDefault (S3 - the WS route must run the same ban and
+// rate-limit checks the HTTP path runs in processGraphQLRequest, instead of
+// bypassing them by going straight to wsp.HandleWebSocket)
+// ---------------------------------------------------------------------------
+
+// setTestWebSocketProxy overrides the package-global WebSocket proxy
+// singleton for the duration of a test, restoring the previous value via
+// t.Cleanup. handleWebSocketOrDefault resolves the proxy through
+// GetWebSocketProxy(), and InitializeWebSocketProxy's sync.Once means only
+// the first caller in the whole test binary can set it that way - assigning
+// the global directly here keeps these tests deterministic regardless of
+// what other tests in this package already initialized.
+func setTestWebSocketProxy(t *testing.T, wsp *WebSocketProxy) {
+	t.Helper()
+	prev := webSocketProxy
+	webSocketProxy = wsp
+	t.Cleanup(func() { webSocketProxy = prev })
+}
+
+// newWSUpgradeRequest builds a GET request carrying the minimum headers
+// IsWebSocketRequest recognises as an upgrade attempt.
+func newWSUpgradeRequest() *http.Request {
+	req := httptest.NewRequest("GET", "/v1/graphql", nil)
+	req.Header.Set("Connection", "Upgrade")
+	req.Header.Set("Upgrade", "websocket")
+	return req
+}
+
+func TestHandleWebSocketOrDefault_BannedUserRejected(t *testing.T) {
+	parseConfig()
+
+	replaceBannedUsers(map[string]string{defaultValue: "testing"})
+	t.Cleanup(func() { replaceBannedUsers(map[string]string{}) })
+
+	setTestWebSocketProxy(t, NewWebSocketProxy("http://127.0.0.1:1", WebSocketConfig{Enabled: true}, cfg.Logger, nil))
+
+	app := fiber.New()
+	app.Get("/v1/graphql", handleWebSocketOrDefault)
+
+	resp, err := app.Test(newWSUpgradeRequest(), fiber.TestConfig{Timeout: 5000 * time.Millisecond})
+	if err != nil {
+		t.Fatalf("app.Test: %v", err)
+	}
+	defer func() { _ = resp.Body.Close() }()
+
+	// Same status processGraphQLRequest returns for a banned user.
+	if resp.StatusCode != fiber.StatusForbidden {
+		t.Fatalf("want %d (banned, same as HTTP path), got %d", fiber.StatusForbidden, resp.StatusCode)
+	}
+}
+
+func TestHandleWebSocketOrDefault_RateLimitedUserRejected(t *testing.T) {
+	parseConfig()
+
+	replaceBannedUsers(map[string]string{})
+
+	cfgMutex.Lock()
+	cfg.Client.RoleRateLimit = true
+	cfgMutex.Unlock()
+	t.Cleanup(func() {
+		cfgMutex.Lock()
+		cfg.Client.RoleRateLimit = false
+		cfgMutex.Unlock()
+	})
+
+	// Req: 0 denies the request outright - after the first Incr the ticker's
+	// rate is already > 0, so this role is denied deterministically on the
+	// very first call, no warm-up requests needed.
+	rateLimitMu.Lock()
+	rateLimits = map[string]RateLimitConfig{
+		defaultValue: {
+			RateCounterTicker: goratecounter.NewRateCounter().WithConfig(goratecounter.RateCounterConfig{
+				Interval: 1 * time.Second,
+			}),
+			Interval: 1 * time.Second,
+			Req:      0,
+		},
+	}
+	rateLimitConfigAtomic.Store(rateLimits)
+	rateLimitMu.Unlock()
+
+	setTestWebSocketProxy(t, NewWebSocketProxy("http://127.0.0.1:1", WebSocketConfig{Enabled: true}, cfg.Logger, nil))
+
+	app := fiber.New()
+	app.Get("/v1/graphql", handleWebSocketOrDefault)
+
+	resp, err := app.Test(newWSUpgradeRequest(), fiber.TestConfig{Timeout: 5000 * time.Millisecond})
+	if err != nil {
+		t.Fatalf("app.Test: %v", err)
+	}
+	defer func() { _ = resp.Body.Close() }()
+
+	// Same status processGraphQLRequest returns for a rate-limited user.
+	if resp.StatusCode != fiber.StatusTooManyRequests {
+		t.Fatalf("want %d (rate limited, same as HTTP path), got %d", fiber.StatusTooManyRequests, resp.StatusCode)
+	}
+}
+
+// TestHandleWebSocketOrDefault_AllowedUserProceedsToHandleWebSocket verifies
+// an allowed user (not banned, not rate-limited) is not blocked before
+// reaching wsp.HandleWebSocket. wsp.enabled=false is used as a cheap,
+// deterministic probe: HandleWebSocket then returns a distinct 501 (see
+// WebSocketProxy.HandleWebSocket), a status neither the ban (403) nor
+// rate-limit (429) rejection, nor a plain proxied request, would produce -
+// so seeing 501 proves control reached HandleWebSocket.
+func TestHandleWebSocketOrDefault_AllowedUserProceedsToHandleWebSocket(t *testing.T) {
+	parseConfig()
+
+	replaceBannedUsers(map[string]string{})
+
+	cfgMutex.Lock()
+	cfg.Client.RoleRateLimit = false
+	cfgMutex.Unlock()
+
+	setTestWebSocketProxy(t, NewWebSocketProxy("http://127.0.0.1:1", WebSocketConfig{Enabled: false}, cfg.Logger, nil))
+
+	app := fiber.New()
+	app.Get("/v1/graphql", handleWebSocketOrDefault)
+
+	resp, err := app.Test(newWSUpgradeRequest(), fiber.TestConfig{Timeout: 5000 * time.Millisecond})
+	if err != nil {
+		t.Fatalf("app.Test: %v", err)
+	}
+	defer func() { _ = resp.Body.Close() }()
+
+	if resp.StatusCode != fiber.StatusNotImplemented {
+		t.Fatalf("want %d (proves control reached wsp.HandleWebSocket), got %d", fiber.StatusNotImplemented, resp.StatusCode)
+	}
+}