refactor download functions

Author: lucarin91

internal/updater/download_image.go

@@ -16,17 +16,14 @@
 package updater
 
 import (
-	"bytes"
 	"context"
-	"crypto/sha256"
-	"encoding/hex"
 	"fmt"
-	"io"
 	"os"
 
 	"github.com/arduino/go-paths-helper"
 	"github.com/codeclysm/extract/v4"
 	"github.com/schollz/progressbar/v3"
+	"go.bug.st/downloader/v2"
 
 	"github.com/arduino/arduino-flasher-cli/cmd/feedback"
 	"github.com/arduino/arduino-flasher-cli/cmd/i18n"
@@ -65,58 +62,25 @@ func DownloadAndExtract(ctx context.Context, targetVersion string, temp *paths.P
 }
 
 func DownloadImage(ctx context.Context, targetVersion string, downloadPath *paths.Path) (*paths.Path, string, error) {
-	var err error
-
 	client := NewClient()
-	manifest, err := client.GetInfoManifest(ctx)
-	if err != nil {
-		return nil, "", err
-	}
-
-	var rel *Release
-	if targetVersion == "latest" || targetVersion == manifest.Latest.Version {
-		rel = &manifest.Latest
-	} else {
-		for _, r := range manifest.Releases {
-			if targetVersion == r.Version {
-				rel = &r
-				break
-			}
-		}
-	}
-
-	if rel == nil {
-		return nil, "", fmt.Errorf("could not find Debian image %s", targetVersion)
-	}
-
-	download, size, err := client.FetchZip(ctx, rel.Url)
+	rel, err := client.GetReleaseByVersion(ctx, targetVersion)
 	if err != nil {
-		return nil, "", fmt.Errorf("could not fetch Debian image: %w", err)
+		return nil, "", fmt.Errorf("could not get release info: %w", err)
 	}
-	defer download.Close()
 
 	tmpZip := downloadPath.Join("arduino-unoq-debian-image-" + rel.Version + ".tar.zst")
-	tmpZipFile, err := tmpZip.Create()
-	if err != nil {
-		return nil, "", err
-	}
-	defer tmpZipFile.Close()
 
-	// Download and keep track of the progress
 	bar := progressbar.DefaultBytes(
-		size,
+		0,
 		i18n.Tr("Downloading Debian image version %s", rel.Version),
 	)
-	checksum := sha256.New()
-	if _, err := io.Copy(io.MultiWriter(checksum, tmpZipFile, bar), download); err != nil {
-		return nil, "", err
+	callback := func(current, total int64) {
+		bar.AddMax64(total)
+		_ = bar.Set64(current)
 	}
 
-	// Check the hash
-	if sha256Byte, err := hex.DecodeString(rel.Sha256); err != nil {
-		return nil, "", fmt.Errorf("could not convert sha256 from hex to bytes: %w", err)
-	} else if s := checksum.Sum(nil); !bytes.Equal(s, sha256Byte) {
-		return nil, "", fmt.Errorf("bad hash: %x (expected %x)", s, sha256Byte)
+	if err := client.DownloadFile(ctx, tmpZip, rel, callback, downloader.Config{}); err != nil {
+		return nil, "", fmt.Errorf("could not download Debian image: %w", err)
 	}
 
 	return tmpZip, rel.Version, nil

internal/updater/http_client.go

@@ -28,11 +28,11 @@ import (
 	"time"
 
 	"github.com/arduino/go-paths-helper"
+	"github.com/shirou/gopsutil/v4/disk"
 	"go.bug.st/downloader/v2"
 	"go.bug.st/f"
 
 	"github.com/arduino/arduino-flasher-cli/cmd/i18n"
-	"github.com/arduino/arduino-flasher-cli/rpc/cc/arduino/flasher/v1"
 )
 
 var baseURL = f.Must(url.Parse("https://downloads.arduino.cc"))
@@ -41,15 +41,10 @@ const pathRelease = "debian-im/Stable"
 
 // Client holds the base URL, command name, allows custom HTTP client, and optional headers.
 type Client struct {
-	HTTPClient HTTPDoer
+	HTTPClient *http.Client
 	Headers    map[string]string // Optional headers to add to each request
 }
 
-// HTTPDoer is an interface for http.Client or mocks.
-type HTTPDoer interface {
-	Do(req *http.Request) (*http.Response, error)
-}
-
 // Option is a functional option for configuring Client.
 type Option func(*Client)
 
@@ -61,7 +56,7 @@ func WithHeaders(headers map[string]string) Option {
 }
 
 // WithHTTPClient sets a custom HTTP client for the Client.
-func WithHTTPClient(client HTTPDoer) Option {
+func WithHTTPClient(client *http.Client) Option {
 	return func(c *Client) {
 		c.HTTPClient = client
 	}
@@ -116,45 +111,51 @@ func (c *Client) GetInfoManifest(ctx context.Context) (Manifest, error) {
 	return res, nil
 }
 
-// FetchZip fetches the Debian image archive.
-func (c *Client) FetchZip(ctx context.Context, zipURL string) (io.ReadCloser, int64, error) {
-	req, err := http.NewRequestWithContext(ctx, "GET", zipURL, nil)
+func (c *Client) GetReleaseByVersion(ctx context.Context, version string) (Release, error) {
+	manifest, err := c.GetInfoManifest(ctx)
 	if err != nil {
-		return nil, 0, fmt.Errorf("failed to create request: %w", err)
+		return Release{}, err
 	}
-	c.addHeaders(req)
-	// #nosec G107 -- zipURL is constructed from trusted config and parameters
-	resp, err := c.HTTPClient.Do(req)
-	if err != nil {
-		return nil, 0, fmt.Errorf("failed to GET zip: %w", err)
-	}
-	if resp.StatusCode != http.StatusOK {
-		resp.Body.Close()
-		return nil, 0, fmt.Errorf("bad http status from %s: %v", zipURL, resp.Status)
+
+	if version == "latest" || version == manifest.Latest.Version {
+		return manifest.Latest, nil
+	} else {
+		for _, r := range manifest.Releases {
+			if version == r.Version {
+				return r, nil
+			}
+		}
 	}
-	return resp.Body, resp.ContentLength, nil
+
+	return Release{}, fmt.Errorf("could not find Debian image %s", version)
 }
 
+type downloadCallback func(current, total int64)
+
 // DownloadFile downloads a file from a URL into the specified path. An optional config and options may be passed (or nil to use the defaults).
 // A DownloadProgressCB callback function must be passed to monitor download progress.
 // If a not empty queryParameter is passed, it is appended to the URL for analysis purposes.
-func DownloadFile(ctx context.Context, path *paths.Path, rel *Release, downloadCB flasher.DownloadProgressCB, config downloader.Config, options ...downloader.DownloadOptions) (returnedError error) {
-	downloadCB.Start(rel.Url, rel.Version)
-	defer func() {
-		if returnedError == nil {
-			downloadCB.End(true, "")
-		} else {
-			downloadCB.End(false, returnedError.Error())
-		}
-	}()
+func (c *Client) DownloadFile(ctx context.Context, path *paths.Path, rel Release, cb downloadCallback, config downloader.Config, options ...downloader.DownloadOptions) (returnedError error) {
+
+	// Check if there is enough free disk space before downloading and extracting an image
+	dk, err := disk.Usage(path.String())
+	if err != nil {
+		return err
+	}
+	// TODO: improve disk space check with Content-Length header
+	if dk.Free/GiB < DownloadDiskSpace {
+		return fmt.Errorf("download and extraction requires up to %d GiB of free space", DownloadDiskSpace)
+	}
 
+	config.HttpClient = *c.HTTPClient
+	// TODO: add headers to downloader's http client
 	d, err := downloader.DownloadWithConfigAndContext(ctx, path.String(), rel.Url, config, options...)
 	if err != nil {
 		return err
 	}
 
 	err = d.RunAndPoll(func(downloaded int64) {
-		downloadCB.Update(downloaded, d.Size())
+		cb(downloaded, d.Size())
 	}, 250*time.Millisecond)
 	if err != nil {
 		return err

rpc/cc/arduino/flasher/v1/commands.pb.go

@@ -41,6 +41,7 @@ message FlashRequest {
   // The version of the Debian image to download.
   string version = 1;
   // The path in which the image will be downloaded and extracted.
+  // If an image file is already present, it will not be downloaded again.
   string temp_path = 2;
   // Preserve user partition if possible.
   bool preserve_user = 3;

rpc/cc/arduino/flasher/v1/commands.proto

@@ -20,7 +20,6 @@ import (
 
 	"github.com/arduino/go-paths-helper"
 	"github.com/codeclysm/extract/v4"
-	"github.com/shirou/gopsutil/v4/disk"
 	"go.bug.st/downloader/v2"
 
 	"github.com/arduino/arduino-flasher-cli/internal/updater"
@@ -37,7 +36,7 @@ func (s *flasherServerImpl) Flash(req *flasher.FlashRequest, stream flasher.Flas
 		responseCallback = func(*flasher.FlashResponse) error { return nil }
 	}
 	ctx := stream.Context()
-	downloadCB := func(msg *flasher.DownloadProgress) {
+	var downloadCB flasher.DownloadProgressCB = func(msg *flasher.DownloadProgress) {
 		_ = responseCallback(&flasher.FlashResponse{
 			Message: &flasher.FlashResponse_DownloadProgress{
 				DownloadProgress: msg,
@@ -59,43 +58,23 @@ func (s *flasherServerImpl) Flash(req *flasher.FlashRequest, stream flasher.Flas
 		})
 	}
 
-	// Check if there is enough free disk space before downloading and extracting an image
-	d, err := disk.Usage(req.TempPath)
-	if err != nil {
-		return err
-	}
-	if d.Free/updater.GiB < updater.DownloadDiskSpace {
-		return fmt.Errorf("download and extraction requires up to %d GiB of free space", updater.DownloadDiskSpace)
-	}
-
 	client := updater.NewClient()
-	manifest, err := client.GetInfoManifest(ctx)
-	if err != nil {
-		return err
-	}
-
-	var rel *updater.Release
-	if req.Version == "latest" || req.Version == manifest.Latest.Version {
-		rel = &manifest.Latest
-	} else {
-		for _, r := range manifest.Releases {
-			if req.Version == r.Version {
-				rel = &r
-				break
-			}
-		}
-	}
 
-	if rel == nil {
-		return fmt.Errorf("could not find Debian image %s", req.Version)
+	rel, err := client.GetReleaseByVersion(ctx, req.GetVersion())
+	if err != nil {
+		return fmt.Errorf("could not get release info: %w", err)
 	}
 
 	tmpZip := paths.New(req.GetTempPath(), "arduino-unoq-debian-image-"+rel.Version+".tar.zst")
 	defer func() { _ = tmpZip.RemoveAll() }()
 
-	if err := updater.DownloadFile(ctx, tmpZip, rel, downloadCB, downloader.Config{}); err != nil {
+	downloadCB.Start(rel.Url, rel.Version)
+	if err := client.DownloadFile(ctx, tmpZip, rel, downloadCB.Update, downloader.Config{}); err != nil {
+		// FIXME: Maybe this is redundant?
+		downloadCB.End(false, err.Error())
 		return err
 	}
+	downloadCB.End(true, "")
 
 	tmpZipFile, err := tmpZip.Open()
 	if err != nil {