diff --git a/fleetspeak/src/client/stdinservice/proto/fleetspeak_stdinservice/config.pb.go b/fleetspeak/src/client/stdinservice/proto/fleetspeak_stdinservice/config.pb.go index 3cfb92de7..5d06ee1e1 100644 --- a/fleetspeak/src/client/stdinservice/proto/fleetspeak_stdinservice/config.pb.go +++ b/fleetspeak/src/client/stdinservice/proto/fleetspeak_stdinservice/config.pb.go @@ -20,14 +20,21 @@ const ( _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) ) -// The configuration information expected by stdinservice.Factory -// in ClientServiceConfig.config. +// The configuration information expected by stdinservice.Factory in +// ClientServiceConfig.config. type Config struct { state protoimpl.MessageState sizeCache protoimpl.SizeCache unknownFields protoimpl.UnknownFields + // The path of the executable to run. Cmd string `protobuf:"bytes,1,opt,name=cmd,proto3" json:"cmd,omitempty"` + // If true, the command will be ran with arguments provided by the + // InputMessage. + AcceptArgs bool `protobuf:"varint,2,opt,name=accept_args,json=acceptArgs,proto3" json:"accept_args,omitempty"` + // If true, the command will be ran with the input provided by the + // InputMessage. + AcceptStdin bool `protobuf:"varint,3,opt,name=accept_stdin,json=acceptStdin,proto3" json:"accept_stdin,omitempty"` } func (x *Config) Reset() { @@ -69,6 +76,20 @@ func (x *Config) GetCmd() string { return "" } +func (x *Config) GetAcceptArgs() bool { + if x != nil { + return x.AcceptArgs + } + return false +} + +func (x *Config) GetAcceptStdin() bool { + if x != nil { + return x.AcceptStdin + } + return false +} + var File_fleetspeak_src_client_stdinservice_proto_fleetspeak_stdinservice_config_proto protoreflect.FileDescriptor var file_fleetspeak_src_client_stdinservice_proto_fleetspeak_stdinservice_config_proto_rawDesc = []byte{ @@ -78,15 +99,20 @@ var file_fleetspeak_src_client_stdinservice_proto_fleetspeak_stdinservice_config 0x73, 0x70, 0x65, 0x61, 0x6b, 0x5f, 0x73, 0x74, 0x64, 0x69, 0x6e, 0x73, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x63, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x12, 0x17, 0x66, 0x6c, 0x65, 0x65, 0x74, 0x73, 0x70, 0x65, 0x61, 0x6b, 0x2e, 0x73, 0x74, 0x64, 0x69, - 0x6e, 0x73, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x22, 0x1a, 0x0a, 0x06, 0x43, 0x6f, 0x6e, 0x66, + 0x6e, 0x73, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x22, 0x5e, 0x0a, 0x06, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x12, 0x10, 0x0a, 0x03, 0x63, 0x6d, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, - 0x03, 0x63, 0x6d, 0x64, 0x42, 0x5f, 0x5a, 0x5d, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, - 0x6f, 0x6d, 0x2f, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2f, 0x66, 0x6c, 0x65, 0x65, 0x74, 0x73, - 0x70, 0x65, 0x61, 0x6b, 0x2f, 0x66, 0x6c, 0x65, 0x65, 0x74, 0x73, 0x70, 0x65, 0x61, 0x6b, 0x2f, - 0x73, 0x72, 0x63, 0x2f, 0x63, 0x6c, 0x69, 0x65, 0x6e, 0x74, 0x2f, 0x73, 0x74, 0x64, 0x69, 0x6e, - 0x73, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2f, 0x66, 0x6c, - 0x65, 0x65, 0x74, 0x73, 0x70, 0x65, 0x61, 0x6b, 0x5f, 0x73, 0x74, 0x64, 0x69, 0x6e, 0x73, 0x65, - 0x72, 0x76, 0x69, 0x63, 0x65, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, + 0x03, 0x63, 0x6d, 0x64, 0x12, 0x1f, 0x0a, 0x0b, 0x61, 0x63, 0x63, 0x65, 0x70, 0x74, 0x5f, 0x61, + 0x72, 0x67, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x08, 0x52, 0x0a, 0x61, 0x63, 0x63, 0x65, 0x70, + 0x74, 0x41, 0x72, 0x67, 0x73, 0x12, 0x21, 0x0a, 0x0c, 0x61, 0x63, 0x63, 0x65, 0x70, 0x74, 0x5f, + 0x73, 0x74, 0x64, 0x69, 0x6e, 0x18, 0x03, 0x20, 0x01, 0x28, 0x08, 0x52, 0x0b, 0x61, 0x63, 0x63, + 0x65, 0x70, 0x74, 0x53, 0x74, 0x64, 0x69, 0x6e, 0x42, 0x5f, 0x5a, 0x5d, 0x67, 0x69, 0x74, 0x68, + 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2f, 0x66, 0x6c, + 0x65, 0x65, 0x74, 0x73, 0x70, 0x65, 0x61, 0x6b, 0x2f, 0x66, 0x6c, 0x65, 0x65, 0x74, 0x73, 0x70, + 0x65, 0x61, 0x6b, 0x2f, 0x73, 0x72, 0x63, 0x2f, 0x63, 0x6c, 0x69, 0x65, 0x6e, 0x74, 0x2f, 0x73, + 0x74, 0x64, 0x69, 0x6e, 0x73, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x70, 0x72, 0x6f, 0x74, + 0x6f, 0x2f, 0x66, 0x6c, 0x65, 0x65, 0x74, 0x73, 0x70, 0x65, 0x61, 0x6b, 0x5f, 0x73, 0x74, 0x64, + 0x69, 0x6e, 0x73, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, + 0x33, } var ( diff --git a/fleetspeak/src/client/stdinservice/proto/fleetspeak_stdinservice/config.proto b/fleetspeak/src/client/stdinservice/proto/fleetspeak_stdinservice/config.proto index b7e7d195c..ffdb1bef4 100644 --- a/fleetspeak/src/client/stdinservice/proto/fleetspeak_stdinservice/config.proto +++ b/fleetspeak/src/client/stdinservice/proto/fleetspeak_stdinservice/config.proto @@ -4,8 +4,17 @@ package fleetspeak.stdinservice; option go_package = "github.com/google/fleetspeak/fleetspeak/src/client/stdinservice/proto/fleetspeak_stdinservice"; -// The configuration information expected by stdinservice.Factory -// in ClientServiceConfig.config. +// The configuration information expected by stdinservice.Factory in +// ClientServiceConfig.config. message Config { + // The path of the executable to run. string cmd = 1; + + // If true, the command will be ran with arguments provided by the + // InputMessage. + bool accept_args = 2; + + // If true, the command will be ran with the input provided by the + // InputMessage. + bool accept_stdin = 3; } diff --git a/fleetspeak/src/client/stdinservice/stdinservice.go b/fleetspeak/src/client/stdinservice/stdinservice.go index 3a06d143d..14d237ac9 100644 --- a/fleetspeak/src/client/stdinservice/stdinservice.go +++ b/fleetspeak/src/client/stdinservice/stdinservice.go @@ -81,8 +81,14 @@ func (s *StdinService) ProcessMessage(ctx context.Context, m *fspb.Message) erro var stdout, stderr bytes.Buffer - cmd := exec.CommandContext(ctx, s.ssConf.Cmd, im.Args...) - cmd.Stdin = bytes.NewBuffer(im.Input) + var args []string + if s.ssConf.AcceptArgs { + args = im.Args + } + cmd := exec.CommandContext(ctx, s.ssConf.Cmd, args...) + if s.ssConf.AcceptStdin { + cmd.Stdin = bytes.NewBuffer(im.Input) + } cmd.Stdout = &stdout cmd.Stderr = &stderr diff --git a/fleetspeak/src/client/stdinservice/stdinservice_test.go b/fleetspeak/src/client/stdinservice/stdinservice_test.go index 7d271b27a..c1ec049af 100644 --- a/fleetspeak/src/client/stdinservice/stdinservice_test.go +++ b/fleetspeak/src/client/stdinservice/stdinservice_test.go @@ -28,12 +28,11 @@ import ( fspb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak" ) -func TestStdinServiceWithEcho(t *testing.T) { +func mustProcessStdinService(t *testing.T, conf *sspb.Config, im *sspb.InputMessage) *sspb.OutputMessage { + t.Helper() s, err := Factory(&fspb.ClientServiceConfig{ - Name: "EchoService", - Config: anypbtest.New(t, &sspb.Config{ - Cmd: "python", - }), + Name: "TestService", + Config: anypbtest.New(t, conf), }) if err != nil { t.Fatal(err) @@ -50,9 +49,7 @@ func TestStdinServiceWithEcho(t *testing.T) { err = s.ProcessMessage(context.Background(), &fspb.Message{ MessageType: "StdinServiceInputMessage", - Data: anypbtest.New(t, &sspb.InputMessage{ - Args: []string{"-c", `print("foo bar")`}, - }), + Data: anypbtest.New(t, im), }) if err != nil { t.Fatal(err) @@ -62,150 +59,92 @@ func TestStdinServiceWithEcho(t *testing.T) { select { case output = <-outChan: default: - t.Fatal(".ProcessMessage (/bin/echo foo bar) expected to produce message, but none found") + t.Fatal("ProcessMessage() expected to produce message, but none found") } om := &sspb.OutputMessage{} if err := output.Data.UnmarshalTo(om); err != nil { t.Fatal(err) } - - wantStdout := []byte("foo bar\n") - wantStdoutWin := []byte("foo bar\r\n") - if !bytes.Equal(om.Stdout, wantStdout) && - !bytes.Equal(om.Stdout, wantStdoutWin) { - t.Fatalf("unexpected output; got %q, want %q", om.Stdout, wantStdout) - } + return om } -func TestStdinServiceWithCat(t *testing.T) { - s, err := Factory(&fspb.ClientServiceConfig{ - Name: "CatService", - Config: anypbtest.New(t, &sspb.Config{ - Cmd: "python", - }), - }) - if err != nil { - t.Fatal(err) +func TestStdinService_AcceptsArgs(t *testing.T) { + conf := &sspb.Config{ + Cmd: "echo", + AcceptArgs: true, + AcceptStdin: true, } - - outChan := make(chan *fspb.Message, 1) - err = s.Start(&clitesting.MockServiceContext{ - OutChan: outChan, - }) - if err != nil { - t.Fatal(err) + im := &sspb.InputMessage{ + Args: []string{"foo bar"}, } - err = s.ProcessMessage(context.Background(), - &fspb.Message{ - MessageType: "StdinServiceInputMessage", - Data: anypbtest.New(t, &sspb.InputMessage{ - Args: []string{"-c", ` -try: - my_input = raw_input # Python2 compat -except NameError: - my_input = input - -try: - while True: - print(my_input()) -except EOFError: - pass - `}, - Input: []byte("foo bar"), - }), - }) - if err != nil { - t.Fatalf("s.ProcessMessage(...) = %q, want success", err) + om := mustProcessStdinService(t, conf, im) + wantStdout := []byte("foo bar\n") + if !bytes.Equal(om.Stdout, wantStdout) { + t.Fatalf("unexpected output; got %q, want %q", om.Stdout, wantStdout) } +} - var output *fspb.Message - select { - case output = <-outChan: - default: - t.Fatal(".ProcessMessage (/bin/cat <<< 'foo bar') expected to produce message, but none found") +func TestStdinService_AcceptsStdin(t *testing.T) { + conf := &sspb.Config{ + Cmd: "cat", + AcceptArgs: true, + AcceptStdin: true, } - - om := &sspb.OutputMessage{} - if err := output.Data.UnmarshalTo(om); err != nil { - t.Fatal(err) + im := &sspb.InputMessage{ + Input: []byte("foo bar"), } - wantStdout := []byte("foo bar\n") - wantStdoutWin := []byte("foo bar\r\n") - if !bytes.Equal(om.Stdout, wantStdout) && - !bytes.Equal(om.Stdout, wantStdoutWin) { + om := mustProcessStdinService(t, conf, im) + + wantStdout := []byte("foo bar") + if !bytes.Equal(om.Stdout, wantStdout) { t.Fatalf("unexpected output; got %q, want %q", om.Stdout, wantStdout) } } -func TestStdinServiceReportsResourceUsage(t *testing.T) { - s, err := Factory(&fspb.ClientServiceConfig{ - Name: "BashService", - Config: anypbtest.New(t, &sspb.Config{ - Cmd: "python", - }), - }) - if err != nil { - t.Fatal(err) +func TestStdinService_RejectsArgs(t *testing.T) { + conf := &sspb.Config{ + Cmd: "echo", + AcceptArgs: false, + AcceptStdin: false, } - - outChan := make(chan *fspb.Message, 1) - err = s.Start(&clitesting.MockServiceContext{ - OutChan: outChan, - }) - if err != nil { - t.Fatal(err) + im := &sspb.InputMessage{ + Args: []string{"don't print this"}, } + om := mustProcessStdinService(t, conf, im) - err = s.ProcessMessage(context.Background(), - &fspb.Message{ - MessageType: "StdinServiceInputMessage", - Data: anypbtest.New(t, &sspb.InputMessage{ - // Generate some system (os.listdir) and user (everything else) execution time... - Args: []string{"-c", ` -import os -import time - -t0 = time.time() -while time.time() - t0 < 1.: - os.listdir(".") - `}, - }), - }) - if err != nil { - t.Fatal(err) + wantStdout := []byte("\n") + if !bytes.Equal(om.Stdout, wantStdout) { + t.Fatalf("unexpected output; got %q, want %q", om.Stdout, wantStdout) } +} - var output *fspb.Message - select { - case output = <-outChan: - default: - t.Fatal(".ProcessMessage (/bin/bash ...) expected to produce message, but none found") +func TestStdinService_RejectsStdin(t *testing.T) { + conf := &sspb.Config{ + Cmd: "cat", + AcceptArgs: false, + AcceptStdin: false, } - - om := &sspb.OutputMessage{} - if err := output.Data.UnmarshalTo(om); err != nil { - t.Fatal(err) + im := &sspb.InputMessage{ + Input: []byte("don't print this"), } + om := mustProcessStdinService(t, conf, im) - // We don't test for ResourceUsage.MeanResidentMemory because memory is currently not being - // queried after the process has terminated. It's only queried right after launching the command - // in which case it can be recorded as "0" which would be indistinguishable from it not being set - // at all, resulting in a flaky test case. The fact that the other resource usage metrics have - // been set here is good enough for now. - - if om.Timestamp.Seconds <= 0 { - t.Fatalf("unexpected output; StdinServiceOutputMessage.timestamp.seconds not set: %q", om) + wantStdout := []byte("") + if !bytes.Equal(om.Stdout, wantStdout) { + t.Fatalf("unexpected output; got %q, want %q", om.Stdout, wantStdout) } } -func TestStdinServiceCancellation(t *testing.T) { +func TestStdinService_Cancelation(t *testing.T) { s, err := Factory(&fspb.ClientServiceConfig{ Name: "SleepService", Config: anypbtest.New(t, &sspb.Config{ - Cmd: "python", + Cmd: "python", + AcceptArgs: true, + AcceptStdin: true, }), }) if err != nil {