Fix monobeam client not authenticating properly (#2148)

* Monobeam client was recreating a standard HTTP
client rather than using the one injected.

Author: 8W9aG

pkg/monobeam/client.go

@@ -88,12 +88,11 @@ func (c *Client) UploadFile(ctx context.Context, objectType string, digest strin
 
 	// Start upload
 	uploadUrl := startUploadURL(objectType, digest)
-	client := &http.Client{}
 	req, err := http.NewRequestWithContext(ctx, http.MethodPost, uploadUrl.String(), nil)
 	if err != nil {
 		return err
 	}
-	resp, err := client.Do(req)
+	resp, err := c.client.Do(req)
 	if err != nil {
 		return err
 	}
@@ -103,7 +102,7 @@ func (c *Client) UploadFile(ctx context.Context, objectType string, digest strin
 	if resp.StatusCode == http.StatusConflict {
 		return nil
 	} else if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusCreated {
-		return errors.New("Bad response: " + strconv.Itoa(resp.StatusCode))
+		return errors.New("Bad response from monobeam: " + strconv.Itoa(resp.StatusCode))
 	}
 
 	// Decode the JSON payload
@@ -150,7 +149,7 @@ func (c *Client) UploadFile(ctx context.Context, objectType string, digest strin
 	if err != nil {
 		return err
 	}
-	beginResp, err := client.Do(req)
+	beginResp, err := c.client.Do(req)
 	if err != nil {
 		return err
 	}
@@ -165,7 +164,7 @@ func (c *Client) UploadFile(ctx context.Context, objectType string, digest strin
 		return err
 	}
 	for i := 0; i < 100; i++ {
-		final, err := checkVerificationStatus(req, client)
+		final, err := c.checkVerificationStatus(req)
 		if final {
 			return err
 		}
@@ -175,27 +174,8 @@ func (c *Client) UploadFile(ctx context.Context, objectType string, digest strin
 	return nil
 }
 
-func baseURL() url.URL {
-	return url.URL{
-		Scheme: env.SchemeFromEnvironment(),
-		Host:   HostFromEnvironment(),
-	}
-}
-
-func startUploadURL(objectType string, digest string) url.URL {
-	uploadUrl := baseURL()
-	uploadUrl.Path = strings.Join([]string{"", "uploads", objectType, "sha256", digest}, "/")
-	return uploadUrl
-}
-
-func verificationURL(objectType string, digest string, uuid string) url.URL {
-	verificationUrl := baseURL()
-	verificationUrl.Path = strings.Join([]string{"", "uploads", objectType, "sha256", digest, uuid, "verification"}, "/")
-	return verificationUrl
-}
-
-func checkVerificationStatus(req *http.Request, client *http.Client) (bool, error) {
-	checkResp, err := client.Do(req)
+func (c *Client) checkVerificationStatus(req *http.Request) (bool, error) {
+	checkResp, err := c.client.Do(req)
 	if err != nil {
 		return true, err
 	}
@@ -219,3 +199,22 @@ func checkVerificationStatus(req *http.Request, client *http.Client) (bool, erro
 
 	return false, nil
 }
+
+func baseURL() url.URL {
+	return url.URL{
+		Scheme: env.SchemeFromEnvironment(),
+		Host:   HostFromEnvironment(),
+	}
+}
+
+func startUploadURL(objectType string, digest string) url.URL {
+	uploadUrl := baseURL()
+	uploadUrl.Path = strings.Join([]string{"", "uploads", objectType, "sha256", digest}, "/")
+	return uploadUrl
+}
+
+func verificationURL(objectType string, digest string, uuid string) url.URL {
+	verificationUrl := baseURL()
+	verificationUrl.Path = strings.Join([]string{"", "uploads", objectType, "sha256", digest, uuid, "verification"}, "/")
+	return verificationUrl
+}

pkg/monobeam/client_test.go

@@ -13,13 +13,16 @@ import (
 	"github.com/stretchr/testify/require"
 	"github.com/vbauerster/mpb/v8"
 
+	"github.com/replicate/cog/pkg/docker/dockertest"
 	"github.com/replicate/cog/pkg/env"
+	r8HTTP "github.com/replicate/cog/pkg/http"
 	"github.com/replicate/cog/pkg/weights"
 )
 
 func TestUploadFile(t *testing.T) {
 	// Setup mock http server
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		require.Equal(t, r8HTTP.UserAgent(), r.Header.Get(r8HTTP.UserAgentHeader))
 		w.WriteHeader(http.StatusConflict)
 	}))
 	defer server.Close()
@@ -44,7 +47,14 @@ func TestUploadFile(t *testing.T) {
 		require.NoError(t, err)
 	}
 
-	client := NewClient(http.DefaultClient)
+	// Setup mock command
+	command := dockertest.NewMockCommand()
+
+	// Setup http client
+	httpClient, err := r8HTTP.ProvideHTTPClient(command)
+	require.NoError(t, err)
+
+	client := NewClient(httpClient)
 	ctx := context.Background()
 	p := mpb.New(
 		mpb.WithRefreshRate(180 * time.Millisecond),