diff --git a/internal/proxy/timeouts_test.go b/internal/proxy/timeouts_test.go index e964d57..8a78a2b 100644 --- a/internal/proxy/timeouts_test.go +++ b/internal/proxy/timeouts_test.go @@ -5,8 +5,10 @@ import ( "io" "net" "net/http" + "net/http/httptest" "strconv" "sync" + "sync/atomic" "testing" "time" @@ -66,32 +68,47 @@ func TestRequestTimeouts(t *testing.T) { t.Parallel() var ( + app *httptest.Server appURL string sendRequest func(*testing.T, string) net.Conn + appGotBody atomic.Bool ) if tc.appTakesNothing { appURL, sendRequest = startAppThatTakesNothing(t), sendLargeBody } else { - appURL, sendRequest = startApp(t, readBody).URL, sendPartOfBody + app = startApp(t, func(_ http.ResponseWriter, r *http.Request) { + n, _ := io.Copy(io.Discard, r.Body) + appGotBody.Store(n > 0) + }) + appURL, sendRequest = app.URL, sendPartOfBody } addr, out := startProxy(t, appURL, tc.env) start := time.Now() - conn := sendRequest(t, addr) - - wantStatus(t, readResponse(t, conn), tc.want) + got := readResponse(t, sendRequest(t, addr)) wantTimedOut(t, start) - wantLine(t, out.requestLine(t), tc.want, requestlog.ActionTimedOut) + + want := tc.want + + if app != nil { + // Close returns once the app has finished with the request. + app.Close() + + // Until some of the body has reached the app, smallwebwaf + // waits on the app, and SPEC.md asks for 504; the timeout + // runs out then only if the test process is held up. + if !appGotBody.Load() { + want = http.StatusGatewayTimeout + } + } + + wantStatus(t, got, want) + wantLine(t, out.requestLine(t), want, requestlog.ActionTimedOut) }) } } -// readBody is an app that reads the request body, then answers. -func readBody(_ http.ResponseWriter, r *http.Request) { - _, _ = io.Copy(io.Discard, r.Body) -} - // startAppThatTakesNothing starts an app that accepts connections and // never reads from them, and returns its URL. func startAppThatTakesNothing(t *testing.T) string {