diff --git a/managedplugin/download.go b/managedplugin/download.go index f5eb9e7..5bc1bf4 100644 --- a/managedplugin/download.go +++ b/managedplugin/download.go @@ -93,8 +93,11 @@ func getURLLocation(ctx context.Context, org string, name string, version string if err != nil { return fmt.Errorf("failed create request %s: %w", urlForLog, redactURLError(err)) } - resp, err := http.DefaultClient.Do(req) + resp, err := downloadClient.Do(req) if err != nil { + if terr := downloadTimeoutError(ctx, urlForLog, err); terr != nil { + return terr + } return fmt.Errorf("failed to get url %s: %w", urlForLog, redactURLError(err)) } resp.Body.Close() @@ -370,8 +373,11 @@ func downloadFile(ctx context.Context, localPath string, downloadURL string, dop } // Do http request - resp, err := http.DefaultClient.Do(req) + resp, err := downloadClient.Do(req) if err != nil { + if terr := downloadTimeoutError(ctx, urlForLog, err); terr != nil { + return terr + } return fmt.Errorf("failed to get url %s: %w", urlForLog, redactURLError(err)) } defer resp.Body.Close() @@ -400,6 +406,9 @@ func downloadFile(ctx context.Context, localPath string, downloadURL string, dop // Write the body to file written, err := io.Copy(io.MultiWriter(writers...), resp.Body) if err != nil { + if terr := downloadTimeoutError(ctx, urlForLog, err); terr != nil { + return terr + } return fmt.Errorf("failed to copy body to file %s: %w", out.Name(), err) } if resp.ContentLength >= 0 && written != resp.ContentLength { diff --git a/managedplugin/download_client.go b/managedplugin/download_client.go new file mode 100644 index 0000000..c071934 --- /dev/null +++ b/managedplugin/download_client.go @@ -0,0 +1,53 @@ +package managedplugin + +import ( + "context" + "net" + "net/http" + "time" +) + +var ( + downloadDialTimeout = 10 * time.Second + downloadTLSHandshakeTimeout = 10 * time.Second + downloadResponseHeaderTimeout = 30 * time.Second + downloadIdleReadTimeout = 30 * time.Second +) + +var downloadClient = newDownloadClient() + +type idleTimeoutConn struct { + net.Conn + idle time.Duration +} + +func (c *idleTimeoutConn) Read(b []byte) (int, error) { + if err := c.SetReadDeadline(time.Now().Add(c.idle)); err != nil { + return 0, err + } + return c.Conn.Read(b) +} + +func newDownloadClient() *http.Client { + dialer := &net.Dialer{ + Timeout: downloadDialTimeout, + KeepAlive: 30 * time.Second, + } + idle := downloadIdleReadTimeout + return &http.Client{ + Transport: &http.Transport{ + Proxy: http.ProxyFromEnvironment, + DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) { + conn, err := dialer.DialContext(ctx, network, addr) + if err != nil { + return nil, err + } + return &idleTimeoutConn{Conn: conn, idle: idle}, nil + }, + TLSHandshakeTimeout: downloadTLSHandshakeTimeout, + ResponseHeaderTimeout: downloadResponseHeaderTimeout, + IdleConnTimeout: downloadIdleReadTimeout, + MaxIdleConnsPerHost: http.DefaultMaxIdleConnsPerHost, + }, + } +} diff --git a/managedplugin/download_client_test.go b/managedplugin/download_client_test.go new file mode 100644 index 0000000..950057f --- /dev/null +++ b/managedplugin/download_client_test.go @@ -0,0 +1,165 @@ +package managedplugin + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func fastStall(t *testing.T, d time.Duration) { + t.Helper() + + prevHeader, prevIdle, prevClient := downloadResponseHeaderTimeout, downloadIdleReadTimeout, downloadClient + downloadResponseHeaderTimeout, downloadIdleReadTimeout = d, d + downloadClient = newDownloadClient() + t.Cleanup(func() { + downloadResponseHeaderTimeout, downloadIdleReadTimeout = prevHeader, prevIdle + downloadClient = prevClient + }) +} + +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.GreaterOrEqual(t, attempts.Load(), int32(2)) + require.Equal(t, sha256Hex(body), checksum) +} + +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.GreaterOrEqual(t, attempts.Load(), int32(2)) + + written, err := os.ReadFile(localPath) + require.NoError(t, err) + require.Equal(t, body, written) + require.Equal(t, sha256Hex(body), checksum) +} + +func TestDownloadFileSlowButProgressingSucceeds(t *testing.T) { + fastRetries(t) + fastStall(t, 300*time.Millisecond) + + body := []byte("slow-but-moving-plugin-payload") + + var attempts atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + attempts.Add(1) + w.Header().Set("Content-Length", fmt.Sprintf("%d", len(body))) + w.WriteHeader(http.StatusOK) + for _, chunk := range [][]byte{body[:6], body[6:12], body[12:18], body[18:24], body[24:]} { + _, _ = w.Write(chunk) + w.(http.Flusher).Flush() + time.Sleep(50 * time.Millisecond) + } + })) + t.Cleanup(server.Close) + + 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, 1, attempts.Load(), "a transfer that keeps making progress must not trip the idle deadline") + require.Equal(t, sha256Hex(body), checksum) +} + +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.ErrorIs(t, err, context.Canceled) + require.False(t, IsTransientDownloadError(err), "the caller's own cancellation must not be retried") +} + +func TestDownloadFileCallerDeadlineIsNotRetried(t *testing.T) { + fastRetries(t) + fastStall(t, 10*time.Second) + + release := make(chan struct{}) + + var attempts atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) { + attempts.Add(1) + <-release + })) + t.Cleanup(server.Close) + t.Cleanup(func() { close(release) }) + + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + + localPath := filepath.Join(t.TempDir(), "plugin.zip") + _, err := downloadFile(ctx, localPath, server.URL, DownloaderOptions{NoProgress: true}) + require.Error(t, err) + require.NotErrorIs(t, err, errDownloadStalled) + require.False(t, IsTransientDownloadError(err), "the caller's own deadline must not be retried") +} diff --git a/managedplugin/download_retry.go b/managedplugin/download_retry.go index 056186f..849cead 100644 --- a/managedplugin/download_retry.go +++ b/managedplugin/download_retry.go @@ -13,8 +13,9 @@ import ( ) 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. @@ -81,7 +82,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), @@ -114,6 +116,17 @@ func isRetryableDownloadError(err error) bool { return false } +func downloadTimeoutError(ctx context.Context, urlForLog string, err error) error { + if ctx.Err() != nil { + return nil + } + var netErr net.Error + if !errors.As(err, &netErr) || !netErr.Timeout() { + return nil + } + return fmt.Errorf("%w: no data from %s: %v", errDownloadStalled, urlForLog, redactURLError(err)) +} + // redactURLQuery strips the query string so that the signed download token never // reaches stdout or a log aggregator. func redactURLQuery(rawURL string) string { diff --git a/managedplugin/download_retry_test.go b/managedplugin/download_retry_test.go index ba62afd..ddc19dc 100644 --- a/managedplugin/download_retry_test.go +++ b/managedplugin/download_retry_test.go @@ -58,6 +58,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: i/o timeout", 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},