refactor(internal/fetch): reorganize download logic (#3022)

Refactor the tarball download code to split download logic into three
functions:

- DownloadTarball: public API that verifies SHA256 against expected
value
- downloadTarball: handles retries with exponential backoff, returns
SHA256
- downloadAttempt: performs a single download attempt, returns SHA256

This allows us to reuse these functions more easily when implementing
caching.

Author: julieqiu

internal/fetch/fetch.go

@@ -19,6 +19,7 @@ import (
 	"archive/tar"
 	"compress/gzip"
 	"crypto/sha256"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
@@ -28,6 +29,8 @@ import (
 	"time"
 )
 
+var errChecksumMismatch = errors.New("checksum mismatch")
+
 // Endpoints defines the endpoints used to access GitHub.
 type Endpoints struct {
 	// API defines the endpoint used to make API calls.
@@ -113,48 +116,78 @@ func TarballLink(githubDownload string, repo *Repo, sha string) string {
 	return fmt.Sprintf("%s/%s/%s/archive/%s.tar.gz", githubDownload, repo.Org, repo.Repo, sha)
 }
 
-// DownloadTarball downloads a tarball from the given source URL to the target path,
-// verifying its SHA256 checksum matches expectedSha256. It retries up to 3 times
-// with exponential backoff on failure.
-func DownloadTarball(target, source, expectedSha256 string) error {
+// DownloadTarball downloads a tarball from the given url to the target
+// path, verifying its SHA256 checksum matches expectedSha256. It retries up to
+// 3 times with exponential backoff on failure.
+func DownloadTarball(target, url, expectedSha256 string) error {
 	if fileExists(target) {
 		return nil
 	}
+	if err := os.MkdirAll(filepath.Dir(target), 0755); err != nil {
+		return err
+	}
+	tempFile, err := os.CreateTemp(filepath.Dir(target), "temp-")
+	if err != nil {
+		return err
+	}
+	tempPath := tempFile.Name()
+	defer func() {
+		tempFile.Close()
+		cerr := os.Remove(tempPath)
+		if err == nil && cerr != nil && !os.IsNotExist(cerr) {
+			err = cerr
+		}
+	}()
+
+	if err := downloadTarball(tempPath, url); err != nil {
+		return err
+	}
+
+	sha, err := computeSHA256(tempPath)
+	if err != nil {
+		return err
+	}
+	if sha != expectedSha256 {
+		return fmt.Errorf("%w: expected=%s, got=%s", errChecksumMismatch, expectedSha256, sha)
+	}
+	if err := os.MkdirAll(filepath.Dir(target), 0755); err != nil {
+		return err
+	}
+	return os.Rename(tempPath, target)
+}
+
+// downloadTarball downloads a tarball from the given source URL to the target
+// path. It retries up to 3 times with exponential backoff on failure.
+func downloadTarball(target, source string) error {
 	var err error
 	backoff := 10 * time.Second
 	for i := range 3 {
 		if i != 0 {
 			time.Sleep(backoff)
 			backoff = 2 * backoff
 		}
-		if err = downloadAttempt(target, source, expectedSha256); err == nil {
+		if err = downloadAttempt(target, source); err == nil {
 			return nil
 		}
 	}
 	return fmt.Errorf("download failed after 3 attempts, last error=%w", err)
 }
 
-func downloadAttempt(target, source, expectedSha256 string) (err error) {
-	if err := os.MkdirAll(filepath.Dir(target), 0755); err != nil {
-		return err
-	}
-	tempFile, err := os.CreateTemp(filepath.Dir(target), "temp-")
+func downloadAttempt(target, source string) (err error) {
+	file, err := os.Create(target)
 	if err != nil {
 		return err
 	}
 	defer func() {
-		cerr := tempFile.Close()
+		cerr := file.Close()
 		if err == nil {
 			err = cerr
 		}
 		if err != nil {
-			os.Remove(tempFile.Name())
+			os.Remove(target)
 		}
 	}()
 
-	hasher := sha256.New()
-	writer := io.MultiWriter(tempFile, hasher)
-
 	client := http.Client{Timeout: 60 * time.Second}
 	response, err := client.Get(source)
 	if err != nil {
@@ -165,15 +198,11 @@ func downloadAttempt(target, source, expectedSha256 string) (err error) {
 		return fmt.Errorf("http error in download %s", response.Status)
 	}
 
-	if _, err := io.Copy(writer, response.Body); err != nil {
+	if _, err := io.Copy(file, response.Body); err != nil {
 		return err
 	}
 
-	got := fmt.Sprintf("%x", hasher.Sum(nil))
-	if expectedSha256 != got {
-		return fmt.Errorf("mismatched hash on download, expected=%s, got=%s", expectedSha256, got)
-	}
-	return os.Rename(tempFile.Name(), target)
+	return nil
 }
 
 func fileExists(name string) bool {

internal/fetch/fetch_test.go

@@ -19,6 +19,7 @@ import (
 	"bytes"
 	"compress/gzip"
 	"crypto/sha256"
+	"errors"
 	"fmt"
 	"net/http"
 	"net/http/httptest"
@@ -201,24 +202,19 @@ func TestTarballLink(t *testing.T) {
 
 func TestDownloadTarballTgzExists(t *testing.T) {
 	testDir := t.TempDir()
-
 	tarball := makeTestContents(t)
-
 	target := path.Join(testDir, "existing-file")
 	if err := os.WriteFile(target, tarball.Contents, 0644); err != nil {
 		t.Fatal(err)
 	}
-
 	if err := DownloadTarball(target, "https://unused/placeholder.tar.gz", tarball.Sha256); err != nil {
 		t.Fatal(err)
 	}
 }
 
 func TestDownloadTarballNeedsDownload(t *testing.T) {
 	testDir := t.TempDir()
-
 	tarball := makeTestContents(t)
-
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if r.URL.Path != "/placeholder.tar.gz" {
 			t.Errorf("Expected to request '/placeholder.tar.gz', got: %s", r.URL.Path)
@@ -241,6 +237,27 @@ func TestDownloadTarballNeedsDownload(t *testing.T) {
 	}
 }
 
+func TestDownloadTarballChecksumMismatch(t *testing.T) {
+	testDir := t.TempDir()
+	tarball := makeTestContents(t)
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		w.Write(tarball.Contents)
+	}))
+	defer server.Close()
+
+	target := path.Join(testDir, "target-file")
+	wrongSha := "0000000000000000000000000000000000000000000000000000000000000000"
+
+	err := DownloadTarball(target, server.URL+"/test.tar.gz", wrongSha)
+	if !errors.Is(err, errChecksumMismatch) {
+		t.Fatalf("expected errChecksumMismatch, got: %v", err)
+	}
+	if _, err := os.Stat(target); !os.IsNotExist(err) {
+		t.Errorf("target file should not exist after checksum failure: %v", err)
+	}
+}
+
 type contents struct {
 	Sha256   string
 	Contents []byte