From 880f3c9b17044e847aefa66af524c8753fa0a6f0 Mon Sep 17 00:00:00 2001
From: Will Sackfield <sackfield@replicate.com>
Date: Tue, 21 Jan 2025 12:53:45 -0500
Subject: [PATCH] Add Login Token to HTTP POST request

---
 pkg/docker/fast_push.go | 21 ++++++++++++-----
 pkg/docker/login.go     | 52 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 66 insertions(+), 7 deletions(-)

diff --git a/pkg/docker/fast_push.go b/pkg/docker/fast_push.go
index b8d1822b01..6f080c9268 100644
--- a/pkg/docker/fast_push.go
+++ b/pkg/docker/fast_push.go
@@ -21,6 +21,11 @@ const FILES_OBJECT_TYPE = "files"
 func FastPush(image string, projectDir string) error {
 	var wg sync.WaitGroup
 
+	token, err := LoadLoginToken(dockerfile.BaseImageRegistry)
+	if err != nil {
+		return err
+	}
+
 	tmpDir := filepath.Join(projectDir, ".cog", "tmp")
 	// Setting uploadCount to 1 because we always upload the /src directory.
 	uploadCount := 1
@@ -51,7 +56,7 @@ func FastPush(image string, projectDir string) error {
 	// Upload weights
 	for _, weight := range weights {
 		wg.Add(1)
-		go uploadFile(WEIGHTS_OBJECT_TYPE, weight.Digest, weight.Path, &wg, resultChan)
+		go uploadFile(WEIGHTS_OBJECT_TYPE, weight.Digest, weight.Path, token, &wg, resultChan)
 	}
 
 	// Upload apt tar file
@@ -62,7 +67,7 @@ func FastPush(image string, projectDir string) error {
 			return err
 		}
 
-		go uploadFile(FILES_OBJECT_TYPE, hash, aptTarFile, &wg, resultChan)
+		go uploadFile(FILES_OBJECT_TYPE, hash, aptTarFile, token, &wg, resultChan)
 	}
 
 	// Upload python packages.
@@ -78,7 +83,7 @@ func FastPush(image string, projectDir string) error {
 			return err
 		}
 
-		go uploadFile(FILES_OBJECT_TYPE, hash, pythonTar, &wg, resultChan)
+		go uploadFile(FILES_OBJECT_TYPE, hash, pythonTar, token, &wg, resultChan)
 	}
 
 	// Upload user /src.
@@ -91,7 +96,7 @@ func FastPush(image string, projectDir string) error {
 	if err != nil {
 		return err
 	}
-	go uploadFile(FILES_OBJECT_TYPE, hash, srcTar, &wg, resultChan)
+	go uploadFile(FILES_OBJECT_TYPE, hash, srcTar, token, &wg, resultChan)
 
 	wg.Wait()
 
@@ -118,7 +123,7 @@ func startUploadURL(objectType string, digest string, size int64) url.URL {
 	return url
 }
 
-func uploadFile(objectType string, digest string, path string, wg *sync.WaitGroup, resultChan chan<- bool) {
+func uploadFile(objectType string, digest string, path string, token string, wg *sync.WaitGroup, resultChan chan<- bool) {
 	console.Debug("uploading file: " + path)
 	defer wg.Done()
 
@@ -130,7 +135,10 @@ func uploadFile(objectType string, digest string, path string, wg *sync.WaitGrou
 	}
 
 	url := startUploadURL(objectType, digest, info.Size())
-	resp, err := http.Post(url.String(), "", nil)
+	client := &http.Client{}
+	req, _ := http.NewRequest("POST", url.String(), nil)
+	req.Header.Set("Authorization", "Bearer "+token)
+	resp, err := client.Do(req)
 	if err != nil {
 		console.Debug("failed to post file: " + path)
 		resultChan <- false
@@ -149,6 +157,7 @@ func uploadFile(objectType string, digest string, path string, wg *sync.WaitGrou
 	}
 
 	console.Debug("multi-part uploading file: " + path)
+	resultChan <- false
 }
 
 func createPythonPackagesTarFile(image string, tmpDir string) (string, error) {
diff --git a/pkg/docker/login.go b/pkg/docker/login.go
index f550508d4e..751fbd7ae7 100644
--- a/pkg/docker/login.go
+++ b/pkg/docker/login.go
@@ -3,6 +3,7 @@ package docker
 import (
 	"encoding/json"
 	"fmt"
+	"io"
 	"os"
 	"os/exec"
 	"strings"
@@ -29,6 +30,19 @@ func SaveLoginToken(registryHost string, username string, token string) error {
 	return saveAuthToCredentialsStore(credsStore, registryHost, username, token)
 }
 
+func LoadLoginToken(registryHost string) (string, error) {
+	conf := config.LoadDefaultConfigFile(os.Stderr)
+	credsStore := conf.CredentialsStore
+	if credsStore == "" {
+		return loadAuthFromConfig(conf, registryHost)
+	}
+	return loadAuthFromCredentialsStore(credsStore, registryHost)
+}
+
+func dockerCredentialBinary(credsStore string) string {
+	return "docker-credential-" + credsStore
+}
+
 func saveAuthToConfig(conf *configfile.ConfigFile, registryHost string, username string, token string) error {
 	// conf.Save() will base64 encode username and password
 	conf.AuthConfigs[registryHost] = types.AuthConfig{
@@ -42,7 +56,7 @@ func saveAuthToConfig(conf *configfile.ConfigFile, registryHost string, username
 }
 
 func saveAuthToCredentialsStore(credsStore string, registryHost string, username string, token string) error {
-	binary := "docker-credential-" + credsStore
+	binary := dockerCredentialBinary(credsStore)
 	input := credentialHelperInput{
 		Username:  username,
 		Secret:    token,
@@ -70,3 +84,39 @@ func saveAuthToCredentialsStore(credsStore string, registryHost string, username
 	}
 	return nil
 }
+
+func loadAuthFromConfig(conf *configfile.ConfigFile, registryHost string) (string, error) {
+	return conf.AuthConfigs[registryHost].Password, nil
+}
+
+func loadAuthFromCredentialsStore(credsStore string, registryHost string) (string, error) {
+	var out strings.Builder
+	binary := dockerCredentialBinary(credsStore)
+	cmd := exec.Command(binary, "get")
+	cmd.Env = os.Environ()
+	cmd.Stdout = &out
+	cmd.Stderr = &out
+	stdin, err := cmd.StdinPipe()
+	if err != nil {
+		return "", err
+	}
+	defer stdin.Close()
+	console.Debug("$ " + strings.Join(cmd.Args, " "))
+	err = cmd.Start()
+	if err != nil {
+		return "", err
+	}
+	_, err = io.WriteString(stdin, registryHost)
+	if err != nil {
+		return "", err
+	}
+	err = stdin.Close()
+	if err != nil {
+		return "", err
+	}
+	err = cmd.Run()
+	if err != nil {
+		return "", err
+	}
+	return out.String(), nil
+}