diff --git a/managedplugin/download.go b/managedplugin/download.go index f5eb9e7..fb75226 100644 --- a/managedplugin/download.go +++ b/managedplugin/download.go @@ -12,6 +12,7 @@ import ( "path/filepath" "runtime" "strings" + "sync/atomic" "time" "github.com/avast/retry-go/v5" @@ -363,8 +364,28 @@ func downloadFile(ctx context.Context, localPath string, downloadURL string, dop return err } + // A server that stops sending — headers or body — otherwise holds the + // attempt until some middlebox kills the connection minutes later. The + // watchdog cancels the attempt after downloadStallTimeout without + // progress, and stallErr turns that cancellation into a retryable error + // instead of the terminal context.Canceled the caller's own cancel gets. + attemptCtx, cancel := context.WithCancel(ctx) + defer cancel() + var stalled atomic.Bool + watchdog := time.AfterFunc(downloadStallTimeout, func() { + stalled.Store(true) + cancel() + }) + defer watchdog.Stop() + stallErr := func(err error) error { + if !stalled.Load() { + return nil + } + return fmt.Errorf("%w: no data from %s for %s: %v", errDownloadStalled, urlForLog, downloadStallTimeout, redactURLError(err)) + } + // Get the data - req, err := http.NewRequestWithContext(ctx, http.MethodGet, downloadURL, nil) + req, err := http.NewRequestWithContext(attemptCtx, http.MethodGet, downloadURL, nil) if err != nil { return fmt.Errorf("failed create request %s: %w", urlForLog, redactURLError(err)) } @@ -372,9 +393,13 @@ func downloadFile(ctx context.Context, localPath string, downloadURL string, dop // Do http request resp, err := http.DefaultClient.Do(req) if err != nil { + if serr := stallErr(err); serr != nil { + return serr + } return fmt.Errorf("failed to get url %s: %w", urlForLog, redactURLError(err)) } defer resp.Body.Close() + watchdog.Reset(downloadStallTimeout) // Check server response if resp.StatusCode == http.StatusNotFound { return errNotFound @@ -398,8 +423,11 @@ func downloadFile(ctx context.Context, localPath string, downloadURL string, dop } // Write the body to file - written, err := io.Copy(io.MultiWriter(writers...), resp.Body) + written, err := io.Copy(io.MultiWriter(writers...), &stallResetReader{r: resp.Body, watchdog: watchdog, timeout: downloadStallTimeout}) if err != nil { + if serr := stallErr(err); serr != nil { + return serr + } return fmt.Errorf("failed to copy body to file %s: %w", out.Name(), err) } if resp.ContentLength >= 0 && written != resp.ContentLength { @@ -417,6 +445,23 @@ func downloadFile(ctx context.Context, localPath string, downloadURL string, dop return checksum, nil } +// stallResetReader defers the stall watchdog every time bytes arrive, so it +// only fires when the peer stops sending entirely — a slow but moving +// download never trips it. +type stallResetReader struct { + r io.Reader + watchdog *time.Timer + timeout time.Duration +} + +func (s *stallResetReader) Read(p []byte) (n int, err error) { + n, err = s.r.Read(p) + if n > 0 { + s.watchdog.Reset(s.timeout) + } + return n, err +} + func truncateFile(f *os.File) error { if err := f.Truncate(0); err != nil { return fmt.Errorf("failed to truncate file %s: %w", f.Name(), err) diff --git a/managedplugin/download_retry.go b/managedplugin/download_retry.go index 056186f..889091c 100644 --- a/managedplugin/download_retry.go +++ b/managedplugin/download_retry.go @@ -10,11 +10,13 @@ import ( "net/url" "strings" "syscall" + "time" ) var ( - errNotFound = errors.New("not found") - errShortRead = errors.New("truncated response body") + errNotFound = errors.New("not found") + errShortRead = errors.New("truncated response body") + errDownloadStalled = errors.New("download stalled") ) // Overridable so tests do not pay the real backoff. @@ -22,6 +24,7 @@ var ( downloadRetryAttempts = uint(RetryAttempts) downloadRetryDelay = RetryWaitTime downloadRetryMaxDelay = MaxRetryWaitTime + downloadStallTimeout = 30 * time.Second ) type httpStatusError struct { @@ -81,7 +84,8 @@ func isRetryableDownloadError(err error) bool { } switch { - case errors.Is(err, errShortRead), + case errors.Is(err, errDownloadStalled), + errors.Is(err, errShortRead), errors.Is(err, io.ErrUnexpectedEOF), errors.Is(err, io.EOF), errors.Is(err, syscall.ECONNRESET), diff --git a/managedplugin/download_retry_test.go b/managedplugin/download_retry_test.go index ba62afd..110861d 100644 --- a/managedplugin/download_retry_test.go +++ b/managedplugin/download_retry_test.go @@ -11,6 +11,7 @@ import ( "net/http/httptest" "os" "path/filepath" + "sync/atomic" "syscall" "testing" "time" @@ -58,6 +59,7 @@ func TestIsRetryableDownloadError(t *testing.T) { {name: "bad gateway 502", err: &httpStatusError{statusCode: http.StatusBadGateway}, want: true}, {name: "service unavailable 503", err: &httpStatusError{statusCode: http.StatusServiceUnavailable}, want: true}, + {name: "stalled download", err: fmt.Errorf("%w: no data from host for 30s", errDownloadStalled), want: true}, {name: "context canceled", err: fmt.Errorf("get url: %w", context.Canceled), want: false}, {name: "context deadline exceeded", err: fmt.Errorf("get url: %w", context.DeadlineExceeded), want: false}, {name: "checksum mismatch is permanent", err: errors.New("checksum mismatch: expected abc, got def"), want: false}, @@ -208,6 +210,115 @@ func fastRetries(t *testing.T) { }) } +// fastStall shrinks the stall watchdog so the stall tests do not wait 30s per +// attempt. +func fastStall(t *testing.T, d time.Duration) { + t.Helper() + + prev := downloadStallTimeout + downloadStallTimeout = d + t.Cleanup(func() { + downloadStallTimeout = prev + }) +} + +// TestDownloadFileRetriesStalledBody reproduces the assets.cloudquery.io outage +// mode: the server sends part of the body and then goes silent without closing +// the connection. Without the watchdog the attempt hangs until a middlebox +// kills the stream minutes later; with it the attempt dies quickly and the +// retry succeeds. +func TestDownloadFileRetriesStalledBody(t *testing.T) { + fastRetries(t) + fastStall(t, 150*time.Millisecond) + + body := []byte("cloudquery-plugin-binary-payload") + release := make(chan struct{}) + + var attempts atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + if attempts.Add(1) == 1 { + w.Header().Set("Content-Length", fmt.Sprintf("%d", len(body))) + w.WriteHeader(http.StatusOK) + _, _ = w.Write(body[:5]) + w.(http.Flusher).Flush() + <-release + return + } + _, _ = w.Write(body) + })) + t.Cleanup(server.Close) + t.Cleanup(func() { close(release) }) + + localPath := filepath.Join(t.TempDir(), "plugin.zip") + checksum, err := downloadFile(context.Background(), localPath, server.URL, DownloaderOptions{NoProgress: true}) + require.NoError(t, err) + require.EqualValues(t, 2, attempts.Load()) + + written, err := os.ReadFile(localPath) + require.NoError(t, err) + require.Equal(t, body, written) + require.Equal(t, sha256Hex(body), checksum) +} + +// TestDownloadFileRetriesStalledHeaders covers a stall before any response +// arrives — the watchdog must cut the header wait too, not just the body copy. +func TestDownloadFileRetriesStalledHeaders(t *testing.T) { + fastRetries(t) + fastStall(t, 150*time.Millisecond) + + body := []byte("payload") + release := make(chan struct{}) + + var attempts atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + if attempts.Add(1) == 1 { + <-release + return + } + _, _ = w.Write(body) + })) + t.Cleanup(server.Close) + t.Cleanup(func() { close(release) }) + + localPath := filepath.Join(t.TempDir(), "plugin.zip") + checksum, err := downloadFile(context.Background(), localPath, server.URL, DownloaderOptions{NoProgress: true}) + require.NoError(t, err) + require.EqualValues(t, 2, attempts.Load()) + require.Equal(t, sha256Hex(body), checksum) +} + +// TestDownloadFileCallerCancelIsNotRetried pins the boundary between the +// watchdog's own cancellation (retryable) and the caller's (terminal). +func TestDownloadFileCallerCancelIsNotRetried(t *testing.T) { + fastRetries(t) + + release := make(chan struct{}) + + var attempts atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + attempts.Add(1) + w.Header().Set("Content-Length", "32") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("early")) + w.(http.Flusher).Flush() + <-release + })) + t.Cleanup(server.Close) + t.Cleanup(func() { close(release) }) + + ctx, cancel := context.WithCancel(context.Background()) + go func() { + time.Sleep(100 * time.Millisecond) + cancel() + }() + + localPath := filepath.Join(t.TempDir(), "plugin.zip") + _, err := downloadFile(ctx, localPath, server.URL, DownloaderOptions{NoProgress: true}) + require.Error(t, err) + require.EqualValues(t, 1, attempts.Load(), "the caller's own cancellation must not be retried") + require.NotErrorIs(t, err, errDownloadStalled) +} + // TestDownloadFileRedactsSignedTokenFromTransportErrors covers the leak that // url.Error reopens: it prints its URL verbatim, so wrapping one puts the signed // token back into the message once per attempt.