From 55a7bb4946336391ac1458518a02eca7445073a7 Mon Sep 17 00:00:00 2001
From: nguyen-phillip <nguyenphillip@google.com>
Date: Mon, 5 Jan 2026 16:24:55 +0000
Subject: [PATCH] Allow packageHTTP and CanResume to support oauth-prefixed
 URLs.

---
 client/client.go     | 45 ++++++++++++++++++++++++++++----------------
 download/download.go |  6 +++---
 2 files changed, 32 insertions(+), 19 deletions(-)

diff --git a/client/client.go b/client/client.go
index a08a6c2..dd56c58 100644
--- a/client/client.go
+++ b/client/client.go
@@ -202,9 +202,35 @@ func (d *Downloader) unmarshalRepoPackages(ctx context.Context, p, cacheDir stri
 	return d.unmarshalRepoPackagesHTTP(ctx, p, cf)
 }
 
+// NewRequest returns an http.Request, adding auth headers if the URL is
+// prefixed with "oauth-".
+func (d *Downloader) NewRequest(ctx context.Context, method, url string, body io.Reader) (*http.Request, error) {
+	u, useOauth := strings.CutPrefix(url, "oauth-")
+	req, err := http.NewRequestWithContext(ctx, method, u, body)
+	if err != nil {
+		return nil, err
+	}
+	if useOauth {
+		creds, err := google.FindDefaultCredentials(ctx)
+		if err != nil {
+			return nil, fmt.Errorf("failed to obtain creds: %v", err)
+		}
+		token, err := creds.TokenSource.Token()
+		if err != nil {
+			return nil, fmt.Errorf("failed to obtain access token: %v", err)
+		}
+		req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token.AccessToken))
+	}
+	return req, nil
+}
+
 // CanResume returns true if we can resume downloading the specified url.
-func (d *Downloader) CanResume(url string) (bool, int64, error) {
-	resp, err := d.HTTPClient.Head(url)
+func (d *Downloader) CanResume(ctx context.Context, url string) (bool, int64, error) {
+	req, err := d.NewRequest(ctx, http.MethodHead, url, nil)
+	if err != nil {
+		return false, 0, err
+	}
+	resp, err := d.HTTPClient.Do(req)
 	if err != nil {
 		return false, 0, err
 	}
@@ -220,23 +246,10 @@ func (d *Downloader) CanResume(url string) (bool, int64, error) {
 
 // Get gets a url using an optional proxy server, retrying once on any error.
 func (d *Downloader) Get(ctx context.Context, path string) (*http.Response, error) {
-	useOauth := strings.HasPrefix(path, "oauth-")
-	path = strings.TrimPrefix(path, "oauth-")
-	req, err := http.NewRequest(http.MethodGet, path, nil)
+	req, err := d.NewRequest(ctx, http.MethodGet, path, nil)
 	if err != nil {
 		return nil, err
 	}
-	if useOauth {
-		creds, err := google.FindDefaultCredentials(ctx)
-		if err != nil {
-			return nil, fmt.Errorf("failed to obtain creds: %v", err)
-		}
-		token, err := creds.TokenSource.Token()
-		if err != nil {
-			return nil, fmt.Errorf("failed to obtain access token: %v", err)
-		}
-		req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token.AccessToken))
-	}
 	resp, err := d.HTTPClient.Do(req)
 	// We retry on any error once as this mitigates some
 	// connection issues in certain situations.
diff --git a/download/download.go b/download/download.go
index 286ee43..97d24fe 100644
--- a/download/download.go
+++ b/download/download.go
@@ -30,7 +30,7 @@ import (
 	"strings"
 
 	"cloud.google.com/go/storage"
-	humanize "github.com/dustin/go-humanize"
+	"github.com/dustin/go-humanize"
 	"github.com/google/googet/v2/client"
 	"github.com/google/googet/v2/goolib"
 	"github.com/google/googet/v2/oswrap"
@@ -76,11 +76,11 @@ func packageHTTP(ctx context.Context, url, dst, chksum string, downloader *clien
 	// Check that the server supports ranged requests and that the
 	// existing file is smaller than what we want to download.
 	logger.Infof("existing file size: %d", size)
-	ok, length, err := downloader.CanResume(url)
+	ok, length, err := downloader.CanResume(ctx, url)
 	if err != nil {
 		logger.Errorf("CanResume: %v", err)
 	}
-	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
+	req, err := downloader.NewRequest(ctx, http.MethodGet, url, nil)
 	if err != nil {
 		return err
 	}