internal/gomote: enable downloads from GCS bucket in WriteTGZFromURL

This change enables downloads from the gomote transfer GCS bucket for
gomote instances that do not have permission to read from that bucket.
Any URL that is passed in for that bucket will have a signed URL
created for it and that will be passed on to the gomote instance.

For golang/go#47521
Updates golang/go#48742

Change-Id: I5874efd1349f4154aea58677c734653b18cd88c9
Reviewed-on: https://go-review.googlesource.com/c/build/+/397597
Reviewed-by: Heschi Kreinick <[email protected]>
Reviewed-by: Carlos Amedee <[email protected]>
Reviewed-by: Carlos Amedee <[email protected]>
Run-TryBot: Carlos Amedee <[email protected]>
Auto-Submit: Carlos Amedee <[email protected]>
TryBot-Result: Gopher Robot <[email protected]>

Author: cagedmantis

internal/gomote/gomote.go

@@ -466,7 +466,18 @@ func (s *Server) WriteTGZFromURL(ctx context.Context, req *protos.WriteTGZFromUR
 		// the helper function returns meaningful GRPC error.
 		return nil, err
 	}
-	if err = bc.PutTarFromURL(ctx, req.GetUrl(), req.GetDirectory()); err != nil {
+	url := req.GetUrl()
+	if onObjectStore(s.gceBucketName, url) {
+		object, err := objectFromURL(s.gceBucketName, url)
+		if err != nil {
+			return nil, status.Errorf(codes.InvalidArgument, "invalid URL")
+		}
+		url, err = s.signURLForDownload(object)
+		if err != nil {
+			return nil, status.Errorf(codes.Aborted, "unable to sign url for download: %s", err)
+		}
+	}
+	if err := bc.PutTarFromURL(ctx, url, req.GetDirectory()); err != nil {
 		return nil, status.Errorf(codes.FailedPrecondition, "unable to write tar.gz: %s", err)
 	}
 	return &protos.WriteTGZFromURLResponse{}, nil

internal/gomote/gomote_test.go

@@ -723,19 +723,6 @@ func TestUploadFileError(t *testing.T) {
 	}
 }
 
-func TestWriteTGZFromURL(t *testing.T) {
-	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
-	client := setupGomoteTest(t, context.Background())
-	gomoteID := mustCreateInstance(t, client, fakeIAP())
-	if _, err := client.WriteTGZFromURL(ctx, &protos.WriteTGZFromURLRequest{
-		GomoteId:  gomoteID,
-		Directory: "foo",
-		Url:       `https://go.dev/dl/go1.17.6.linux-amd64.tar.gz`,
-	}); err != nil {
-		t.Fatalf("client.WriteTGZFromURL(ctx, req) = response, %s; want no error", err)
-	}
-}
-
 // TODO(go.dev/issue/48737) add test for files on GCS
 func TestWriteFileFromURL(t *testing.T) {
 	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
@@ -817,6 +804,32 @@ func TestWriteFileFromURLError(t *testing.T) {
 	}
 }
 
+func TestWriteTGZFromURL(t *testing.T) {
+	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
+	client := setupGomoteTest(t, context.Background())
+	gomoteID := mustCreateInstance(t, client, fakeIAP())
+	if _, err := client.WriteTGZFromURL(ctx, &protos.WriteTGZFromURLRequest{
+		GomoteId:  gomoteID,
+		Directory: "foo",
+		Url:       `https://go.dev/dl/go1.17.6.linux-amd64.tar.gz`,
+	}); err != nil {
+		t.Fatalf("client.WriteTGZFromURL(ctx, req) = response, %s; want no error", err)
+	}
+}
+
+func TestWriteTGZFromURLGomoteStaging(t *testing.T) {
+	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
+	client := setupGomoteTest(t, context.Background())
+	gomoteID := mustCreateInstance(t, client, fakeIAP())
+	if _, err := client.WriteTGZFromURL(ctx, &protos.WriteTGZFromURLRequest{
+		GomoteId:  gomoteID,
+		Directory: "foo",
+		Url:       fmt.Sprintf("https://storage.googleapis.com/%s/go1.17.6.linux-amd64.tar.gz?field=x", testBucketName),
+	}); err != nil {
+		t.Fatalf("client.WriteTGZFromURL(ctx, req) = response, %s; want no error", err)
+	}
+}
+
 func TestWriteTGZFromURLError(t *testing.T) {
 	// This test will create a gomote instance and attempt to call TestWriteTGZFromURL.
 	// If overrideID is set to true, the test will use a different gomoteID than the
@@ -862,6 +875,12 @@ func TestWriteTGZFromURLError(t *testing.T) {
 			url:        "go.dev/dl/1_14.tar.gz",
 			wantCode:   codes.PermissionDenied,
 		},
+		{
+			desc:     "invalid gomote staging bucket URL",
+			ctx:      access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
+			url:      fmt.Sprintf("https://storage.googleapis.com/%s/go1.17.6.linux-amd64.tar.gz", testBucketName),
+			wantCode: codes.InvalidArgument,
+		},
 	}
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {