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 {