internal/gomote, internal/gomote/protos: implement write files from URL

This change adds the implementation for the WriteFileFromURL endpoint.
The caller will be able to add the contents of an HTTP get call to a
file on the gomote instance. They must set the permissions on the
file. Files located on the gomote transfer GCS bucket will be
retrieved using the GCS storage package and authentication vs a
vanilla HTTP call.

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

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

Author: cagedmantis

buildlet/fakebuildletclient.go

@@ -147,7 +147,10 @@ func (fc *FakeClient) ProxyTCP(port int) (io.ReadWriteCloser, error) { return ni
 // Put places a file on a fake buildlet.
 func (fc *FakeClient) Put(ctx context.Context, r io.Reader, path string, mode os.FileMode) error {
 	// TODO(go.dev/issue/48742) add a file system implementation which would enable proper testing.
-	return errUnimplemented
+	if path == "" {
+		errors.New("invalid argument")
+	}
+	return nil
 }
 
 // PutTar fakes putting  a tar zipped file on a buildldet.

internal/gomote/gomote.go

@@ -11,7 +11,10 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
+	"io/fs"
 	"log"
+	"net/http"
 	"regexp"
 	"strings"
 	"time"
@@ -39,6 +42,8 @@ type scheduler interface {
 // bucketHandle interface used to enable testing of the storage.bucketHandle.
 type bucketHandle interface {
 	GenerateSignedPostPolicyV4(object string, opts *storage.PostPolicyV4Options) (*storage.PostPolicyV4, error)
+	SignedURL(object string, opts *storage.SignedURLOptions) (string, error)
+	Object(name string) *storage.ObjectHandle
 }
 
 // Server is a gomote server implementation.
@@ -380,6 +385,68 @@ func (s *Server) signURLForUpload(object string) (url string, fields map[string]
 	return pv4.URL, pv4.Fields, nil
 }
 
+// signURLForDownload generates a signed URL and fields to be used to upload an object to GCS without authenticating.
+func (s *Server) signURLForDownload(object string) (url string, err error) {
+	url, err = s.bucket.SignedURL(object, &storage.SignedURLOptions{
+		Expires: time.Now().Add(10 * time.Minute),
+		Method:  http.MethodGet,
+		Scheme:  storage.SigningSchemeV4,
+	})
+	if err != nil {
+		return "", fmt.Errorf("unable to generate signed url: %w", err)
+	}
+	return url, err
+}
+
+// WriteFileFromURL initiates an HTTP request to the passed in URL and streams the contents of the request to the gomote instance.
+func (s *Server) WriteFileFromURL(ctx context.Context, req *protos.WriteFileFromURLRequest) (*protos.WriteFileFromURLResponse, error) {
+	creds, err := access.IAPFromContext(ctx)
+	if err != nil {
+		log.Printf("WriteTGZFromURL access.IAPFromContext(ctx) = nil, %s", err)
+		return nil, status.Errorf(codes.Unauthenticated, "request does not contain the required authentication")
+	}
+	_, bc, err := s.sessionAndClient(req.GetGomoteId(), creds.ID)
+	if err != nil {
+		// the helper function returns meaningful GRPC error.
+		return nil, err
+	}
+	var rc io.ReadCloser
+	// objects stored in the gomote staging bucket are only accessible when you have been granted explicit permissions. A builder
+	// requires a signed URL in order to access objects stored in the the gomote staging bucket.
+	if onObjectStore(s.gceBucketName, req.GetUrl()) {
+		object, err := objectFromURL(s.gceBucketName, req.GetUrl())
+		if err != nil {
+			return nil, status.Errorf(codes.InvalidArgument, "invalid object URL")
+		}
+		rc, err = s.bucket.Object(object).NewReader(ctx)
+		if err != nil {
+			return nil, status.Errorf(codes.Internal, "unable to create object reader: %s", err)
+		}
+	} else {
+		httpRequest, err := http.NewRequestWithContext(ctx, http.MethodGet, req.GetUrl(), nil)
+		// TODO(amedee) find sane client defaults, possibly rely on context timeout in request.
+		client := &http.Client{
+			Timeout: 30 * time.Second,
+			Transport: &http.Transport{
+				TLSHandshakeTimeout: 5 * time.Second,
+			},
+		}
+		resp, err := client.Do(httpRequest)
+		if err != nil {
+			return nil, status.Errorf(codes.Aborted, "failed to get file from URL: %s", err)
+		}
+		if resp.StatusCode != http.StatusOK {
+			return nil, status.Errorf(codes.Aborted, "unable to get file from URL: response code: %d", resp.StatusCode)
+		}
+		rc = resp.Body
+	}
+	defer rc.Close()
+	if err := bc.Put(ctx, rc, req.GetFilename(), fs.FileMode(req.GetMode())); err != nil {
+		return nil, status.Errorf(codes.Aborted, "failed to send the file to the gomote instance: %s", err)
+	}
+	return &protos.WriteFileFromURLResponse{}, nil
+}
+
 // WriteTGZFromURL will instruct the gomote instance to download the tar.gz from the provided URL. The tar.gz file will be unpacked in the work directory
 // relative to the directory provided.
 func (s *Server) WriteTGZFromURL(ctx context.Context, req *protos.WriteTGZFromURLRequest) (*protos.WriteTGZFromURLResponse, error) {
@@ -452,3 +519,21 @@ func emailToUser(email string) (string, error) {
 	}
 	return email[strings.Index(email, ":")+1 : strings.LastIndex(email, "@")], nil
 }
+
+// onObjectStore returns true if the the url is for an object on GCS.
+func onObjectStore(bucketName, url string) bool {
+	return strings.HasPrefix(url, fmt.Sprintf("https://storage.googleapis.com/%s/", bucketName))
+}
+
+// objectFromURL returns the object name for an object on GCS.
+func objectFromURL(bucketName, url string) (string, error) {
+	if !onObjectStore(bucketName, url) {
+		return "", errors.New("URL not for gomote transfer bucket")
+	}
+	url = strings.TrimPrefix(url, fmt.Sprintf("https://storage.googleapis.com/%s/", bucketName))
+	pos := strings.Index(url, "?")
+	if pos == -1 {
+		return "", errors.New("invalid object store URL")
+	}
+	return url[:pos], nil
+}

internal/gomote/gomote_test.go

@@ -736,6 +736,87 @@ func TestWriteTGZFromURL(t *testing.T) {
 	}
 }
 
+// TODO(go.dev/issue/48737) add test for files on GCS
+func TestWriteFileFromURL(t *testing.T) {
+	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
+	client := setupGomoteTest(t, context.Background())
+	gomoteID := mustCreateInstance(t, client, fakeIAP())
+	if _, err := client.WriteFileFromURL(ctx, &protos.WriteFileFromURLRequest{
+		GomoteId: gomoteID,
+		Url:      `https://go.dev/dl/go1.17.6.linux-amd64.tar.gz`,
+		Filename: "foo",
+		Mode:     0777,
+	}); err != nil {
+		t.Fatalf("client.WriteFileFromURL(ctx, req) = response, %s; want no error", err)
+	}
+}
+
+func TestWriteFileFromURLError(t *testing.T) {
+	// This test will create a gomote instance and attempt to call TestWriteFileFromURL.
+	// If overrideID is set to true, the test will use a different gomoteID than the
+	// the one created for the test.
+	testCases := []struct {
+		desc       string
+		ctx        context.Context
+		overrideID bool
+		gomoteID   string // Used iff overrideID is true.
+		url        string
+		filename   string
+		mode       uint32
+		wantCode   codes.Code
+	}{
+		{
+			desc:     "unauthenticated request",
+			ctx:      context.Background(),
+			wantCode: codes.Unauthenticated,
+		},
+		{
+			desc:       "missing gomote id",
+			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
+			overrideID: true,
+			gomoteID:   "",
+			wantCode:   codes.NotFound,
+		},
+		{
+			desc:       "gomote does not exist",
+			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("foo", "bar")),
+			overrideID: true,
+			gomoteID:   "chucky",
+			url:        "go.dev/dl/1_14.tar.gz",
+			wantCode:   codes.NotFound,
+		},
+		{
+			desc:       "wrong gomote id",
+			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("foo", "bar")),
+			overrideID: false,
+			url:        "go.dev/dl/1_14.tar.gz",
+			wantCode:   codes.PermissionDenied,
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			client := setupGomoteTest(t, context.Background())
+			gomoteID := mustCreateInstance(t, client, fakeIAP())
+			if tc.overrideID {
+				gomoteID = tc.gomoteID
+			}
+			req := &protos.WriteFileFromURLRequest{
+				GomoteId: gomoteID,
+				Url:      tc.url,
+				Filename: tc.filename,
+				Mode:     0,
+			}
+			got, err := client.WriteFileFromURL(tc.ctx, req)
+			if err != nil && status.Code(err) != tc.wantCode {
+				t.Fatalf("unexpected error: %s; want %s", err, tc.wantCode)
+			}
+			if err == nil {
+				t.Fatalf("client.WriteFileFromURL(ctx, %v) = %v, nil; want error", req, got)
+			}
+		})
+	}
+}
+
 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
@@ -812,6 +893,29 @@ func TestIsPrivilegedUser(t *testing.T) {
 	}
 }
 
+func TestObjectFromURL(t *testing.T) {
+	url := `https://storage.googleapis.com/example-bucket/cat.jpeg?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=example...`
+	bucket := "example-bucket"
+	wantObject := "cat.jpeg"
+	object, err := objectFromURL(bucket, url)
+	if err != nil {
+		t.Fatalf("urlToBucketObject(url) = %q, %s; want %q, no error", object, err, wantObject)
+	}
+	if object != wantObject {
+		t.Fatalf("urlToBucketObject(url) = %q; want %q", object, wantObject)
+	}
+}
+
+func TestObjectFromURLError(t *testing.T) {
+	bucket := "example-bucket"
+	object := "cat.jpeg"
+	url := fmt.Sprintf("https://storage.googleapis.com/%s/%s", bucket, object)
+	got, err := objectFromURL(bucket, url)
+	if err == nil {
+		t.Fatalf("urlToBucketObject(url) = %q, nil; want \"\", error", got)
+	}
+}
+
 func TestEmailToUser(t *testing.T) {
 	testCases := []struct {
 		desc  string
@@ -933,3 +1037,14 @@ func (fbc *fakeBucketHandler) GenerateSignedPostPolicyV4(object string, opts *st
 		},
 	}, nil
 }
+
+func (fbc *fakeBucketHandler) SignedURL(object string, opts *storage.SignedURLOptions) (string, error) {
+	if object == "" || opts == nil {
+		return "", errors.New("invalid arguments")
+	}
+	return fmt.Sprintf("https://localhost/%s?X-Goog-Algorithm=GOOG4-yyy", object), nil
+}
+
+func (fbc *fakeBucketHandler) Object(name string) *storage.ObjectHandle {
+	return &storage.ObjectHandle{}
+}