next into main #87
@@ -198,19 +198,24 @@ func (s *Server) handleStatusJSON() http.HandlerFunc {
|
||||
}()
|
||||
|
||||
var dbStats database.Stats
|
||||
var err error
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
err = ctx.Err()
|
||||
case err = <-errChan:
|
||||
case dbStats = <-statsChan:
|
||||
}
|
||||
if deadlinePassed(ctx) {
|
||||
s.logger.Error("Database stats timeout in status.json")
|
||||
writeJSONError(w, http.StatusRequestTimeout, "Database timeout")
|
||||
writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
|
||||
|
||||
return
|
||||
case err := <-errChan:
|
||||
}
|
||||
if err != nil {
|
||||
s.logger.Error("Failed to get database stats", "error", err)
|
||||
writeJSONError(w, http.StatusInternalServerError, err.Error())
|
||||
|
||||
return
|
||||
case dbStats = <-statsChan:
|
||||
// Success
|
||||
}
|
||||
|
||||
uptime := time.Since(metrics.ConnectedSince).Truncate(time.Second).String()
|
||||
@@ -231,6 +236,11 @@ func (s *Server) handleStatusJSON() http.HandlerFunc {
|
||||
var whoisStats *WHOISStatsInfo
|
||||
if s.asnFetcher != nil {
|
||||
whoisStats = s.getWHOISStats(ctx)
|
||||
if deadlinePassed(ctx) {
|
||||
writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
|
||||
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
stats := Stats{
|
||||
@@ -386,7 +396,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:
|
||||
@@ -412,18 +422,24 @@ func (s *Server) handleStats() http.HandlerFunc {
|
||||
}()
|
||||
|
||||
var dbStats database.Stats
|
||||
var err error
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
err = ctx.Err()
|
||||
case err = <-errChan:
|
||||
case dbStats = <-statsChan:
|
||||
}
|
||||
if deadlinePassed(ctx) {
|
||||
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:
|
||||
}
|
||||
if err != nil {
|
||||
s.logger.Error("Failed to get database stats", "error", err)
|
||||
writeJSONError(w, http.StatusInternalServerError, err.Error())
|
||||
|
||||
return
|
||||
case dbStats = <-statsChan:
|
||||
// Success
|
||||
}
|
||||
|
||||
uptime := time.Since(metrics.ConnectedSince).Truncate(time.Second).String()
|
||||
@@ -462,6 +478,11 @@ func (s *Server) handleStats() http.HandlerFunc {
|
||||
var whoisStats *WHOISStatsInfo
|
||||
if s.asnFetcher != nil {
|
||||
whoisStats = s.getWHOISStats(ctx)
|
||||
if deadlinePassed(ctx) {
|
||||
writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
|
||||
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Calculate connection duration
|
||||
@@ -613,6 +634,11 @@ func (s *Server) handleIPInfo() http.HandlerFunc {
|
||||
} else {
|
||||
// It's a hostname - resolve it
|
||||
resolved, err := net.DefaultResolver.LookupHost(ctx, t)
|
||||
if deadlinePassed(ctx) {
|
||||
writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
|
||||
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
response.Errors = append(response.Errors, t+": "+err.Error())
|
||||
|
||||
@@ -634,6 +660,11 @@ func (s *Server) handleIPInfo() http.HandlerFunc {
|
||||
// Look up each IP
|
||||
for _, ip := range ipsToLookup {
|
||||
ipInfo, err := s.db.GetIPInfoContext(ctx, ip)
|
||||
if deadlinePassed(ctx) {
|
||||
writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
|
||||
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
response.Errors = append(response.Errors, ip+": "+err.Error())
|
||||
|
||||
@@ -642,6 +673,11 @@ func (s *Server) handleIPInfo() http.HandlerFunc {
|
||||
|
||||
// Do PTR lookup for this IP
|
||||
ptrs, err := net.DefaultResolver.LookupAddr(ctx, ip)
|
||||
if deadlinePassed(ctx) {
|
||||
writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
|
||||
|
||||
return
|
||||
}
|
||||
if err == nil && len(ptrs) > 0 {
|
||||
// Remove trailing dots from PTR records
|
||||
for i, ptr := range ptrs {
|
||||
@@ -702,6 +738,11 @@ func (s *Server) handleASDetailJSON() http.HandlerFunc {
|
||||
}
|
||||
|
||||
asInfo, prefixes, err := s.db.GetASDetailsContext(r.Context(), asn)
|
||||
if deadlinePassed(r.Context()) {
|
||||
writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
|
||||
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
if errors.Is(err, database.ErrNoRoute) {
|
||||
writeJSONError(w, http.StatusNotFound, err.Error())
|
||||
@@ -760,6 +801,11 @@ func (s *Server) handlePrefixDetailJSON() http.HandlerFunc {
|
||||
prefix := parsed.Masked().String()
|
||||
|
||||
routes, _, err := s.db.GetPrefixDetailsContext(r.Context(), prefix)
|
||||
if deadlinePassed(r.Context()) {
|
||||
writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
|
||||
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
if errors.Is(err, database.ErrNoRoute) {
|
||||
writeJSONError(w, http.StatusNotFound, err.Error())
|
||||
@@ -802,6 +848,11 @@ func (s *Server) handleASDetail() http.HandlerFunc {
|
||||
}
|
||||
|
||||
asInfo, prefixes, err := s.db.GetASDetailsContext(r.Context(), asn)
|
||||
if deadlinePassed(r.Context()) {
|
||||
http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout)
|
||||
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
if errors.Is(err, database.ErrNoRoute) {
|
||||
http.Error(w, "AS not found", http.StatusNotFound)
|
||||
@@ -815,6 +866,11 @@ func (s *Server) handleASDetail() http.HandlerFunc {
|
||||
|
||||
// Get peers
|
||||
peers, err := s.db.GetASPeersContext(r.Context(), asn)
|
||||
if deadlinePassed(r.Context()) {
|
||||
http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout)
|
||||
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
s.logger.Error("Failed to get AS peers", "error", err)
|
||||
// Continue without peers rather than failing the whole request
|
||||
@@ -931,6 +987,11 @@ func (s *Server) handlePrefixDetail() http.HandlerFunc {
|
||||
prefix := parsed.Masked().String()
|
||||
|
||||
routes, originASNs, err := s.db.GetPrefixDetailsContext(r.Context(), prefix)
|
||||
if deadlinePassed(r.Context()) {
|
||||
http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout)
|
||||
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
if errors.Is(err, database.ErrNoRoute) {
|
||||
http.Error(w, "Prefix not found", http.StatusNotFound)
|
||||
@@ -1079,6 +1140,11 @@ func (s *Server) handlePrefixLength() http.HandlerFunc {
|
||||
// Get random sample of prefixes
|
||||
const maxPrefixes = 500
|
||||
prefixes, err := s.db.GetRandomPrefixesByLengthContext(r.Context(), maskLength, ipVersion, maxPrefixes)
|
||||
if deadlinePassed(r.Context()) {
|
||||
http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout)
|
||||
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
s.logger.Error("Failed to get prefixes by length", "error", err)
|
||||
http.Error(w, "Internal server error", http.StatusInternalServerError)
|
||||
@@ -1170,6 +1236,11 @@ func (s *Server) handlePrefixLength6() http.HandlerFunc {
|
||||
// Get random sample of prefixes
|
||||
const maxPrefixes = 500
|
||||
prefixes, err := s.db.GetRandomPrefixesByLengthContext(r.Context(), maskLength, ipVersion, maxPrefixes)
|
||||
if deadlinePassed(r.Context()) {
|
||||
http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout)
|
||||
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
s.logger.Error("Failed to get prefixes by length", "error", err)
|
||||
http.Error(w, "Internal server error", http.StatusInternalServerError)
|
||||
|
||||
@@ -270,6 +270,257 @@ func TestPrefixDetailShowsOriginNames(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 then answers as the real store does when the time
|
||||
// limit cuts it off: the IP read with the context's error, the prefix read
|
||||
// that there is no route, and the others with what they had read, here
|
||||
// nothing, and no error. Every other Store method is unused and would panic
|
||||
// if called.
|
||||
type slowDB struct {
|
||||
database.Store
|
||||
}
|
||||
|
||||
func (slowDB) GetStatsContext(ctx context.Context) (database.Stats, error) {
|
||||
<-ctx.Done()
|
||||
|
||||
return database.Stats{}, nil
|
||||
}
|
||||
|
||||
func (slowDB) GetIPInfoContext(ctx context.Context, _ string) (*database.IPInfo, error) {
|
||||
<-ctx.Done()
|
||||
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
|
||||
func (slowDB) GetASDetailsContext(ctx context.Context, asn int) (*database.ASN, []database.LiveRoute, error) {
|
||||
<-ctx.Done()
|
||||
|
||||
return &database.ASN{ASN: asn}, nil, nil
|
||||
}
|
||||
|
||||
func (slowDB) GetPrefixDetailsContext(ctx context.Context, _ string) ([]database.LiveRoute, map[int]database.ASN, error) {
|
||||
<-ctx.Done()
|
||||
|
||||
return nil, nil, database.ErrNoRoute
|
||||
}
|
||||
|
||||
func (slowDB) GetRandomPrefixesByLengthContext(ctx context.Context, _, _, _ int) ([]database.LiveRoute, error) {
|
||||
<-ctx.Done()
|
||||
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// TestRoutesAnswer504WhenDatabaseReadIsCutOff sends a request through the
|
||||
// server's router for each route that reads the database, with a store whose
|
||||
// reads are cut off when the request's time limit passes. Each route answers
|
||||
// 504 with the time limit message. Without its check after the read, a route
|
||||
// answers with what the cut-off read returned: 500, 404, 200, or, on a page
|
||||
// that stops before rendering once the request has ended, 504 with an empty
|
||||
// body. 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 {
|
||||
checkRouteAnswers504(t, s, tt.path, tt.wantJSONBody)
|
||||
}
|
||||
}
|
||||
|
||||
// slowSecondReadDB embeds database.Store (left nil) and overrides the reads
|
||||
// of the routes that read the database twice. The first read answers at once;
|
||||
// the second, the AS page's peers read or the stats routes' WHOIS read, waits
|
||||
// until its context ends and then answers with the context's error, as the
|
||||
// real store does.
|
||||
type slowSecondReadDB struct {
|
||||
database.Store
|
||||
}
|
||||
|
||||
func (slowSecondReadDB) GetASDetailsContext(_ context.Context, asn int) (*database.ASN, []database.LiveRoute, error) {
|
||||
return &database.ASN{ASN: asn}, nil, nil
|
||||
}
|
||||
|
||||
func (slowSecondReadDB) GetASPeersContext(ctx context.Context, _ int) ([]database.ASPeer, error) {
|
||||
<-ctx.Done()
|
||||
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
|
||||
func (slowSecondReadDB) GetStatsContext(_ context.Context) (database.Stats, error) {
|
||||
return database.Stats{}, nil
|
||||
}
|
||||
|
||||
func (slowSecondReadDB) GetWHOISStats(ctx context.Context, _ time.Duration) (*database.WHOISStats, error) {
|
||||
<-ctx.Done()
|
||||
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
|
||||
// idleASNFetcher is a WHOIS fetcher that does nothing. The stats handlers
|
||||
// read the WHOIS stats only when the server has a fetcher.
|
||||
type idleASNFetcher struct{}
|
||||
|
||||
func (idleASNFetcher) QueueImmediate(int) {}
|
||||
|
||||
func (idleASNFetcher) GetStats() ASNFetcherStats {
|
||||
return ASNFetcherStats{}
|
||||
}
|
||||
|
||||
// TestRoutesAnswer504WhenSecondDatabaseReadIsCutOff sends a request through
|
||||
// the server's router for each route that reads the database twice, with a
|
||||
// store that answers the first read and cuts off the second when the
|
||||
// request's time limit passes. Each route answers 504 with the time limit
|
||||
// message. Without the check after the second read, the AS page answers 504
|
||||
// with an empty body and the stats routes answer 200 without the WHOIS stats.
|
||||
func TestRoutesAnswer504WhenSecondDatabaseReadIsCutOff(t *testing.T) {
|
||||
s := New(slowSecondReadDB{}, streamer.New(logger.New(), metrics.New()), logger.New(), &config.Config{})
|
||||
s.SetASNFetcher(idleASNFetcher{})
|
||||
|
||||
tests := []struct {
|
||||
path string
|
||||
wantJSONBody bool
|
||||
}{
|
||||
{path: "/as/64500"},
|
||||
{path: "/status.json", wantJSONBody: true},
|
||||
{path: "/api/v1/stats", wantJSONBody: true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
checkRouteAnswers504(t, s, tt.path, tt.wantJSONBody)
|
||||
}
|
||||
}
|
||||
|
||||
// checkRouteAnswers504 sends a request for path through the server's router
|
||||
// with a short deadline and checks that it answers 504 with the time limit
|
||||
// message: the JSON error body on JSON routes, the plain message on pages.
|
||||
func checkRouteAnswers504(t *testing.T, s *Server, path string, wantJSONBody bool) {
|
||||
t.Helper()
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), shortRequestDeadline)
|
||||
defer cancel()
|
||||
rec := httptest.NewRecorder()
|
||||
s.router.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, path, nil).WithContext(ctx))
|
||||
|
||||
if rec.Code != http.StatusGatewayTimeout {
|
||||
t.Errorf("%s: status %d, want %d; body %s", path, rec.Code, http.StatusGatewayTimeout, rec.Body)
|
||||
}
|
||||
if wantJSONBody {
|
||||
checkJSONErrorCode(t, path, rec, http.StatusGatewayTimeout)
|
||||
} else if body := strings.TrimSpace(rec.Body.String()); body != requestTimeoutMessage {
|
||||
t.Errorf("%s: body %q, want %q", path, body, requestTimeoutMessage)
|
||||
}
|
||||
}
|
||||
|
||||
// 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 {
|
||||
|
||||
+20
-130
@@ -4,6 +4,7 @@ import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
@@ -109,142 +110,28 @@ 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).
|
||||
// deadlinePassed reports whether ctx ended because its deadline passed.
|
||||
//
|
||||
// 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)
|
||||
}
|
||||
})
|
||||
}
|
||||
// A handler calls it right after each database read or host name lookup,
|
||||
// before it looks at what came back, error or not, and answers 504 if it
|
||||
// reports true. A read cut off by the time limit can return an error, or part
|
||||
// of its rows or none of them without one, so only this check tells a cut-off
|
||||
// read from a complete one.
|
||||
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 +180,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{}{
|
||||
|
||||
@@ -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())
|
||||
|
||||
Reference in New Issue
Block a user