diff --git a/internal/gomote/swarming.go b/internal/gomote/swarming.go index 8b6e014b39..1edb20e4f6 100644 --- 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 index 8c347dfc35..30bf0d0d42 100644 --- 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())