internal/gomote: add swarming execute command

This change introduces the execute endpoint to the swarming gomote implementation. Fixes golang/go#63776 Change-Id: I084a8b058266bf9296eb54e183e2f8f753fc8f28 Reviewed-on: https://go-review.googlesource.com/c/build/+/537899 LUCI-TryBot-Result: Go LUCI <[email protected]> Reviewed-by: Dmitri Shuralyov <[email protected]> Reviewed-by: Dmitri Shuralyov <[email protected]> Auto-Submit: Carlos Amedee <[email protected]>
golang · Oct 30, 2023 · 7f5600b · 7f5600b
1 parent cabf17f
commit 7f5600b
Show file tree

Hide file tree

Showing 2 changed files with 169 additions and 0 deletions.
diff --git a/internal/gomote/swarming.go b/internal/gomote/swarming.go
@@ -191,6 +191,48 @@ func (ss *SwarmingServer) DestroyInstance(ctx context.Context, req *protos.Destr
 	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 (ss *SwarmingServer) 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")
+	}
+	ses, bc, err := ss.sessionAndClient(stream.Context(), req.GetGomoteId(), creds.ID)
+	if err != nil {
+		// the helper function returns meaningful GRPC error.
+		return err
+	}
+	builderType := req.GetImitateHostType()
+	if builderType == "" {
+		builderType = ses.BuilderType
+	}
+	remoteErr, execErr := bc.Exec(stream.Context(), req.GetCommand(), buildlet.ExecOpts{
+		Dir:         req.GetDirectory(),
+		SystemLevel: req.GetSystemLevel(),
+		Output: &streamWriter{writeFunc: func(p []byte) (int, error) {
+			err := stream.Send(&protos.ExecuteCommandResponse{
+				Output: p,
+			})
+			if err != nil {
+				return 0, fmt.Errorf("unable to send data=%w", err)
+			}
+			return len(p), nil
+		}},
+		Args:  req.GetArgs(),
+		Debug: req.GetDebug(),
+		Path:  req.GetPath(),
+	})
+	if execErr != nil {
+		// there were system errors preventing the command from being started or seen to completion.
+		return status.Errorf(codes.Aborted, "unable to execute command: %s", execErr)
+	}
+	if remoteErr != nil {
+		// the command failed remotely
+		return status.Errorf(codes.Unknown, "command execution failed: %s", remoteErr)
+	}
+	return nil
+}
+
 // ListSwarmingBuilders lists all of the swarming builders which run for gotip. The requester must be authenticated.
 func (ss *SwarmingServer) ListSwarmingBuilders(ctx context.Context, req *protos.ListSwarmingBuildersRequest) (*protos.ListSwarmingBuildersResponse, error) {
 	_, err := access.IAPFromContext(ctx)
@@ -246,6 +288,24 @@ func (ss *SwarmingServer) session(gomoteID, ownerID string) (*remote.Session, er
 	return session, nil
 }
 
+// sessionAndClient is a helper function that retrieves a session and buildlet client for the
+// associated gomoteID and ownerID. The gomote instance timeout is renewed if the gomote id and owner id
+// are valid.
+func (ss *SwarmingServer) sessionAndClient(ctx context.Context, gomoteID, ownerID string) (*remote.Session, buildlet.Client, error) {
+	session, err := ss.session(gomoteID, ownerID)
+	if err != nil {
+		return nil, nil, err
+	}
+	bc, err := ss.buildlets.BuildletClient(gomoteID)
+	if err != nil {
+		return nil, nil, status.Errorf(codes.NotFound, "specified gomote instance does not exist")
+	}
+	if err := ss.buildlets.KeepAlive(ctx, gomoteID); err != nil {
+		log.Printf("gomote: unable to keep alive %s: %s", gomoteID, err)
+	}
+	return session, bc, nil
+}
+
 // SwarmOpts provides additional options for swarming task creation.
 type SwarmOpts struct {
 	// OnInstanceRequested optionally specifies a hook to run synchronously

diff --git a/internal/gomote/swarming_test.go b/internal/gomote/swarming_test.go
@@ -291,6 +291,115 @@ func TestSwarmingDestroyInstanceError(t *testing.T) {
 	}
 }
 
+func TestSwarmingExecuteCommand(t *testing.T) {
+	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
+	client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
+	gomoteID := mustCreateSwarmingInstance(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)
+	}
+	var out []byte
+	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 TestSwarmingExecuteCommandError(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 different gomoteID than
+	// 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 := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
+			gomoteID := mustCreateSwarmingInstance(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 TestSwarmingListInstance(t *testing.T) {
 	client := setupGomoteSwarmingTest(t, context.Background(), mockSwarmClientSimple())
 	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())