diff --git a/admin_dashboard_test.go b/admin_dashboard_test.go index 1f66f8e..c08dd2b 100644 --- a/admin_dashboard_test.go +++ b/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 diff --git a/api.go b/api.go index 9b1124f..c73d946 100644 --- a/api.go +++ b/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 } }() diff --git a/server.go b/server.go index 821ec4c..473799e 100644 --- a/server.go +++ b/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) diff --git a/server_handlers_test.go b/server_handlers_test.go index a2dfe5c..c939c88 100644 --- a/server_handlers_test.go +++ b/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") + } +} diff --git a/websocket_test.go b/websocket_test.go index d0a49bc..6d8735c 100644 --- a/websocket_test.go +++ b/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) + } +}