fix: component download (#1394)

Author: jon-stewart

lwcomponent/component.go

@@ -257,7 +257,14 @@ func (s State) Install(component *Component, version string) error {
 		return err
 	}
 
-	err = DownloadFile(path, artifact.URL, 5*time.Minute)
+	// Slow S3 downloads
+	downloadTimeout := 0 * time.Minute
+	switch component.Name {
+	case "sast":
+		downloadTimeout = 5 * time.Minute
+	}
+
+	err = DownloadFile(path, artifact.URL, time.Duration(downloadTimeout))
 	if err != nil {
 		return errors.Wrap(err, "unable to download component artifact")
 	}

lwcomponent/http.go

@@ -30,37 +30,47 @@ const (
 	defaultMaxRetry = 2
 )
 
-// downloadFile is an internal helper that downloads a file to the provided file path
-func DownloadFile(filepath string, url string, timeout time.Duration) error {
-	var (
-		resp     *http.Response
-		err      error
-		_timeout time.Duration = timeout
-	)
-
-	if _timeout == 0 {
-		_timeout = defaultTimeout
+func DownloadFile(filepath string, url string, timeout time.Duration) (err error) {
+	if timeout == 0 {
+		timeout = defaultTimeout
 	}
 
-	client := &http.Client{Timeout: _timeout}
+	client := &http.Client{Timeout: timeout}
 
-	resp, err = client.Get(url)
+	file, err := os.Create(filepath)
 	if err != nil {
-		for retry := 0; retry < defaultMaxRetry && os.IsTimeout(err); retry++ {
-			resp, err = client.Get(url)
-		}
+		return err
+	}
+	defer file.Close()
+
+	err = downloadFile(client, url, file)
+
+	for retry := 0; retry < defaultMaxRetry && os.IsTimeout(err); retry++ {
+		_, err = file.Seek(0, io.SeekStart)
 		if err != nil {
 			return err
 		}
+
+		err = downloadFile(client, url, file)
 	}
-	defer resp.Body.Close()
 
-	out, err := os.Create(filepath)
+	return
+}
+
+// client.Get returns on receiving HTTP headers and we stream the HTTP data to the output file.
+// A timeout will interrupt io.Copy.
+func downloadFile(client *http.Client, url string, file *os.File) (err error) {
+	var (
+		resp *http.Response
+	)
+
+	resp, err = client.Get(url)
 	if err != nil {
 		return err
 	}
-	defer out.Close()
+	defer resp.Body.Close()
+
+	_, err = io.Copy(file, resp.Body)
 
-	_, err = io.Copy(out, resp.Body)
-	return err
+	return
 }

lwcomponent/http_test.go

@@ -43,9 +43,22 @@ func TestDownloadFile(t *testing.T) {
 	})
 
 	t.Run("timeout error", func(t *testing.T) {
-		err = lwcomponent.DownloadFile(file.Name(), fmt.Sprintf("%s%s", server.URL, urlPath), 1*time.Microsecond)
+		var (
+			count int = 0
+		)
+
+		mux.HandleFunc("/slow", func(w http.ResponseWriter, r *http.Request) {
+			if assert.Equal(t, "GET", r.Method, "Get() should be a GET method") {
+				count += 1
+				time.Sleep(10 * time.Millisecond)
+				fmt.Fprint(w, content)
+			}
+		})
+
+		err = lwcomponent.DownloadFile(file.Name(), fmt.Sprintf("%s%s", server.URL, "/slow"), 2*time.Millisecond)
 		assert.NotNil(t, err)
 		assert.True(t, os.IsTimeout(err))
+		assert.Equal(t, 3, count)
 	})
 
 	t.Run("non-timeout error", func(t *testing.T) {