From bc7c5e18a536986d445faa93ab931150dcf26a90 Mon Sep 17 00:00:00 2001 From: sneak Date: Wed, 14 Oct 2026 13:17:52 +0000 Subject: [PATCH] Answer 504, not 408, when a request outlasts its time limit (closes #70) The hand-written time limit middleware answered 408 and wrote that answer from its own goroutine while the handler could still be writing. It is replaced by chi's middleware.Timeout, placed after JSONResponseMiddleware so its 504 is not lost behind a 200. When a handler's database read or other lookup fails, it asks deadlinePassed whether the request's time limit passed and, if so, answers 504 instead of its 500 or 404. JSONValidationMiddleware answers 504 with the JSON error body when the limit has passed and the handler wrote nothing. status.json and /api/v1/stats answer 504 when their 4-second database read runs out. Model: opus-5-5 --- internal/server/handlers.go | 77 +++++++++++--- internal/server/handlers_test.go | 173 +++++++++++++++++++++++++++++++ internal/server/middleware.go | 148 +++----------------------- internal/server/routes.go | 7 +- 4 files changed, 255 insertions(+), 150 deletions(-) diff --git a/internal/server/handlers.go b/internal/server/handlers.go index eed2271..dec6355 100644 --- a/internal/server/handlers.go +++ b/internal/server/handlers.go @@ -201,12 +201,16 @@ func (s *Server) handleStatusJSON() http.HandlerFunc { select { case <-ctx.Done(): s.logger.Error("Database stats timeout in status.json") - writeJSONError(w, http.StatusRequestTimeout, "Database timeout") + writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage) return case err := <-errChan: s.logger.Error("Failed to get database stats", "error", err) - writeJSONError(w, http.StatusInternalServerError, err.Error()) + if deadlinePassed(ctx) { + writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage) + } else { + writeJSONError(w, http.StatusInternalServerError, err.Error()) + } return case dbStats = <-statsChan: @@ -386,7 +390,7 @@ func (s *Server) handleStats() http.HandlerFunc { // Check if context is already cancelled select { case <-ctx.Done(): - http.Error(w, "Request timeout", http.StatusRequestTimeout) + writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage) return default: @@ -415,11 +419,16 @@ func (s *Server) handleStats() http.HandlerFunc { select { case <-ctx.Done(): s.logger.Error("Database stats timeout") - // Don't write response here - timeout middleware already handles it + writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage) + return case err := <-errChan: s.logger.Error("Failed to get database stats", "error", err) - writeJSONError(w, http.StatusInternalServerError, err.Error()) + if deadlinePassed(ctx) { + writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage) + } else { + writeJSONError(w, http.StatusInternalServerError, err.Error()) + } return case dbStats = <-statsChan: @@ -622,6 +631,13 @@ func (s *Server) handleIPInfo() http.HandlerFunc { } } + // The time limit passed while resolving host names. + if deadlinePassed(ctx) { + writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage) + + return + } + if len(ipsToLookup) == 0 { writeJSONError(w, http.StatusBadRequest, "No valid IPs or hostnames provided") @@ -664,6 +680,13 @@ func (s *Server) handleIPInfo() http.HandlerFunc { } } + // The time limit passed while looking up the IPs. + if deadlinePassed(ctx) { + writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage) + + return + } + // Return response (even if no results, include errors) if len(response.Results) == 0 && len(response.Errors) > 0 { writeJSONError(w, http.StatusNotFound, "No routes found: "+response.Errors[0]) @@ -703,9 +726,12 @@ func (s *Server) handleASDetailJSON() http.HandlerFunc { asInfo, prefixes, err := s.db.GetASDetailsContext(r.Context(), asn) if err != nil { - if errors.Is(err, database.ErrNoRoute) { + switch { + case errors.Is(err, database.ErrNoRoute): writeJSONError(w, http.StatusNotFound, err.Error()) - } else { + case deadlinePassed(r.Context()): + writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage) + default: writeJSONError(w, http.StatusInternalServerError, err.Error()) } @@ -761,9 +787,12 @@ func (s *Server) handlePrefixDetailJSON() http.HandlerFunc { routes, err := s.db.GetPrefixDetailsContext(r.Context(), prefix) if err != nil { - if errors.Is(err, database.ErrNoRoute) { + switch { + case errors.Is(err, database.ErrNoRoute): writeJSONError(w, http.StatusNotFound, err.Error()) - } else { + case deadlinePassed(r.Context()): + writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage) + default: writeJSONError(w, http.StatusInternalServerError, err.Error()) } @@ -803,9 +832,12 @@ func (s *Server) handleASDetail() http.HandlerFunc { asInfo, prefixes, err := s.db.GetASDetailsContext(r.Context(), asn) if err != nil { - if errors.Is(err, database.ErrNoRoute) { + switch { + case errors.Is(err, database.ErrNoRoute): http.Error(w, "AS not found", http.StatusNotFound) - } else { + case deadlinePassed(r.Context()): + http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout) + default: s.logger.Error("Failed to get AS details", "error", err) http.Error(w, "Internal server error", http.StatusInternalServerError) } @@ -932,9 +964,12 @@ func (s *Server) handlePrefixDetail() http.HandlerFunc { routes, err := s.db.GetPrefixDetailsContext(r.Context(), prefix) if err != nil { - if errors.Is(err, database.ErrNoRoute) { + switch { + case errors.Is(err, database.ErrNoRoute): http.Error(w, "Prefix not found", http.StatusNotFound) - } else { + case deadlinePassed(r.Context()): + http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout) + default: s.logger.Error("Failed to get prefix details", "error", err) http.Error(w, "Internal server error", http.StatusInternalServerError) } @@ -1088,8 +1123,12 @@ func (s *Server) handlePrefixLength() http.HandlerFunc { const maxPrefixes = 500 prefixes, err := s.db.GetRandomPrefixesByLengthContext(r.Context(), maskLength, ipVersion, maxPrefixes) if err != nil { - s.logger.Error("Failed to get prefixes by length", "error", err) - http.Error(w, "Internal server error", http.StatusInternalServerError) + if deadlinePassed(r.Context()) { + http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout) + } else { + s.logger.Error("Failed to get prefixes by length", "error", err) + http.Error(w, "Internal server error", http.StatusInternalServerError) + } return } @@ -1179,8 +1218,12 @@ func (s *Server) handlePrefixLength6() http.HandlerFunc { const maxPrefixes = 500 prefixes, err := s.db.GetRandomPrefixesByLengthContext(r.Context(), maskLength, ipVersion, maxPrefixes) if err != nil { - s.logger.Error("Failed to get prefixes by length", "error", err) - http.Error(w, "Internal server error", http.StatusInternalServerError) + if deadlinePassed(r.Context()) { + http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout) + } else { + s.logger.Error("Failed to get prefixes by length", "error", err) + http.Error(w, "Internal server error", http.StatusInternalServerError) + } return } diff --git a/internal/server/handlers_test.go b/internal/server/handlers_test.go index 927938c..72f46e2 100644 --- a/internal/server/handlers_test.go +++ b/internal/server/handlers_test.go @@ -200,6 +200,179 @@ func TestPrefixDetailHandlersParsePrefix(t *testing.T) { } } +// shortRequestDeadline is the deadline the tests below give a request. The +// server's 30-second limit and the stats handlers' 4-second limit inherit it, +// so the tests do not wait for either. +const shortRequestDeadline = 50 * time.Millisecond + +// TestRequestPastTimeLimitAnswers504 sends a request through the server's +// router to a page and to a JSON route whose handlers wait for the time limit +// to pass and return without answering. The page answers 504; the JSON route +// answers 504 with the JSON error body. Before the fix both answered 408, +// written by the time limit middleware from its own goroutine while the +// handler could still be writing. +func TestRequestPastTimeLimitAnswers504(t *testing.T) { + s := New(nil, streamer.New(logger.New(), metrics.New()), logger.New(), &config.Config{}) + + waitForTimeLimit := func(_ http.ResponseWriter, r *http.Request) { + <-r.Context().Done() + } + // Registered the way setupRoutes registers pages and JSON routes. + s.router.Get("/slow-page", waitForTimeLimit) + s.router.Get("/slow.json", JSONValidationMiddleware(http.HandlerFunc(waitForTimeLimit)).ServeHTTP) + + tests := []struct { + path string + wantJSONBody bool + }{ + {path: "/slow-page"}, + {path: "/slow.json", wantJSONBody: true}, + } + + for _, tt := range tests { + ctx, cancel := context.WithTimeout(context.Background(), shortRequestDeadline) + rec := httptest.NewRecorder() + s.router.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, tt.path, nil).WithContext(ctx)) + cancel() + + if rec.Code != http.StatusGatewayTimeout { + t.Errorf("%s: status %d, want %d; body %s", tt.path, rec.Code, http.StatusGatewayTimeout, rec.Body) + } + if tt.wantJSONBody { + checkJSONErrorCode(t, tt.path, rec, http.StatusGatewayTimeout) + } + } +} + +// TestStatsHandlersAnswer504WhenDatabaseReadRunsOut calls each stats handler +// with a store that does not answer before the handler's 4-second limit runs +// out, and checks that the handler itself answers 504 with the JSON error +// body. Before the fix status.json answered 408, and /api/v1/stats wrote +// nothing, which the JSON check on its route turned into 500. +func TestStatsHandlersAnswer504WhenDatabaseReadRunsOut(t *testing.T) { + release := make(chan struct{}) + defer close(release) + db := blockingStatsDB{release: release} + s := New(db, streamer.New(logger.New(), metrics.New()), logger.New(), &config.Config{}) + + handlers := map[string]http.HandlerFunc{ + "status.json": s.handleStatusJSON(), + "stats": s.handleStats(), + } + + for name, handler := range handlers { + ctx, cancel := context.WithTimeout(context.Background(), shortRequestDeadline) + rec := httptest.NewRecorder() + handler(rec, httptest.NewRequest(http.MethodGet, "/", nil).WithContext(ctx)) + cancel() + + if rec.Code != http.StatusGatewayTimeout { + t.Errorf("%s: status %d, want %d; body %s", name, rec.Code, http.StatusGatewayTimeout, rec.Body) + } + checkJSONErrorCode(t, name, rec, http.StatusGatewayTimeout) + } +} + +// slowDB embeds database.Store (left nil) and overrides the reads the routes +// in TestRoutesAnswer504WhenDatabaseReadIsCutOff make. Each read waits until +// its context ends and returns the context's error, as a slow database read +// cut off by a time limit does. +type slowDB struct { + database.Store +} + +func (slowDB) GetStatsContext(ctx context.Context) (database.Stats, error) { + <-ctx.Done() + + return database.Stats{}, ctx.Err() +} + +func (slowDB) GetIPInfoContext(ctx context.Context, _ string) (*database.IPInfo, error) { + <-ctx.Done() + + return nil, ctx.Err() +} + +func (slowDB) GetASDetailsContext(ctx context.Context, _ int) (*database.ASN, []database.LiveRoute, error) { + <-ctx.Done() + + return nil, nil, ctx.Err() +} + +func (slowDB) GetPrefixDetailsContext(ctx context.Context, _ string) ([]database.LiveRoute, error) { + <-ctx.Done() + + return nil, ctx.Err() +} + +func (slowDB) GetRandomPrefixesByLengthContext(ctx context.Context, _, _, _ int) ([]database.LiveRoute, error) { + <-ctx.Done() + + return nil, ctx.Err() +} + +// TestRoutesAnswer504WhenDatabaseReadIsCutOff sends a request through the +// server's router for each route that reads the database, with a store whose +// reads fail when the request's time limit passes. Each route answers 504, the +// JSON routes with the JSON error body. Before the fix they answered 408 or, +// when the handler wrote first, its own error: 500, or 404 on the IP lookups. +// The health check is left out: it answers 503 when its own 2-second ping runs +// out. +func TestRoutesAnswer504WhenDatabaseReadIsCutOff(t *testing.T) { + s := New(slowDB{}, streamer.New(logger.New(), metrics.New()), logger.New(), &config.Config{}) + + tests := []struct { + path string + wantJSONBody bool + }{ + {path: "/as/64500"}, + {path: "/prefix/192.0.2.0/24"}, + {path: "/prefixlength/24"}, + {path: "/prefixlength6/48"}, + {path: "/status.json", wantJSONBody: true}, + {path: "/ip/192.0.2.1", wantJSONBody: true}, + {path: "/api/v1/stats", wantJSONBody: true}, + {path: "/api/v1/ip/192.0.2.1", wantJSONBody: true}, + {path: "/api/v1/as/64500", wantJSONBody: true}, + {path: "/api/v1/prefix/192.0.2.0/24", wantJSONBody: true}, + } + + for _, tt := range tests { + ctx, cancel := context.WithTimeout(context.Background(), shortRequestDeadline) + rec := httptest.NewRecorder() + s.router.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, tt.path, nil).WithContext(ctx)) + cancel() + + if rec.Code != http.StatusGatewayTimeout { + t.Errorf("%s: status %d, want %d; body %s", tt.path, rec.Code, http.StatusGatewayTimeout, rec.Body) + } + if tt.wantJSONBody { + checkJSONErrorCode(t, tt.path, rec, http.StatusGatewayTimeout) + } + } +} + +// checkJSONErrorCode checks that rec holds the JSON error body with the given +// code. +func checkJSONErrorCode(t *testing.T, name string, rec *httptest.ResponseRecorder, want int) { + t.Helper() + + var body struct { + Status string `json:"status"` + Error struct { + Code int `json:"code"` + } `json:"error"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil { + t.Errorf("%s: decoding the answer %q: %v", name, rec.Body, err) + + return + } + if body.Status != "error" || body.Error.Code != want { + t.Errorf("%s: answer %s, want status \"error\" and error code %d", name, rec.Body, want) + } +} + // settledGoroutineCount lets transient goroutines finish, then reports the // current count. func settledGoroutineCount() int { diff --git a/internal/server/middleware.go b/internal/server/middleware.go index dd61ee8..24711c2 100644 --- a/internal/server/middleware.go +++ b/internal/server/middleware.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "encoding/json" + "errors" "fmt" "log/slog" "net/http" @@ -109,142 +110,24 @@ func JSONResponseMiddleware(next http.Handler) http.Handler { }) } -// timeoutWriter wraps ResponseWriter to prevent concurrent writes after timeout -type timeoutWriter struct { - http.ResponseWriter - mu sync.Mutex - written bool - header http.Header // cached header to prevent concurrent access -} +// requestTimeoutMessage is the error message sent with 504 Gateway Timeout. +const requestTimeoutMessage = "Request timeout" -func (tw *timeoutWriter) Write(b []byte) (int, error) { - tw.mu.Lock() - defer tw.mu.Unlock() - - if tw.written { - return 0, nil // Discard writes after timeout - } - - return tw.ResponseWriter.Write(b) -} - -func (tw *timeoutWriter) WriteHeader(statusCode int) { - tw.mu.Lock() - defer tw.mu.Unlock() - - if tw.written { - return // Discard writes after timeout - } - - tw.ResponseWriter.WriteHeader(statusCode) -} - -func (tw *timeoutWriter) Header() http.Header { - tw.mu.Lock() - defer tw.mu.Unlock() - - if tw.written { - // Return a copy to prevent modifications after timeout - if tw.header == nil { - tw.header = make(http.Header) - } - - return tw.header - } - - return tw.ResponseWriter.Header() -} - -func (tw *timeoutWriter) markWritten() { - tw.mu.Lock() - defer tw.mu.Unlock() - tw.written = true -} - -// TimeoutMiddleware creates an HTTP middleware that enforces a request timeout. -// If the handler does not complete within the specified duration, the middleware -// returns a JSON error response with HTTP status 408 (Request Timeout). -// -// The timeout parameter specifies the maximum duration allowed for request processing. -// The returned middleware handles panics from the wrapped handler by re-panicking -// after cleanup, and prevents concurrent writes to the response after timeout occurs. -func TimeoutMiddleware(timeout time.Duration, logger *slog.Logger) func(http.Handler) http.Handler { - return func(next http.Handler) http.Handler { - return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - startTime := time.Now() - - ctx, cancel := context.WithTimeout(r.Context(), timeout) - defer cancel() - - tw := &timeoutWriter{ - ResponseWriter: w, - header: make(http.Header), - } - - done := make(chan struct{}) - panicChan := make(chan interface{}, 1) - - go func() { - defer func() { - if p := recover(); p != nil { - panicChan <- p - } - }() - - next.ServeHTTP(tw, r.WithContext(ctx)) - close(done) - }() - - select { - case p := <-panicChan: - panic(p) - case <-done: - return - case <-ctx.Done(): - tw.markWritten() // Prevent the handler from writing after timeout - execTime := time.Since(startTime) - - // Log the timeout as a warning - logger.Warn("Request timeout", - "method", r.Method, - "path", r.URL.Path, - "duration_ms", execTime.Milliseconds(), - "remote_addr", r.RemoteAddr, - ) - - // Write directly to the underlying writer since we've marked tw as written - // This is safe because markWritten() prevents the handler from writing - tw.mu.Lock() - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(http.StatusRequestTimeout) - tw.mu.Unlock() - - response := map[string]interface{}{ - "status": "error", - "error": map[string]interface{}{ - "msg": "Request timeout", - "code": http.StatusRequestTimeout, - }, - "@meta": map[string]interface{}{ - "time_zone": "UTC", - "api_version": 1, - "execution_time": fmt.Sprintf("%d ms", execTime.Milliseconds()), - }, - } - - _ = json.NewEncoder(w).Encode(response) - } - }) - } +// deadlinePassed reports whether ctx ended because its deadline passed. When +// a handler's database read or other lookup fails, it uses this to tell a +// request cut off by its time limit, which answers 504, from other failures. +func deadlinePassed(ctx context.Context) bool { + return errors.Is(ctx.Err(), context.DeadlineExceeded) } // JSONValidationMiddleware is an HTTP middleware that validates JSON API responses. // It ensures that responses with Content-Type "application/json" contain valid JSON. // // If a response is not valid JSON or is empty when JSON is expected, the middleware -// returns a properly formatted JSON error response. For timeout errors (status 408), -// the error message will be "Request timeout". For other errors, it returns -// "Internal server error" with status 500 if the original status was 200. +// returns a properly formatted JSON error response. If the request's time limit +// has passed, it answers 504 with the error message "Request timeout". For other +// errors, it returns "Internal server error" with status 500 if the original +// status was 200. // // Non-JSON responses are passed through unchanged. func JSONValidationMiddleware(next http.Handler) http.Handler { @@ -293,11 +176,14 @@ func JSONValidationMiddleware(next http.Handler) http.Handler { if statusCode == http.StatusOK { statusCode = http.StatusInternalServerError } + if deadlinePassed(r.Context()) { + statusCode = http.StatusGatewayTimeout + } w.WriteHeader(statusCode) errorMsg := "Internal server error" - if statusCode == http.StatusRequestTimeout { - errorMsg = "Request timeout" + if statusCode == http.StatusGatewayTimeout { + errorMsg = requestTimeoutMessage } response := map[string]interface{}{ diff --git a/internal/server/routes.go b/internal/server/routes.go index cde3293..9227772 100644 --- a/internal/server/routes.go +++ b/internal/server/routes.go @@ -16,9 +16,12 @@ func (s *Server) setupRoutes() { r.Use(middleware.RealIP) r.Use(RequestLoggerMiddleware(s.logger.Logger)) // Structured request logging r.Use(middleware.Recoverer) - const requestTimeout = 30 * time.Second // Increased from 8s for slow queries - r.Use(TimeoutMiddleware(requestTimeout, s.logger.Logger)) r.Use(JSONResponseMiddleware) + // Timeout must come after JSONResponseMiddleware, which sends its answer + // when the handler returns and would send a 200 ahead of Timeout's 504 for + // a handler that wrote nothing. + const requestTimeout = 30 * time.Second // Increased from 8s for slow queries + r.Use(middleware.Timeout(requestTimeout)) // Routes r.Get("/", s.handleIndex())