fix: Bound plugin downloads with transport timeouts

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 {

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,
+		},
+	}
+}

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")
+}

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 {

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},