golang.org/x/build@v0.0.0-20240506185731-218518f32b70/internal/gomote/gomote.go (about)

     1  // Copyright 2021 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  //go:build linux || darwin
     6  
     7  package gomote
     8  
     9  import (
    10  	"context"
    11  	"errors"
    12  	"fmt"
    13  	"io"
    14  	"io/fs"
    15  	"log"
    16  	"net/http"
    17  	"regexp"
    18  	"strings"
    19  	"time"
    20  
    21  	"cloud.google.com/go/storage"
    22  	"github.com/google/uuid"
    23  	"golang.org/x/build/buildenv"
    24  	"golang.org/x/build/buildlet"
    25  	"golang.org/x/build/dashboard"
    26  	"golang.org/x/build/internal/access"
    27  	"golang.org/x/build/internal/coordinator/pool/queue"
    28  	"golang.org/x/build/internal/coordinator/remote"
    29  	"golang.org/x/build/internal/coordinator/schedule"
    30  	"golang.org/x/build/internal/envutil"
    31  	"golang.org/x/build/internal/gomote/protos"
    32  	"golang.org/x/build/types"
    33  	"golang.org/x/crypto/ssh"
    34  	"google.golang.org/grpc/codes"
    35  	"google.golang.org/grpc/status"
    36  )
    37  
    38  type scheduler interface {
    39  	State() (st schedule.SchedulerState)
    40  	WaiterState(waiter *queue.SchedItem) (ws types.BuildletWaitStatus)
    41  	GetBuildlet(ctx context.Context, si *queue.SchedItem) (buildlet.Client, error)
    42  }
    43  
    44  // bucketHandle interface used to enable testing of the storage.bucketHandle.
    45  type bucketHandle interface {
    46  	GenerateSignedPostPolicyV4(object string, opts *storage.PostPolicyV4Options) (*storage.PostPolicyV4, error)
    47  	SignedURL(object string, opts *storage.SignedURLOptions) (string, error)
    48  	Object(name string) *storage.ObjectHandle
    49  }
    50  
    51  // Server is a gomote server implementation.
    52  type Server struct {
    53  	// embed the unimplemented server.
    54  	protos.UnimplementedGomoteServiceServer
    55  
    56  	bucket                  bucketHandle
    57  	buildlets               *remote.SessionPool
    58  	gceBucketName           string
    59  	scheduler               scheduler
    60  	sshCertificateAuthority ssh.Signer
    61  }
    62  
    63  // New creates a gomote server. If the rawCAPriKey is invalid, the program will exit.
    64  func New(rsp *remote.SessionPool, sched *schedule.Scheduler, rawCAPriKey []byte, gomoteGCSBucket string, storageClient *storage.Client) *Server {
    65  	signer, err := ssh.ParsePrivateKey(rawCAPriKey)
    66  	if err != nil {
    67  		log.Fatalf("unable to parse raw certificate authority private key into signer=%s", err)
    68  	}
    69  	return &Server{
    70  		bucket:                  storageClient.Bucket(gomoteGCSBucket),
    71  		buildlets:               rsp,
    72  		gceBucketName:           gomoteGCSBucket,
    73  		scheduler:               sched,
    74  		sshCertificateAuthority: signer,
    75  	}
    76  }
    77  
    78  // AddBootstrap adds the bootstrap version of Go to an instance and returns the URL for the bootstrap version. If no
    79  // bootstrap version is defined then the returned version URL will be empty.
    80  func (s *Server) AddBootstrap(ctx context.Context, req *protos.AddBootstrapRequest) (*protos.AddBootstrapResponse, error) {
    81  	creds, err := access.IAPFromContext(ctx)
    82  	if err != nil {
    83  		log.Printf("AddBootstrap access.IAPFromContext(ctx) = nil, %s", err)
    84  		return nil, status.Errorf(codes.Unauthenticated, "request does not contain the required authentication")
    85  	}
    86  	ses, bc, err := s.sessionAndClient(ctx, req.GetGomoteId(), creds.ID)
    87  	if err != nil {
    88  		// the helper function returns meaningful GRPC error.
    89  		return nil, err
    90  	}
    91  	bconf, ok := dashboard.Builders[ses.BuilderType]
    92  	if !ok {
    93  		return nil, status.Errorf(codes.Internal, "unknown builder type")
    94  	}
    95  	url := bconf.GoBootstrapURL(buildenv.Production)
    96  	if url == "" {
    97  		return &protos.AddBootstrapResponse{}, nil
    98  	}
    99  	if err = bc.PutTarFromURL(ctx, url, "go1.4"); err != nil {
   100  		return nil, status.Errorf(codes.Internal, "unable to download bootstrap Go")
   101  	}
   102  	return &protos.AddBootstrapResponse{BootstrapGoUrl: url}, nil
   103  }
   104  
   105  // Authenticate will allow the caller to verify that they are properly authenticated and authorized to interact with the
   106  // Service.
   107  func (s *Server) Authenticate(ctx context.Context, req *protos.AuthenticateRequest) (*protos.AuthenticateResponse, error) {
   108  	_, err := access.IAPFromContext(ctx)
   109  	if err != nil {
   110  		log.Printf("Authenticate access.IAPFromContext(ctx) = nil, %s", err)
   111  		return nil, status.Errorf(codes.Unauthenticated, "request does not contain the required authentication")
   112  	}
   113  	return &protos.AuthenticateResponse{}, nil
   114  }
   115  
   116  // CreateInstance will create a gomote instance for the authenticated user.
   117  func (s *Server) CreateInstance(req *protos.CreateInstanceRequest, stream protos.GomoteService_CreateInstanceServer) error {
   118  	creds, err := access.IAPFromContext(stream.Context())
   119  	if err != nil {
   120  		log.Printf("CreateInstance access.IAPFromContext(ctx) = nil, %s", err)
   121  		return status.Errorf(codes.Unauthenticated, "request does not contain the required authentication")
   122  	}
   123  	if req.GetBuilderType() == "" {
   124  		return status.Errorf(codes.InvalidArgument, "invalid builder type")
   125  	}
   126  	bconf, ok := dashboard.Builders[req.GetBuilderType()]
   127  	if !ok {
   128  		return status.Errorf(codes.InvalidArgument, "unknown builder type")
   129  	}
   130  	if ((!bconf.HostConfig().IsHermetic() && bconf.HostConfig().IsGoogle()) || bconf.IsRestricted()) && !isPrivilegedUser(creds.Email) {
   131  		return status.Errorf(codes.PermissionDenied, "user is unable to create gomote of that builder type")
   132  	}
   133  	userName, err := emailToUser(creds.Email)
   134  	if err != nil {
   135  		return status.Errorf(codes.Internal, "invalid user email format")
   136  	}
   137  	si := &queue.SchedItem{
   138  		HostType:  bconf.HostType,
   139  		IsGomote:  true,
   140  		IsRelease: userName == "relui-prod",
   141  		User:      creds.Email,
   142  	}
   143  	type result struct {
   144  		buildletClient buildlet.Client
   145  		err            error
   146  	}
   147  	rc := make(chan result, 1)
   148  	go func() {
   149  		bc, err := s.scheduler.GetBuildlet(stream.Context(), si)
   150  		rc <- result{bc, err}
   151  	}()
   152  	ticker := time.NewTicker(5 * time.Second)
   153  	defer ticker.Stop()
   154  	for {
   155  		select {
   156  		case <-stream.Context().Done():
   157  			return status.Errorf(codes.DeadlineExceeded, "timed out waiting for gomote instance to be created")
   158  		case <-ticker.C:
   159  			st := s.scheduler.WaiterState(si)
   160  			err := stream.Send(&protos.CreateInstanceResponse{
   161  				Status:       protos.CreateInstanceResponse_WAITING,
   162  				WaitersAhead: int64(st.Ahead),
   163  			})
   164  			if err != nil {
   165  				return status.Errorf(codes.Internal, "unable to stream result: %s", err)
   166  			}
   167  		case r := <-rc:
   168  			if r.err != nil {
   169  				log.Printf("error creating gomote buildlet: %v", r.err)
   170  
   171  				return status.Errorf(codes.Unknown, "gomote creation failed: %s", r.err)
   172  			}
   173  			gomoteID := s.buildlets.AddSession(creds.ID, userName, req.GetBuilderType(), bconf.HostType, r.buildletClient)
   174  			log.Printf("created buildlet %v for %v (%s)", gomoteID, userName, r.buildletClient.String())
   175  			session, err := s.buildlets.Session(gomoteID)
   176  			if err != nil {
   177  				return status.Errorf(codes.Internal, "unable to query for gomote timeout") // this should never happen
   178  			}
   179  			wd, err := r.buildletClient.WorkDir(stream.Context())
   180  			if err != nil {
   181  				return status.Errorf(codes.Internal, "could not read working dir: %v", err)
   182  			}
   183  			err = stream.Send(&protos.CreateInstanceResponse{
   184  				Instance: &protos.Instance{
   185  					GomoteId:    gomoteID,
   186  					BuilderType: req.GetBuilderType(),
   187  					HostType:    bconf.HostType,
   188  					Expires:     session.Expires.Unix(),
   189  					WorkingDir:  wd,
   190  				},
   191  				Status:       protos.CreateInstanceResponse_COMPLETE,
   192  				WaitersAhead: 0,
   193  			})
   194  			if err != nil {
   195  				return status.Errorf(codes.Internal, "unable to stream result: %s", err)
   196  			}
   197  			return nil
   198  		}
   199  	}
   200  }
   201  
   202  // InstanceAlive will ensure that the gomote instance is still alive and will extend the timeout. The requester must be authenticated.
   203  func (s *Server) InstanceAlive(ctx context.Context, req *protos.InstanceAliveRequest) (*protos.InstanceAliveResponse, error) {
   204  	creds, err := access.IAPFromContext(ctx)
   205  	if err != nil {
   206  		log.Printf("InstanceAlive access.IAPFromContext(ctx) = nil, %s", err)
   207  		return nil, status.Errorf(codes.Unauthenticated, "request does not contain the required authentication")
   208  	}
   209  	if req.GetGomoteId() == "" {
   210  		return nil, status.Errorf(codes.InvalidArgument, "invalid gomote ID")
   211  	}
   212  	_, err = s.session(req.GetGomoteId(), creds.ID)
   213  	if err != nil {
   214  		// the helper function returns meaningful GRPC error.
   215  		return nil, err
   216  	}
   217  	if err := s.buildlets.RenewTimeout(req.GetGomoteId()); err != nil {
   218  		return nil, status.Errorf(codes.Internal, "unable to renew timeout")
   219  	}
   220  	return &protos.InstanceAliveResponse{}, nil
   221  }
   222  
   223  // ListDirectory lists the contents of the directory on a gomote instance.
   224  func (s *Server) ListDirectory(ctx context.Context, req *protos.ListDirectoryRequest) (*protos.ListDirectoryResponse, error) {
   225  	creds, err := access.IAPFromContext(ctx)
   226  	if err != nil {
   227  		return nil, status.Errorf(codes.Unauthenticated, "request does not contain the required authentication")
   228  	}
   229  	if req.GetGomoteId() == "" || req.GetDirectory() == "" {
   230  		return nil, status.Errorf(codes.InvalidArgument, "invalid arguments")
   231  	}
   232  	_, bc, err := s.sessionAndClient(ctx, req.GetGomoteId(), creds.ID)
   233  	if err != nil {
   234  		// the helper function returns meaningful GRPC error.
   235  		return nil, err
   236  	}
   237  	opt := buildlet.ListDirOpts{
   238  		Recursive: req.GetRecursive(),
   239  		Digest:    req.GetDigest(),
   240  		Skip:      req.GetSkipFiles(),
   241  	}
   242  	var entries []string
   243  	if err = bc.ListDir(context.Background(), req.GetDirectory(), opt, func(bi buildlet.DirEntry) {
   244  		entries = append(entries, bi.String())
   245  	}); err != nil {
   246  		return nil, status.Errorf(codes.Unimplemented, "method ListDirectory not implemented")
   247  	}
   248  	return &protos.ListDirectoryResponse{
   249  		Entries: entries,
   250  	}, nil
   251  }
   252  
   253  // ListInstances will list the gomote instances owned by the requester. The requester must be authenticated.
   254  func (s *Server) ListInstances(ctx context.Context, req *protos.ListInstancesRequest) (*protos.ListInstancesResponse, error) {
   255  	creds, err := access.IAPFromContext(ctx)
   256  	if err != nil {
   257  		log.Printf("ListInstances access.IAPFromContext(ctx) = nil, %s", err)
   258  		return nil, status.Errorf(codes.Unauthenticated, "request does not contain the required authentication")
   259  	}
   260  	res := &protos.ListInstancesResponse{}
   261  	for _, s := range s.buildlets.List() {
   262  		if s.OwnerID != creds.ID {
   263  			continue
   264  		}
   265  		res.Instances = append(res.Instances, &protos.Instance{
   266  			GomoteId:    s.ID,
   267  			BuilderType: s.BuilderType,
   268  			HostType:    s.HostType,
   269  			Expires:     s.Expires.Unix(),
   270  		})
   271  	}
   272  	return res, nil
   273  }
   274  
   275  // DestroyInstance will destroy a gomote instance. It will ensure that the caller is authenticated and is the owner of the instance
   276  // before it destroys the instance.
   277  func (s *Server) DestroyInstance(ctx context.Context, req *protos.DestroyInstanceRequest) (*protos.DestroyInstanceResponse, error) {
   278  	creds, err := access.IAPFromContext(ctx)
   279  	if err != nil {
   280  		log.Printf("DestroyInstance access.IAPFromContext(ctx) = nil, %s", err)
   281  		return nil, status.Errorf(codes.Unauthenticated, "request does not contain the required authentication")
   282  	}
   283  	if req.GetGomoteId() == "" {
   284  		return nil, status.Errorf(codes.InvalidArgument, "invalid gomote ID")
   285  	}
   286  	_, err = s.session(req.GetGomoteId(), creds.ID)
   287  	if err != nil {
   288  		// the helper function returns meaningful GRPC error.
   289  		return nil, err
   290  	}
   291  	if err := s.buildlets.DestroySession(req.GetGomoteId()); err != nil {
   292  		log.Printf("DestroyInstance remote.DestroySession(%s) = %s", req.GetGomoteId(), err)
   293  		return nil, status.Errorf(codes.Internal, "unable to destroy gomote instance")
   294  	}
   295  	return &protos.DestroyInstanceResponse{}, nil
   296  }
   297  
   298  // ExecuteCommand will execute a command on a gomote instance. The output from the command will be streamed back to the caller if the output is set.
   299  func (s *Server) ExecuteCommand(req *protos.ExecuteCommandRequest, stream protos.GomoteService_ExecuteCommandServer) error {
   300  	creds, err := access.IAPFromContext(stream.Context())
   301  	if err != nil {
   302  		return status.Errorf(codes.Unauthenticated, "request does not contain the required authentication")
   303  	}
   304  	ses, bc, err := s.sessionAndClient(stream.Context(), req.GetGomoteId(), creds.ID)
   305  	if err != nil {
   306  		// the helper function returns meaningful GRPC error.
   307  		return err
   308  	}
   309  	builderType := req.GetImitateHostType()
   310  	if builderType == "" {
   311  		builderType = ses.BuilderType
   312  	}
   313  	conf, ok := dashboard.Builders[builderType]
   314  	if !ok {
   315  		return status.Errorf(codes.Internal, "unable to retrieve configuration for instance")
   316  	}
   317  	remoteErr, execErr := bc.Exec(stream.Context(), req.GetCommand(), buildlet.ExecOpts{
   318  		Dir:         req.GetDirectory(),
   319  		SystemLevel: req.GetSystemLevel(),
   320  		Output: &streamWriter{writeFunc: func(p []byte) (int, error) {
   321  			err := stream.Send(&protos.ExecuteCommandResponse{
   322  				Output: p,
   323  			})
   324  			if err != nil {
   325  				return 0, fmt.Errorf("unable to send data=%w", err)
   326  			}
   327  			return len(p), nil
   328  		}},
   329  		Args:     req.GetArgs(),
   330  		ExtraEnv: envutil.Dedup(conf.GOOS(), append(conf.Env(), req.GetAppendEnvironment()...)),
   331  		Debug:    req.GetDebug(),
   332  		Path:     req.GetPath(),
   333  	})
   334  	if execErr != nil {
   335  		// there were system errors preventing the command from being started or seen to completion.
   336  		return status.Errorf(codes.Aborted, "unable to execute command: %s", execErr)
   337  	}
   338  	if remoteErr != nil {
   339  		// the command succeeded remotely
   340  		return status.Errorf(codes.Unknown, "command execution failed: %s", remoteErr)
   341  	}
   342  	return nil
   343  }
   344  
   345  // streamWriter implements the io.Writer interface.
   346  type streamWriter struct {
   347  	writeFunc func(p []byte) (int, error)
   348  }
   349  
   350  // Write calls the writeFunc function with the same arguments passed to the Write function.
   351  func (sw *streamWriter) Write(p []byte) (int, error) {
   352  	return sw.writeFunc(p)
   353  }
   354  
   355  // ReadTGZToURL retrieves a directory from the gomote instance and writes the file to GCS. It returns a signed URL which the caller uses
   356  // to read the file from GCS.
   357  func (s *Server) ReadTGZToURL(ctx context.Context, req *protos.ReadTGZToURLRequest) (*protos.ReadTGZToURLResponse, error) {
   358  	creds, err := access.IAPFromContext(ctx)
   359  	if err != nil {
   360  		return nil, status.Errorf(codes.Unauthenticated, "request does not contain the required authentication")
   361  	}
   362  	_, bc, err := s.sessionAndClient(ctx, req.GetGomoteId(), creds.ID)
   363  	if err != nil {
   364  		// the helper function returns meaningful GRPC error.
   365  		return nil, err
   366  	}
   367  	tgz, err := bc.GetTar(ctx, req.GetDirectory())
   368  	if err != nil {
   369  		return nil, status.Errorf(codes.Aborted, "unable to retrieve tar from gomote instance: %s", err)
   370  	}
   371  	defer tgz.Close()
   372  	objectName := uuid.NewString()
   373  	objectHandle := s.bucket.Object(objectName)
   374  	// A context for writes is used to ensure we can cancel the context if a
   375  	// problem is encountered while writing to the object store. The API documentation
   376  	// states that the context should be canceled to stop writing without saving the data.
   377  	writeCtx, cancel := context.WithCancel(ctx)
   378  	tgzWriter := objectHandle.NewWriter(writeCtx)
   379  	defer cancel()
   380  	if _, err = io.Copy(tgzWriter, tgz); err != nil {
   381  		return nil, status.Errorf(codes.Aborted, "unable to stream tar.gz: %s", err)
   382  	}
   383  	// when close is called, the object is stored in the bucket.
   384  	if err := tgzWriter.Close(); err != nil {
   385  		return nil, status.Errorf(codes.Aborted, "unable to store object: %s", err)
   386  	}
   387  	url, err := s.signURLForDownload(objectName)
   388  	if err != nil {
   389  		return nil, status.Errorf(codes.Internal, "unable to create signed URL for download: %s", err)
   390  	}
   391  	return &protos.ReadTGZToURLResponse{
   392  		Url: url,
   393  	}, nil
   394  }
   395  
   396  // RemoveFiles removes files or directories from the gomote instance.
   397  func (s *Server) RemoveFiles(ctx context.Context, req *protos.RemoveFilesRequest) (*protos.RemoveFilesResponse, error) {
   398  	creds, err := access.IAPFromContext(ctx)
   399  	if err != nil {
   400  		log.Printf("RemoveFiles access.IAPFromContext(ctx) = nil, %s", err)
   401  		return nil, status.Errorf(codes.Unauthenticated, "request does not contain the required authentication")
   402  	}
   403  	// TODO(go.dev/issue/48742) consider what additional path validation should be implemented.
   404  	if req.GetGomoteId() == "" || len(req.GetPaths()) == 0 {
   405  		return nil, status.Errorf(codes.InvalidArgument, "invalid arguments")
   406  	}
   407  	_, bc, err := s.sessionAndClient(ctx, req.GetGomoteId(), creds.ID)
   408  	if err != nil {
   409  		// the helper function returns meaningful GRPC error.
   410  		return nil, err
   411  	}
   412  	if err := bc.RemoveAll(ctx, req.GetPaths()...); err != nil {
   413  		log.Printf("RemoveFiles buildletClient.RemoveAll(ctx, %q) = %s", req.GetPaths(), err)
   414  		return nil, status.Errorf(codes.Unknown, "unable to remove files")
   415  	}
   416  	return &protos.RemoveFilesResponse{}, nil
   417  }
   418  
   419  // SignSSHKey signs the public SSH key with a certificate. The signed public SSH key is intended for use with the gomote service SSH
   420  // server. It will be signed by the certificate authority of the server and will restrict access to the gomote instance that it was
   421  // signed for.
   422  func (s *Server) SignSSHKey(ctx context.Context, req *protos.SignSSHKeyRequest) (*protos.SignSSHKeyResponse, error) {
   423  	creds, err := access.IAPFromContext(ctx)
   424  	if err != nil {
   425  		return nil, status.Errorf(codes.Unauthenticated, "request does not contain the required authentication")
   426  	}
   427  	session, err := s.session(req.GetGomoteId(), creds.ID)
   428  	if err != nil {
   429  		// the helper function returns meaningful GRPC error.
   430  		return nil, err
   431  	}
   432  	signedPublicKey, err := remote.SignPublicSSHKey(ctx, s.sshCertificateAuthority, req.GetPublicSshKey(), session.ID, session.OwnerID, 5*time.Minute)
   433  	if err != nil {
   434  		return nil, status.Errorf(codes.InvalidArgument, "unable to sign ssh key")
   435  	}
   436  	return &protos.SignSSHKeyResponse{
   437  		SignedPublicSshKey: signedPublicKey,
   438  	}, nil
   439  }
   440  
   441  // UploadFile creates a URL and a set of HTTP post fields which are used to upload a file to a staging GCS bucket. Uploaded files are made available to the
   442  // gomote instances via a subsequent call to one of the WriteFromURL endpoints.
   443  func (s *Server) UploadFile(ctx context.Context, req *protos.UploadFileRequest) (*protos.UploadFileResponse, error) {
   444  	_, err := access.IAPFromContext(ctx)
   445  	if err != nil {
   446  		return nil, status.Errorf(codes.Unauthenticated, "request does not contain the required authentication")
   447  	}
   448  	objectName := uuid.NewString()
   449  	url, fields, err := s.signURLForUpload(objectName)
   450  	if err != nil {
   451  		log.Printf("unable to create signed URL: %s", err)
   452  		return nil, status.Errorf(codes.Internal, "unable to create signed url")
   453  	}
   454  	return &protos.UploadFileResponse{
   455  		Url:        url,
   456  		Fields:     fields,
   457  		ObjectName: objectName,
   458  	}, nil
   459  }
   460  
   461  // signURLForUpload generates a signed URL and a set of http Post fields to be used to upload an object to GCS without authenticating.
   462  func (s *Server) signURLForUpload(object string) (url string, fields map[string]string, err error) {
   463  	if object == "" {
   464  		return "", nil, errors.New("invalid object name")
   465  	}
   466  	pv4, err := s.bucket.GenerateSignedPostPolicyV4(object, &storage.PostPolicyV4Options{
   467  		Expires:  time.Now().Add(10 * time.Minute),
   468  		Insecure: false,
   469  	})
   470  	if err != nil {
   471  		return "", nil, fmt.Errorf("unable to generate signed url: %w", err)
   472  	}
   473  	return pv4.URL, pv4.Fields, nil
   474  }
   475  
   476  // signURLForDownload generates a signed URL and fields to be used to upload an object to GCS without authenticating.
   477  func (s *Server) signURLForDownload(object string) (url string, err error) {
   478  	url, err = s.bucket.SignedURL(object, &storage.SignedURLOptions{
   479  		Expires: time.Now().Add(10 * time.Minute),
   480  		Method:  http.MethodGet,
   481  		Scheme:  storage.SigningSchemeV4,
   482  	})
   483  	if err != nil {
   484  		return "", fmt.Errorf("unable to generate signed url: %w", err)
   485  	}
   486  	return url, err
   487  }
   488  
   489  // WriteFileFromURL initiates an HTTP request to the passed in URL and streams the contents of the request to the gomote instance.
   490  func (s *Server) WriteFileFromURL(ctx context.Context, req *protos.WriteFileFromURLRequest) (*protos.WriteFileFromURLResponse, error) {
   491  	creds, err := access.IAPFromContext(ctx)
   492  	if err != nil {
   493  		log.Printf("WriteTGZFromURL access.IAPFromContext(ctx) = nil, %s", err)
   494  		return nil, status.Errorf(codes.Unauthenticated, "request does not contain the required authentication")
   495  	}
   496  	_, bc, err := s.sessionAndClient(ctx, req.GetGomoteId(), creds.ID)
   497  	if err != nil {
   498  		// the helper function returns meaningful GRPC error.
   499  		return nil, err
   500  	}
   501  	var rc io.ReadCloser
   502  	// objects stored in the gomote staging bucket are only accessible when you have been granted explicit permissions. A builder
   503  	// requires a signed URL in order to access objects stored in the gomote staging bucket.
   504  	if onObjectStore(s.gceBucketName, req.GetUrl()) {
   505  		object, err := objectFromURL(s.gceBucketName, req.GetUrl())
   506  		if err != nil {
   507  			return nil, status.Errorf(codes.InvalidArgument, "invalid object URL")
   508  		}
   509  		rc, err = s.bucket.Object(object).NewReader(ctx)
   510  		if err != nil {
   511  			return nil, status.Errorf(codes.Internal, "unable to create object reader: %s", err)
   512  		}
   513  	} else {
   514  		httpRequest, err := http.NewRequestWithContext(ctx, http.MethodGet, req.GetUrl(), nil)
   515  		// TODO(amedee) find sane client defaults, possibly rely on context timeout in request.
   516  		client := &http.Client{
   517  			Timeout: 30 * time.Second,
   518  			Transport: &http.Transport{
   519  				TLSHandshakeTimeout: 5 * time.Second,
   520  			},
   521  		}
   522  		resp, err := client.Do(httpRequest)
   523  		if err != nil {
   524  			return nil, status.Errorf(codes.Aborted, "failed to get file from URL: %s", err)
   525  		}
   526  		if resp.StatusCode != http.StatusOK {
   527  			return nil, status.Errorf(codes.Aborted, "unable to get file from %q: response code: %d", req.GetUrl(), resp.StatusCode)
   528  		}
   529  		rc = resp.Body
   530  	}
   531  	defer rc.Close()
   532  	if err := bc.Put(ctx, rc, req.GetFilename(), fs.FileMode(req.GetMode())); err != nil {
   533  		return nil, status.Errorf(codes.Aborted, "failed to send the file to the gomote instance: %s", err)
   534  	}
   535  	return &protos.WriteFileFromURLResponse{}, nil
   536  }
   537  
   538  // 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
   539  // relative to the directory provided.
   540  func (s *Server) WriteTGZFromURL(ctx context.Context, req *protos.WriteTGZFromURLRequest) (*protos.WriteTGZFromURLResponse, error) {
   541  	creds, err := access.IAPFromContext(ctx)
   542  	if err != nil {
   543  		log.Printf("WriteTGZFromURL access.IAPFromContext(ctx) = nil, %s", err)
   544  		return nil, status.Errorf(codes.Unauthenticated, "request does not contain the required authentication")
   545  	}
   546  	if req.GetGomoteId() == "" {
   547  		return nil, status.Errorf(codes.InvalidArgument, "invalid gomote ID")
   548  	}
   549  	if req.GetUrl() == "" {
   550  		return nil, status.Errorf(codes.InvalidArgument, "missing URL")
   551  	}
   552  	_, bc, err := s.sessionAndClient(ctx, req.GetGomoteId(), creds.ID)
   553  	if err != nil {
   554  		// the helper function returns meaningful GRPC error.
   555  		return nil, err
   556  	}
   557  	url := req.GetUrl()
   558  	if onObjectStore(s.gceBucketName, url) {
   559  		object, err := objectFromURL(s.gceBucketName, url)
   560  		if err != nil {
   561  			return nil, status.Errorf(codes.InvalidArgument, "invalid URL")
   562  		}
   563  		url, err = s.signURLForDownload(object)
   564  		if err != nil {
   565  			return nil, status.Errorf(codes.Aborted, "unable to sign url for download: %s", err)
   566  		}
   567  	}
   568  	if err := bc.PutTarFromURL(ctx, url, req.GetDirectory()); err != nil {
   569  		return nil, status.Errorf(codes.FailedPrecondition, "unable to write tar.gz: %s", err)
   570  	}
   571  	return &protos.WriteTGZFromURLResponse{}, nil
   572  }
   573  
   574  // session is a helper function that retrieves a session associated with the gomoteID and ownerID.
   575  func (s *Server) session(gomoteID, ownerID string) (*remote.Session, error) {
   576  	session, err := s.buildlets.Session(gomoteID)
   577  	if err != nil {
   578  		return nil, status.Errorf(codes.NotFound, "specified gomote instance does not exist")
   579  	}
   580  	if session.OwnerID != ownerID {
   581  		return nil, status.Errorf(codes.PermissionDenied, "not allowed to modify this gomote session")
   582  	}
   583  	return session, nil
   584  }
   585  
   586  // sessionAndClient is a helper function that retrieves a session and buildlet client for the
   587  // associated gomoteID and ownerID. The gomote instance timeout is renewed if the gomote id and owner id
   588  // are valid.
   589  func (s *Server) sessionAndClient(ctx context.Context, gomoteID, ownerID string) (*remote.Session, buildlet.Client, error) {
   590  	session, err := s.session(gomoteID, ownerID)
   591  	if err != nil {
   592  		return nil, nil, err
   593  	}
   594  	bc, err := s.buildlets.BuildletClient(gomoteID)
   595  	if err != nil {
   596  		return nil, nil, status.Errorf(codes.NotFound, "specified gomote instance does not exist")
   597  	}
   598  	if err := s.buildlets.KeepAlive(ctx, gomoteID); err != nil {
   599  		log.Printf("gomote: unable to keep alive %s: %s", gomoteID, err)
   600  	}
   601  	return session, bc, nil
   602  }
   603  
   604  // isPrivilegedUser returns true if the user is trusted to use sensitive machines.
   605  // The user has to be a part of the appropriate IAM group.
   606  func isPrivilegedUser(email string) bool {
   607  	return strings.HasSuffix(email, "@google.com") || strings.HasSuffix(email, "@symbolic-datum-552.iam.gserviceaccount.com") ||
   608  		strings.HasSuffix(email, "@go-security-trybots.iam.gserviceaccount.com")
   609  }
   610  
   611  // iapEmailRE matches the email string returned by Identity Aware Proxy for sessions where
   612  // the authority is Google.
   613  var iapEmailRE = regexp.MustCompile(`^accounts\.google\.com:.+@.+\..+$`)
   614  
   615  // emailToUser returns the displayed user for the IAP email string passed in.
   616  // For example, "accounts.google.com:example@gmail.com" -> "example"
   617  func emailToUser(email string) (string, error) {
   618  	if match := iapEmailRE.MatchString(email); !match {
   619  		return "", errors.New("invalid email format")
   620  	}
   621  	return email[strings.Index(email, ":")+1 : strings.LastIndex(email, "@")], nil
   622  }
   623  
   624  // onObjectStore returns true if the url is for an object on GCS.
   625  func onObjectStore(bucketName, url string) bool {
   626  	return strings.HasPrefix(url, fmt.Sprintf("https://storage.googleapis.com/%s/", bucketName))
   627  }
   628  
   629  // objectFromURL returns the object name for an object on GCS.
   630  func objectFromURL(bucketName, url string) (string, error) {
   631  	if !onObjectStore(bucketName, url) {
   632  		return "", errors.New("URL not for gomote transfer bucket")
   633  	}
   634  	objectName := strings.TrimPrefix(url, fmt.Sprintf("https://storage.googleapis.com/%s/", bucketName))
   635  	return objectName, nil
   636  }