internal/gomote, internal/gomote/protos: add ReadTGZToURL endpoint

This change adds the ReadTGZToURL endpoint implementation. This tars and
zips the directory requested on the gomote instance and uploads it to
GCS. A signed URL is returned which the caller can use to download the file.

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

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

buildlet/fakebuildletclient.go

@@ -109,7 +109,8 @@ func (fc *FakeClient) GCEInstanceName() string { return fc.instanceName }
 
 // GetTar gives a vake tar zipped directory.
 func (fc *FakeClient) GetTar(ctx context.Context, dir string) (io.ReadCloser, error) {
- return nil, errUnimplemented
+ r := strings.NewReader("the gopher goes to the sea and fights the kraken")
+ return io.NopCloser(r), nil
 }
 
 // IPPort provides a fake ip and port pair.

internal/gomote/gomote.go

@@ -273,11 +273,19 @@ func (s *Server) ExecuteCommand(req *protos.ExecuteCommandRequest, stream protos
  remoteErr, execErr := bc.Exec(stream.Context(), req.GetCommand(), buildlet.ExecOpts{
  Dir: req.GetCommand(),
  SystemLevel: req.GetSystemLevel(),
- Output: &execStreamWriter{stream: stream},
- Args: req.GetArgs(),
- ExtraEnv: req.GetAppendEnvironment(),
- Debug: req.GetDebug(),
- Path: req.GetPath(),
+ Output: &streamWriter{writeFunc: func(p []byte) (int, error) {
+ err := stream.Send(&protos.ExecuteCommandResponse{
+ Output: string(p),
+ })
+ if err != nil {
+ return 0, fmt.Errorf("unable to send data=%w", err)
+ }
+ return len(p), nil
+ }},
+ Args: req.GetArgs(),
+ ExtraEnv: req.GetAppendEnvironment(),
+ Debug: req.GetDebug(),
+ Path: req.GetPath(),
  })
  if execErr != nil {
  // there were system errors preventing the command from being started or seen to completition.
@@ -290,21 +298,55 @@ func (s *Server) ExecuteCommand(req *protos.ExecuteCommandRequest, stream protos
  return nil
 }
 
-// execStreamWriter implements the io.Writer interface. Any data writen to it will be streamed
-// as an execute command response.
-type execStreamWriter struct {
- stream protos.GomoteService_ExecuteCommandServer
+// streamWriter implements the io.Writer interface.
+type streamWriter struct {
+ writeFunc func(p []byte) (int, error)
 }
 
-// Write sends data writen to it as an execute command response.
-func (sw *execStreamWriter) Write(p []byte) (int, error) {
- err := sw.stream.Send(&protos.ExecuteCommandResponse{
- Output: string(p),
- })
+// Write calls the writeFunc function with the same arguments passed to the Write function.
+func (sw *streamWriter) Write(p []byte) (int, error) {
+ return sw.writeFunc(p)
+}
+
+// ReadTGZToURL retrieves a directory from the gomote instance and writes the file to GCS. It returns a signed URL which the caller uses
+// to read the file from GCS.
+func (s *Server) ReadTGZToURL(ctx context.Context, req *protos.ReadTGZToURLRequest) (*protos.ReadTGZToURLResponse, error) {
+ creds, err := access.IAPFromContext(ctx)
  if err != nil {
- return 0, fmt.Errorf("unable to send data=%w", err)
+ return nil, status.Errorf(codes.Unauthenticated, "request does not contain the required authentication")
  }
- return len(p), nil
+ _, bc, err := s.sessionAndClient(req.GetGomoteId(), creds.ID)
+ if err != nil {
+ // the helper function returns meaningful GRPC error.
+ return nil, err
+ }
+ tgz, err := bc.GetTar(ctx, req.GetDirectory())
+ if err != nil {
+ return nil, status.Errorf(codes.Aborted, "unable to retrieve tar from gomote instance: %s", err)
+ }
+ defer tgz.Close()
+ objectName := uuid.NewString()
+ objectHandle := s.bucket.Object(objectName)
+ // A context for writes is used to ensure we can cancel the context if a
+ // problem is encountered while writing to the object store. The API documentation
+ // states that the context should be canceled to stop writing without saving the data.
+ writeCtx, cancel := context.WithCancel(ctx)
+ tgzWriter := objectHandle.NewWriter(writeCtx)
+ defer cancel()
+ if _, err = io.Copy(tgzWriter, tgz); err != nil {
+ return nil, status.Errorf(codes.Aborted, "unable to stream tar.gz: %s", err)
+ }
+ // when close is called, the object is stored in the bucket.
+ if err := tgzWriter.Close(); err != nil {
+ return nil, status.Errorf(codes.Aborted, "unable to store object: %s", err)
+ }
+ url, err := s.signURLForDownload(objectName)
+ if err != nil {
+ return nil, status.Errorf(codes.Internal, "unable to create signed URL for download: %s", err)
+ }
+ return &protos.ReadTGZToURLResponse{
+ Url: url,
+ }, nil
 }
 
 // RemoveFiles removes files or directories from the gomote instance.

internal/gomote/gomote_test.go

@@ -522,6 +522,66 @@ func TestExecuteCommandError(t *testing.T) {
  }
 }
 
+func TestReadTGZToURLError(t *testing.T) {
+ // This test will create a gomote instance and attempt to call ReadTGZToURL.
+ // 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.
+ directory string
+ 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",
+ wantCode: codes.NotFound,
+ },
+ {
+ desc: "wrong gomote id",
+ ctx: access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("foo", "bar")),
+ overrideID: false,
+ 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.ReadTGZToURLRequest{
+ GomoteId: gomoteID,
+ Directory: tc.directory,
+ }
+ got, err := client.ReadTGZToURL(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.ReadTGZToURL(ctx, %v) = %v, nil; want error", req, got)
+ }
+ })
+ }
+}
+
 func TestRemoveFiles(t *testing.T) {
  ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
  client := setupGomoteTest(t, context.Background())