diff --git a/internal/rproxy/limits_test.go b/internal/rproxy/limits_test.go index 2af3b1f..4247692 100644 --- a/internal/rproxy/limits_test.go +++ b/internal/rproxy/limits_test.go @@ -12,7 +12,6 @@ import ( "os" "strconv" "strings" - "sync" "sync/atomic" "testing" "time" @@ -918,49 +917,48 @@ func TestAVisitorThatStopsALittleShortHasTheConnectionClosedWhenTheAnswerIsCutOf ) { t.Parallel() - const limit = 300 * time.Millisecond + // A busy host passes on a part sent at once well within this limit. + const limit = 2 * time.Second - // With two parts, rproxy fails to send the first and then waits for the - // second, which the app sends only once rproxy has ended the request to it. + // The app sends the parts in atOnce with its status and headers. It sends the + // parts in late only once rproxy has ended the request to it, which rproxy + // does when it cuts off the answer, so that those never reach the visitor. for _, tc := range []struct { - name string - parts []string + name string + atOnce, late []string }{ - {"with the body in one part", []string{greeting}}, - {"with the body in two parts", []string{greeting[:5], greeting[5:]}}, + {"with the body in one part", nil, []string{greeting}}, + {"with the body in two parts", nil, []string{greeting[:5], greeting[5:]}}, + { + "with the first part of the body at once", + []string{greeting[:5]}, []string{greeting[5:]}, + }, } { t.Run(tc.name, func(t *testing.T) { t.Parallel() - headersRead := make(chan struct{}) - // letAppSendBody also runs when the test ends, as the app's server waits - // for the app. - letAppSendBody := sync.OnceFunc(func() { close(headersRead) }) standIn := startWAF(t, allow) - // The app sends its body once the visitor has the status and headers and - // the response time limit has run out since, so that rproxy cannot send - // it. After each part, it reads the request's body, which ends when rproxy - // ends the request to it. app := startApp(t, func(w http.ResponseWriter, r *http.Request) { rc := http.NewResponseController(w) _ = rc.EnableFullDuplex() w.Header().Set(contentLengthHeader, strconv.Itoa(len(greeting))) + for _, part := range tc.atOnce { + _, _ = io.WriteString(w, part) + } + _ = rc.Flush() - <-headersRead - time.Sleep(limit) + // The request's body ends when rproxy ends the request to the app. + _, _ = io.Copy(io.Discard, r.Body) - for _, part := range tc.parts { + for _, part := range tc.late { _, _ = io.WriteString(w, part) _ = rc.Flush() - _, _ = io.Copy(io.Discard, r.Body) } }) - t.Cleanup(letAppSendBody) - // The request time limit is far longer than dial's deadline. addr, _ := startRProxy(t, standIn.URL, app.URL, map[string]string{ maxRequestTime: "120000", @@ -975,15 +973,14 @@ func TestAVisitorThatStopsALittleShortHasTheConnectionClosedWhenTheAnswerIsCutOf t.Fatalf("read the answer: %v", err) } - letAppSendBody() - // If the connection stays open, dial's deadline ends this read with a // timeout. body, err := io.ReadAll(res.Body) _ = res.Body.Close() - if len(body) != 0 || !errors.Is(err, io.ErrUnexpectedEOF) { - t.Errorf("read %q (%v), want the connection closed", body, err) + want := strings.Join(tc.atOnce, "") + if string(body) != want || !errors.Is(err, io.ErrUnexpectedEOF) { + t.Errorf("read %q (%v), want %q and the connection closed", body, err, want) } }) }