internal/gomote, buildlet: add execute command implementation

This change adds the implementation for the execute command endpoint.
This enables the caller to execute a command on the gomote instance.
The output from the command will stream back to the caller if requested.

This change also adds helper functions that retrieve the session from
the session pool and the buildlet client.

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

Change-Id: Iec1853667992b5674b07be5e972ac145a6349fca
Reviewed-on: https://go-review.googlesource.com/c/build/+/382494
Trust: Carlos Amedee <[email protected]>
Run-TryBot: Carlos Amedee <[email protected]>
TryBot-Result: Gopher Robot <[email protected]>
Reviewed-by: Alex Rakoczy <[email protected]>

Author: cagedmantis

buildlet/fakebuildletclient.go

@@ -7,6 +7,7 @@ package buildlet
 import (
 	"context"
 	"errors"
+	"fmt"
 	"io"
 	"net"
 	"net/http"
@@ -88,7 +89,19 @@ func (fc *FakeClient) DestroyVM(ts oauth2.TokenSource, proj, zone, instance stri
 
 // Exec fakes the execution.
 func (fc *FakeClient) Exec(ctx context.Context, cmd string, opts ExecOpts) (remoteErr, execErr error) {
-	return nil, errUnimplemented
+	if cmd == "" {
+		return nil, errors.New("invalid command")
+	}
+	if opts.Output == nil {
+		return nil, nil
+	}
+	out := []byte("<this is a song that never ends>")
+	for it := 0; it < 3; it++ {
+		if n, err := opts.Output.Write(out); n != len(out) || err != nil {
+			return nil, fmt.Errorf("Output.Write(...) = %d, %q; want %d, no error", n, err, len(out))
+		}
+	}
+	return nil, nil
 }
 
 // GCEInstanceName gives the fake instance name.

internal/gomote/gomote.go

@@ -10,6 +10,7 @@ package gomote
 import (
 	"context"
 	"errors"
+	"fmt"
 	"log"
 	"regexp"
 	"strings"
@@ -149,12 +150,10 @@ func (s *Server) InstanceAlive(ctx context.Context, req *protos.InstanceAliveReq
 	if req.GetGomoteId() == "" {
 		return nil, status.Errorf(codes.InvalidArgument, "invalid gomote ID")
 	}
-	session, err := s.buildlets.Session(req.GetGomoteId())
+	_, err = s.session(req.GetGomoteId(), creds.ID)
 	if err != nil {
-		return nil, status.Errorf(codes.NotFound, "specified gomote instance does not exist")
-	}
-	if session.OwnerID != creds.ID {
-		return nil, status.Errorf(codes.PermissionDenied, "not allowed to modify this gomote session")
+		// the helper function returns meaningful GRPC error.
+		return nil, err
 	}
 	if err := s.buildlets.RenewTimeout(req.GetGomoteId()); err != nil {
 		return nil, status.Errorf(codes.Internal, "unable to renew timeout")
@@ -170,16 +169,10 @@ func (s *Server) ListDirectory(ctx context.Context, req *protos.ListDirectoryReq
 	if req.GetGomoteId() == "" || req.GetDirectory() == "" {
 		return nil, status.Errorf(codes.InvalidArgument, "invalid arguments")
 	}
-	session, err := s.buildlets.Session(req.GetGomoteId())
+	_, bc, err := s.sessionAndClient(req.GetGomoteId(), creds.ID)
 	if err != nil {
-		return nil, status.Errorf(codes.NotFound, "specified gomote instance does not exist")
-	}
-	if session.OwnerID != creds.ID {
-		return nil, status.Errorf(codes.PermissionDenied, "not allowed to modify this gomote session")
-	}
-	bc, err := s.buildlets.BuildletClient(req.GetGomoteId())
-	if err != nil {
-		return nil, status.Errorf(codes.Internal, "unable to retrieve buildlet client")
+		// the helper function returns meaningful GRPC error.
+		return nil, err
 	}
 	opt := buildlet.ListDirOpts{
 		Recursive: req.GetRecursive(),
@@ -230,12 +223,10 @@ func (s *Server) DestroyInstance(ctx context.Context, req *protos.DestroyInstanc
 	if req.GetGomoteId() == "" {
 		return nil, status.Errorf(codes.InvalidArgument, "invalid gomote ID")
 	}
-	session, err := s.buildlets.Session(req.GetGomoteId())
+	_, err = s.session(req.GetGomoteId(), creds.ID)
 	if err != nil {
-		return nil, status.Errorf(codes.NotFound, "specified gomote instance does not exist")
-	}
-	if session.OwnerID != creds.ID {
-		return nil, status.Errorf(codes.PermissionDenied, "not allowed to modify this gomote session")
+		// the helper function returns meaningful GRPC error.
+		return nil, err
 	}
 	if err := s.buildlets.DestroySession(req.GetGomoteId()); err != nil {
 		log.Printf("DestroyInstance remote.DestroySession(%s) = %s", req.GetGomoteId(), err)
@@ -244,6 +235,54 @@ func (s *Server) DestroyInstance(ctx context.Context, req *protos.DestroyInstanc
 	return &protos.DestroyInstanceResponse{}, nil
 }
 
+// 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.
+func (s *Server) ExecuteCommand(req *protos.ExecuteCommandRequest, stream protos.GomoteService_ExecuteCommandServer) error {
+	creds, err := access.IAPFromContext(stream.Context())
+	if err != nil {
+		return 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 err
+	}
+	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(),
+	})
+	if execErr != nil {
+		// there were system errors preventing the command from being started or seen to completition.
+		return status.Errorf(codes.Aborted, "unable to execute command: %s", execErr)
+	}
+	if remoteErr != nil {
+		// the command succeeded remotely
+		return status.Errorf(codes.Unknown, "command execution failed: %s", remoteErr)
+	}
+	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
+}
+
+// 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),
+	})
+	if err != nil {
+		return 0, fmt.Errorf("unable to send data=%w", err)
+	}
+	return len(p), nil
+}
+
 // RemoveFiles removes files or directories from the gomote instance.
 func (s *Server) RemoveFiles(ctx context.Context, req *protos.RemoveFilesRequest) (*protos.RemoveFilesResponse, error) {
 	creds, err := access.IAPFromContext(ctx)
@@ -255,16 +294,10 @@ func (s *Server) RemoveFiles(ctx context.Context, req *protos.RemoveFilesRequest
 	if req.GetGomoteId() == "" || len(req.GetPaths()) == 0 {
 		return nil, status.Errorf(codes.InvalidArgument, "invalid arguments")
 	}
-	session, err := s.buildlets.Session(req.GetGomoteId())
+	_, bc, err := s.sessionAndClient(req.GetGomoteId(), creds.ID)
 	if err != nil {
-		return nil, status.Errorf(codes.NotFound, "specified gomote instance does not exist")
-	}
-	if session.OwnerID != creds.ID {
-		return nil, status.Errorf(codes.PermissionDenied, "not allowed to modify this gomote session")
-	}
-	bc, err := s.buildlets.BuildletClient(req.GetGomoteId())
-	if err != nil {
-		return nil, status.Errorf(codes.Internal, "unable to retrieve buildlet client")
+		// the helper function returns meaningful GRPC error.
+		return nil, err
 	}
 	if err := bc.RemoveAll(ctx, req.GetPaths()...); err != nil {
 		log.Printf("RemoveFiles buildletClient.RemoveAll(ctx, %q) = %s", req.GetPaths(), err)
@@ -287,21 +320,41 @@ func (s *Server) WriteTGZFromURL(ctx context.Context, req *protos.WriteTGZFromUR
 	if req.GetUrl() == "" {
 		return nil, status.Errorf(codes.InvalidArgument, "missing URL")
 	}
-	session, err := s.buildlets.Session(req.GetGomoteId())
+	_, bc, err := s.sessionAndClient(req.GetGomoteId(), creds.ID)
+	if err != nil {
+		// the helper function returns meaningful GRPC error.
+		return nil, err
+	}
+	if err = bc.PutTarFromURL(ctx, req.GetUrl(), req.GetDirectory()); err != nil {
+		return nil, status.Errorf(codes.FailedPrecondition, "unable to write tar.gz: %s", err)
+	}
+	return &protos.WriteTGZFromURLResponse{}, nil
+}
+
+// session is a helper function that retreives a session associated with the gomoteID and ownerID.
+func (s *Server) session(gomoteID, ownerID string) (*remote.Session, error) {
+	session, err := s.buildlets.Session(gomoteID)
 	if err != nil {
 		return nil, status.Errorf(codes.NotFound, "specified gomote instance does not exist")
 	}
-	if session.OwnerID != creds.ID {
+	if session.OwnerID != ownerID {
 		return nil, status.Errorf(codes.PermissionDenied, "not allowed to modify this gomote session")
 	}
-	bc, err := s.buildlets.BuildletClient(req.GetGomoteId())
+	return session, nil
+}
+
+// sessionAndClient is a helper function that retrieves a session and buildlet client for the
+// associated gomoteID and ownerID.
+func (s *Server) sessionAndClient(gomoteID, ownerID string) (*remote.Session, buildlet.Client, error) {
+	session, err := s.session(gomoteID, ownerID)
 	if err != nil {
-		return nil, status.Errorf(codes.NotFound, "specified gomote instance does not exist")
+		return nil, nil, err
 	}
-	if err = bc.PutTarFromURL(ctx, req.GetUrl(), req.GetDirectory()); err != nil {
-		return nil, status.Errorf(codes.FailedPrecondition, "unable to write tar.gz: %s", err)
+	bc, err := s.buildlets.BuildletClient(gomoteID)
+	if err != nil {
+		return nil, nil, status.Errorf(codes.NotFound, "specified gomote instance does not exist")
 	}
-	return &protos.WriteTGZFromURLResponse{}, nil
+	return session, bc, nil
 }
 
 // isPrivilagedUser returns true if the user is using a Google account.

internal/gomote/gomote_test.go

@@ -401,6 +401,115 @@ func TestDestroyInstanceError(t *testing.T) {
 	}
 }
 
+func TestExecuteCommand(t *testing.T) {
+	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
+	client := setupGomoteTest(t, context.Background())
+	gomoteID := mustCreateInstance(t, client, fakeIAP())
+	stream, err := client.ExecuteCommand(ctx, &protos.ExecuteCommandRequest{
+		GomoteId:          gomoteID,
+		Command:           "ls",
+		SystemLevel:       false,
+		Debug:             false,
+		AppendEnvironment: nil,
+		Path:              nil,
+		Directory:         "/workdir",
+		Args:              []string{"-alh"},
+	})
+	if err != nil {
+		t.Fatalf("client.ExecuteCommand(ctx, req) = response, %s; want no error", err)
+	}
+	out := []string{}
+	for {
+		res, err := stream.Recv()
+		if err != nil && err == io.EOF {
+			break
+		}
+		if err != nil {
+			t.Fatalf("stream.Recv() = _, %s; want no error", err)
+		}
+		out = append(out, res.GetOutput())
+	}
+	if len(out) == 0 {
+		t.Fatalf("output: %q, expected non-empty", out)
+	}
+}
+
+func TestExecuteCommandError(t *testing.T) {
+	// This test will create a gomote instance and attempt to call TestExecuteCommand.
+	// If overrideID is set to true, the test will use a diffrent 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.
+		cmd        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:     "missing command",
+			ctx:      access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
+			wantCode: codes.Aborted,
+		},
+		{
+			desc:       "gomote does not exist",
+			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("foo", "bar")),
+			overrideID: true,
+			gomoteID:   "chucky",
+			cmd:        "ls",
+			wantCode:   codes.NotFound,
+		},
+		{
+			desc:       "wrong gomote id",
+			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("foo", "bar")),
+			overrideID: false,
+			cmd:        "ls",
+			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
+			}
+			stream, err := client.ExecuteCommand(tc.ctx, &protos.ExecuteCommandRequest{
+				GomoteId:          gomoteID,
+				Command:           tc.cmd,
+				SystemLevel:       false,
+				Debug:             false,
+				AppendEnvironment: nil,
+				Path:              nil,
+				Directory:         "/workdir",
+				Args:              []string{"-alh"},
+			})
+			if err != nil {
+				t.Fatalf("unexpected error: %s", err)
+			}
+			res, err := stream.Recv()
+			if err != nil && status.Code(err) != tc.wantCode {
+				t.Fatalf("unexpected error: %s", err)
+			}
+			if err == nil {
+				t.Fatalf("client.ExecuteCommand(ctx, req) = %v, nil; want error", res)
+			}
+		})
+	}
+}
+
 func TestRemoveFiles(t *testing.T) {
 	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
 	client := setupGomoteTest(t, context.Background())