diff --git a/internal/protocol/http/fetcher.go b/internal/protocol/http/fetcher.go index 80b1f9e3..6f54377d 100644 --- a/internal/protocol/http/fetcher.go +++ b/internal/protocol/http/fetcher.go @@ -303,8 +303,9 @@ type Fetcher struct { resolveDataPos atomic.Int64 // How many bytes downloaded during resolve // Resolve response - kept open for one-time URLs - resolveResp *http.Response - resolveRespLock sync.Mutex + resolveResp *http.Response + resolveRespCancel context.CancelFunc // Cancels the Resolve request + resolveRespLock sync.Mutex // Async prefetch during resolve phase prefetchFile *os.File // Temporary file for prefetch data @@ -425,8 +426,11 @@ func (f *Fetcher) Resolve(req *base.Request) error { // Build HTTP request WITHOUT Range header (normal request) // This allows the response to be reused for downloading (important for one-time URLs) - httpReq, err := f.buildRequest(context.TODO(), req) + // The body outlives Resolve, so the read timeout gets its own cancel func. + reqCtx, reqCancel := context.WithCancel(context.Background()) + httpReq, err := f.buildRequest(reqCtx, req) if err != nil { + reqCancel() f.setState(stateError) return err } @@ -438,6 +442,7 @@ func (f *Fetcher) Resolve(req *base.Request) error { connStartTime := time.Now() resp, err := client.Do(httpReq) if err != nil { + reqCancel() f.setState(stateError) return err } @@ -452,6 +457,7 @@ func (f *Fetcher) Resolve(req *base.Request) error { if resp.StatusCode != base.HttpCodeOK && resp.StatusCode != base.HttpCodePartialContent { resp.Body.Close() + reqCancel() f.setState(stateError) return NewRequestError(resp.StatusCode) } @@ -469,6 +475,7 @@ func (f *Fetcher) Resolve(req *base.Request) error { start, _, total, err := parseRangeResponse(resp) if err != nil { resp.Body.Close() + reqCancel() f.setState(stateError) return err } @@ -536,8 +543,10 @@ func (f *Fetcher) Resolve(req *base.Request) error { f.resolveRespLock.Lock() if reuseResolveBody { f.resolveResp = resp + f.resolveRespCancel = reqCancel } else { resp.Body.Close() + reqCancel() f.prefetchDone.Store(true) } f.resolveRespLock.Unlock() @@ -574,7 +583,7 @@ func (f *Fetcher) asyncPrefetch() { // Get the resolve response f.resolveRespLock.Lock() - resp := f.resolveResp + resp, cancel := f.resolveResp, f.resolveRespCancel f.resolveRespLock.Unlock() if resp == nil { @@ -585,10 +594,7 @@ func (f *Fetcher) asyncPrefetch() { // Close the response even if temporary-file setup fails or shutdown // rejects a late prefetch. Otherwise the server can remain blocked writing. f.resolveRespLock.Lock() - if f.resolveResp != nil { - f.resolveResp.Body.Close() - f.resolveResp = nil - } + f.closeResolveRespLocked() f.resolveRespLock.Unlock() }() @@ -618,7 +624,7 @@ func (f *Fetcher) asyncPrefetch() { f.prefetchFilePath = tmpFile.Name() buf := make([]byte, 32*1024) // 32KB buffer - reader := f.responseReader(resp) + reader := f.responseReader(resp, cancel) for { select { @@ -881,10 +887,7 @@ func (f *Fetcher) doStart() error { if !f.resolveFallback.Load() { // Also close resolve response if still open. f.resolveRespLock.Lock() - if f.resolveResp != nil { - f.resolveResp.Body.Close() - f.resolveResp = nil - } + f.closeResolveRespLocked() f.resolveRespLock.Unlock() } } @@ -1337,7 +1340,11 @@ func (f *Fetcher) downloadChunkOnce(conn *connection, client *http.Client, buf [ } f.connMu.Unlock() - httpReq, err := f.buildRequest(conn.ctx, f.meta.Req) + // The read timeout cancels this request through its own context. + reqCtx, reqCancel := context.WithCancel(conn.ctx) + defer reqCancel() + + httpReq, err := f.buildRequest(reqCtx, f.meta.Req) if err != nil { return err } @@ -1365,7 +1372,7 @@ func (f *Fetcher) downloadChunkOnce(conn *connection, client *http.Client, buf [ // Check if this might be a redirect URL expiration error // If so, try falling back to the original URL if f.hasRedirectURL() && isRedirectExpiredError(originalErr) { - fallbackResp, fallbackErr := f.tryFallbackToOriginalURL(conn.ctx, client, rangeStart, rangeEnd, rangeRequested, ifRange) + fallbackResp, fallbackErr := f.tryFallbackToOriginalURL(reqCtx, client, rangeStart, rangeEnd, rangeRequested, ifRange) if fallbackErr == nil && fallbackResp != nil { // Fallback succeeded, use this response instead resp = fallbackResp @@ -1486,7 +1493,7 @@ func (f *Fetcher) downloadChunkOnce(conn *connection, client *http.Client, buf [ f.slowStart.onConnectSuccess() } - reader := f.responseReader(resp) + reader := f.responseReader(resp, reqCancel) var responseBytesRead int64 for { if conn.ctx.Err() != nil { @@ -1840,8 +1847,8 @@ func (f *Fetcher) runConnectionWithResolveResp(conn *connection) { // Get the resolve response f.resolveRespLock.Lock() - resp := f.resolveResp - f.resolveResp = nil // Take ownership + resp, cancel := f.resolveResp, f.resolveRespCancel + f.resolveResp, f.resolveRespCancel = nil, nil // Take ownership f.resolveRespLock.Unlock() if resp == nil { @@ -1850,6 +1857,9 @@ func (f *Fetcher) runConnectionWithResolveResp(conn *connection) { return } + if cancel != nil { + defer cancel() + } defer resp.Body.Close() f.connMu.Lock() @@ -1866,7 +1876,7 @@ func (f *Fetcher) runConnectionWithResolveResp(conn *connection) { } // Download data from resolve response - reader := f.responseReader(resp) + reader := f.responseReader(resp, cancel) for { if conn.ctx.Err() != nil { return @@ -2481,10 +2491,7 @@ func (f *Fetcher) Pause() error { // Clean up resolve response if still held f.resolveRespLock.Lock() - if f.resolveResp != nil { - f.resolveResp.Body.Close() - f.resolveResp = nil - } + f.closeResolveRespLocked() f.resolveRespLock.Unlock() f.fileMu.Lock() @@ -2602,9 +2609,25 @@ func (f *Fetcher) Wait() error { // Generated output can pause while its producer waits for processing capacity. // Its producer owns upstream deadlines; cancellation still closes the response. -func (f *Fetcher) responseReader(resp *http.Response) io.Reader { +// Otherwise a stalled read cancels the request via cancel, which must belong +// to the request that produced resp. +func (f *Fetcher) responseReader(resp *http.Response, cancel context.CancelFunc) io.Reader { if f.ctl != nil && f.ctl.ManagedProduction != nil && resp.Request != nil && f.ctl.ManagedProduction(resp.Request.URL.String()) { return resp.Body } - return NewTimeoutReader(resp.Body, readTimeout) + return NewTimeoutReader(resp.Body, readTimeout, cancel) +} + +// closeResolveRespLocked closes the kept Resolve response, if any, and +// releases its request context. The caller must hold resolveRespLock. +func (f *Fetcher) closeResolveRespLocked() { + if f.resolveResp == nil { + return + } + f.resolveResp.Body.Close() + f.resolveResp = nil + if f.resolveRespCancel != nil { + f.resolveRespCancel() + f.resolveRespCancel = nil + } } diff --git a/internal/protocol/http/timeout_reader.go b/internal/protocol/http/timeout_reader.go index f23e16e9..28809a9e 100644 --- a/internal/protocol/http/timeout_reader.go +++ b/internal/protocol/http/timeout_reader.go @@ -6,35 +6,56 @@ import ( "time" ) +// TimeoutReader fails a Read that makes no progress within timeout. +// +// One timer per reader is armed when Read starts and stopped before it +// returns, so time the caller spends between reads is not counted. If the +// timer fires, it calls cancel, which should cancel the HTTP request the +// reader belongs to so the transport unblocks the pending Read. Read still +// waits for the underlying Read to return; that Read and every later one +// return context.DeadlineExceeded. Errors from any other cancellation are +// returned unchanged, and closing the response body is left to the caller. +// +// With a nil cancel the pending Read cannot be interrupted; it still reports +// context.DeadlineExceeded, but only once it returns. +// +// TimeoutReader is not safe for concurrent use. type TimeoutReader struct { - reader io.Reader - timeout time.Duration + reader io.Reader + timeout time.Duration + cancel context.CancelFunc + timer *time.Timer + timedOut bool } -func NewTimeoutReader(r io.Reader, timeout time.Duration) *TimeoutReader { +func NewTimeoutReader(r io.Reader, timeout time.Duration, cancel context.CancelFunc) *TimeoutReader { + if cancel == nil { + cancel = func() {} + } return &TimeoutReader{ reader: r, timeout: timeout, + cancel: cancel, } } func (tr *TimeoutReader) Read(p []byte) (n int, err error) { - ctx, cancel := context.WithTimeout(context.Background(), tr.timeout) - defer cancel() - - done := make(chan struct{}) - var readErr error - var bytesRead int - - go func() { - bytesRead, readErr = tr.reader.Read(p) - close(done) - }() - - select { - case <-done: - return bytesRead, readErr - case <-ctx.Done(): - return 0, ctx.Err() + if tr.timedOut { + return 0, context.DeadlineExceeded } + if tr.timer == nil { + tr.timer = time.AfterFunc(tr.timeout, tr.cancel) + } else { + tr.timer.Reset(tr.timeout) + } + + n, err = tr.reader.Read(p) + + if !tr.timer.Stop() { + // The timer fired during this Read and cancelled the request, so err is + // whatever the interrupted read produced. Report the timeout instead. + tr.timedOut = true + return n, context.DeadlineExceeded + } + return n, err } diff --git a/internal/protocol/http/timeout_reader_test.go b/internal/protocol/http/timeout_reader_test.go index 337d55b3..109c83ae 100644 --- a/internal/protocol/http/timeout_reader_test.go +++ b/internal/protocol/http/timeout_reader_test.go @@ -2,17 +2,26 @@ package http import ( "bytes" + "compress/gzip" "context" + "crypto/tls" + "crypto/x509" "errors" "io" + gohttp "net/http" + "net/http/httptest" + "sync" + "sync/atomic" "testing" "time" + + "github.com/GopeedLab/gopeed/internal/httpclient" ) func TestTimeoutReader_Read(t *testing.T) { data := []byte("Hello, World!") reader := bytes.NewReader(data) - timeoutReader := NewTimeoutReader(reader, 1*time.Second) + timeoutReader := NewTimeoutReader(reader, 1*time.Second, nil) buf := make([]byte, len(data)) n, err := timeoutReader.Read(buf) @@ -29,7 +38,7 @@ func TestTimeoutReader_Read(t *testing.T) { func TestTimeoutReader_ReadTimeout(t *testing.T) { reader := &slowReader{delay: 2 * time.Second} - timeoutReader := NewTimeoutReader(reader, 1*time.Second) + timeoutReader := NewTimeoutReader(reader, 1*time.Second, nil) buf := make([]byte, 8192) _, err := timeoutReader.Read(buf) @@ -49,3 +58,340 @@ func (sr *slowReader) Read(p []byte) (n int, err error) { time.Sleep(sr.delay) return 0, io.EOF } + +// The timer must only cover time spent inside Read, not caller work between reads. +func TestTimeoutReader_IdleBetweenReads(t *testing.T) { + var cancelled atomic.Bool + timeoutReader := NewTimeoutReader(bytes.NewReader([]byte("abcdef")), 50*time.Millisecond, func() { + cancelled.Store(true) + }) + + buf := make([]byte, 3) + if _, err := timeoutReader.Read(buf); err != nil { + t.Fatalf("first read: %v", err) + } + time.Sleep(200 * time.Millisecond) + if _, err := timeoutReader.Read(buf); err != nil { + t.Fatalf("second read after idle: %v", err) + } + if cancelled.Load() { + t.Fatal("request was cancelled while idle between reads") + } +} + +// A timed-out Read must not return while the underlying Read can still write +// into the caller's buffer, and must leave closing the reader to the caller. +func TestTimeoutReader_TimedOutReadDoesNotOutliveCaller(t *testing.T) { + reader := &lateWriteReader{cancelled: make(chan struct{})} + timeoutReader := NewTimeoutReader(reader, 50*time.Millisecond, reader.cancel) + + buf := make([]byte, 8) + _, err := timeoutReader.Read(buf) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("expected %v, got %v", context.DeadlineExceeded, err) + } + if !reader.returned.Load() { + t.Fatal("Read returned before the underlying Read finished") + } + if _, err := timeoutReader.Read(buf); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("read after timeout: expected %v, got %v", context.DeadlineExceeded, err) + } + if reader.closed.Load() { + t.Fatal("reader was closed by the timeout") + } +} + +// lateWriteReader blocks until cancelled, then writes into p before returning. +type lateWriteReader struct { + cancelOnce sync.Once + cancelled chan struct{} + returned atomic.Bool + closed atomic.Bool +} + +func (r *lateWriteReader) Read(p []byte) (int, error) { + <-r.cancelled + time.Sleep(20 * time.Millisecond) + p[0] = 'x' + r.returned.Store(true) + return 0, context.Canceled +} + +func (r *lateWriteReader) cancel() { + r.cancelOnce.Do(func() { close(r.cancelled) }) +} + +func (r *lateWriteReader) Close() error { + r.closed.Store(true) + return nil +} + +type timeoutReaderClientCase struct { + name string + http2 bool + client func(t *testing.T, server *httptest.Server) *gohttp.Client +} + +func timeoutReaderClientCases() []timeoutReaderClientCase { + newClient := func(mode httpclient.ImpersonationMode) func(t *testing.T, server *httptest.Server) *gohttp.Client { + return func(t *testing.T, server *httptest.Server) *gohttp.Client { + var tlsConfig *tls.Config + if server.TLS != nil { + roots := x509.NewCertPool() + roots.AddCert(server.Certificate()) + tlsConfig = &tls.Config{RootCAs: roots} + } + client, err := httpclient.NewClient(httpclient.Options{ + Transport: httpclient.TransportOptions{ + TLSClientConfig: tlsConfig, + }, + Impersonation: httpclient.ImpersonationOptions{ + Mode: mode, + Browser: httpclient.BrowserChrome, + }, + }) + if err != nil { + t.Fatal(err) + } + return client + } + } + var cases []timeoutReaderClientCase + for _, mode := range []httpclient.ImpersonationMode{httpclient.ImpersonationNative, httpclient.ImpersonationFixed} { + cases = append(cases, + timeoutReaderClientCase{name: string(mode) + "/http1", client: newClient(mode)}, + timeoutReaderClientCase{name: string(mode) + "/http2", http2: true, client: newClient(mode)}, + ) + } + return cases +} + +func startTimeoutReaderServer(t *testing.T, useHTTP2 bool, handler gohttp.HandlerFunc) *httptest.Server { + server := httptest.NewUnstartedServer(handler) + if useHTTP2 { + server.EnableHTTP2 = true + server.StartTLS() + } else { + server.Start() + } + t.Cleanup(server.Close) + return server +} + +// readBodyWithTimeout sends the request under its own context derived from +// ctx, as the fetcher does, and reads the body through a TimeoutReader. +func readBodyWithTimeout(t *testing.T, ctx context.Context, client *gohttp.Client, url string, timeout time.Duration, wantProto int) ([]byte, error) { + t.Helper() + reqCtx, cancel := context.WithCancel(ctx) + defer cancel() + req, err := gohttp.NewRequestWithContext(reqCtx, gohttp.MethodGet, url, nil) + if err != nil { + t.Fatal(err) + } + resp, err := client.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if resp.ProtoMajor != wantProto { + t.Fatalf("expected HTTP/%d, got %s", wantProto, resp.Proto) + } + + type result struct { + data []byte + err error + } + done := make(chan result, 1) + go func() { + data, err := io.ReadAll(NewTimeoutReader(resp.Body, timeout, cancel)) + done <- result{data, err} + }() + select { + case r := <-done: + return r.data, r.err + case <-time.After(timeout + 10*time.Second): + t.Fatal("Read did not return after the read timeout") + return nil, nil + } +} + +func TestTimeoutReader_StalledHTTPBodyTimesOut(t *testing.T) { + for _, tc := range timeoutReaderClientCases() { + t.Run(tc.name, func(t *testing.T) { + server := startTimeoutReaderServer(t, tc.http2, func(w gohttp.ResponseWriter, r *gohttp.Request) { + w.Header().Set("Content-Length", "1048576") + w.Write(bytes.Repeat([]byte("a"), 1024)) + w.(gohttp.Flusher).Flush() + <-r.Context().Done() + }) + + wantProto := 1 + if tc.http2 { + wantProto = 2 + } + const timeout = 200 * time.Millisecond + start := time.Now() + data, err := readBodyWithTimeout(t, context.Background(), tc.client(t, server), server.URL, timeout, wantProto) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("expected %v, got %v", context.DeadlineExceeded, err) + } + if len(data) != 1024 { + t.Fatalf("expected the 1024 bytes sent before the stall, got %d", len(data)) + } + if elapsed := time.Since(start); elapsed < timeout { + t.Fatalf("timed out after %v, before the %v read timeout", elapsed, timeout) + } + }) + } +} + +func TestTimeoutReader_SlowSteadyHTTPBodyDoesNotTimeOut(t *testing.T) { + const ( + timeout = 300 * time.Millisecond + interval = 30 * time.Millisecond + chunks = 40 // 1.2s in total, four times the read timeout + ) + for _, tc := range timeoutReaderClientCases() { + t.Run(tc.name, func(t *testing.T) { + server := startTimeoutReaderServer(t, tc.http2, func(w gohttp.ResponseWriter, r *gohttp.Request) { + for i := 0; i < chunks; i++ { + w.Write([]byte("0123456789")) + w.(gohttp.Flusher).Flush() + time.Sleep(interval) + } + }) + + wantProto := 1 + if tc.http2 { + wantProto = 2 + } + data, err := readBodyWithTimeout(t, context.Background(), tc.client(t, server), server.URL, timeout, wantProto) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + if len(data) != chunks*10 { + t.Fatalf("expected %d bytes, got %d", chunks*10, len(data)) + } + }) + } +} + +// A gzip-encoded body that stalls must time out without racing the caller's +// Body.Close against the decompressing reader. +func TestTimeoutReader_StalledGzipHTTPBodyTimesOut(t *testing.T) { + payload := bytes.Repeat([]byte("gopeed"), 1024) + stages := []struct { + name string + write func(w io.Writer) + want int + }{ + { + // The first bytes of a gzip header, so the client blocks inside + // gzip.NewReader. + name: "incomplete-header", + write: func(w io.Writer) { w.Write([]byte{0x1f, 0x8b, 0x08, 0x00}) }, + want: 0, + }, + { + // A complete header and some compressed data, flushed without + // finishing the gzip stream. + name: "stalled-body", + write: func(w io.Writer) { + gz := gzip.NewWriter(w) + gz.Write(payload) + gz.Flush() + }, + want: len(payload), + }, + } + for _, stage := range stages { + for _, tc := range timeoutReaderClientCases() { + t.Run(stage.name+"/"+tc.name, func(t *testing.T) { + server := startTimeoutReaderServer(t, tc.http2, func(w gohttp.ResponseWriter, r *gohttp.Request) { + w.Header().Set("Content-Encoding", "gzip") + stage.write(w) + w.(gohttp.Flusher).Flush() + <-r.Context().Done() + }) + + wantProto := 1 + if tc.http2 { + wantProto = 2 + } + data, err := readBodyWithTimeout(t, context.Background(), tc.client(t, server), server.URL, 200*time.Millisecond, wantProto) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("expected %v, got %v", context.DeadlineExceeded, err) + } + if len(data) != stage.want { + t.Fatalf("expected the %d bytes sent before the stall, got %d", stage.want, len(data)) + } + }) + } + } +} + +// Cancelling the caller's context must not be reported as a read timeout. +func TestTimeoutReader_CallerCancellationIsNotTimeout(t *testing.T) { + for _, tc := range timeoutReaderClientCases() { + t.Run(tc.name, func(t *testing.T) { + server := startTimeoutReaderServer(t, tc.http2, func(w gohttp.ResponseWriter, r *gohttp.Request) { + w.Header().Set("Content-Length", "1048576") + w.Write(bytes.Repeat([]byte("a"), 1024)) + w.(gohttp.Flusher).Flush() + <-r.Context().Done() + }) + + wantProto := 1 + if tc.http2 { + wantProto = 2 + } + // Cancel once the client has read the first bytes, so the next + // Read is the one blocked when the caller cancels. + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + client := tc.client(t, server) + next := client.Transport + client.Transport = roundTripFunc(func(r *gohttp.Request) (*gohttp.Response, error) { + resp, err := next.RoundTrip(r) + if err == nil { + resp.Body = &cancelAfterFirstRead{ReadCloser: resp.Body, cancel: cancel} + } + return resp, err + }) + + const timeout = 5 * time.Second + start := time.Now() + data, err := readBodyWithTimeout(t, ctx, client, server.URL, timeout, wantProto) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected %v, got %v", context.Canceled, err) + } + if len(data) == 0 { + t.Fatal("cancelled before any data was read") + } + if elapsed := time.Since(start); elapsed >= timeout { + t.Fatalf("read returned after %v, not on cancellation", elapsed) + } + }) + } +} + +type roundTripFunc func(*gohttp.Request) (*gohttp.Response, error) + +func (f roundTripFunc) RoundTrip(r *gohttp.Request) (*gohttp.Response, error) { + return f(r) +} + +// cancelAfterFirstRead calls cancel once a Read has returned data. +type cancelAfterFirstRead struct { + io.ReadCloser + once sync.Once + cancel context.CancelFunc +} + +func (b *cancelAfterFirstRead) Read(p []byte) (int, error) { + n, err := b.ReadCloser.Read(p) + if n > 0 { + b.once.Do(b.cancel) + } + return n, err +}