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 3cfb92de..bdb4b87b 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 arguments provided by the InputMessage will be discarded. + // This should be set to true if the target command does not expect arguments. + RejectArgs bool `protobuf:"varint,2,opt,name=reject_args,json=rejectArgs,proto3" json:"reject_args,omitempty"` + // If true, the input provided by the InputMessage will be discarded. + // This should be set to true if the target command does not expect any input. + RejectStdin bool `protobuf:"varint,3,opt,name=reject_stdin,json=rejectStdin,proto3" json:"reject_stdin,omitempty"` } func (x *Config) Reset() { @@ -69,6 +76,20 @@ func (x *Config) GetCmd() string { return "" } +func (x *Config) GetRejectArgs() bool { + if x != nil { + return x.RejectArgs + } + return false +} + +func (x *Config) GetRejectStdin() bool { + if x != nil { + return x.RejectStdin + } + 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, 0x72, 0x65, 0x6a, 0x65, 0x63, 0x74, 0x5f, 0x61, + 0x72, 0x67, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x08, 0x52, 0x0a, 0x72, 0x65, 0x6a, 0x65, 0x63, + 0x74, 0x41, 0x72, 0x67, 0x73, 0x12, 0x21, 0x0a, 0x0c, 0x72, 0x65, 0x6a, 0x65, 0x63, 0x74, 0x5f, + 0x73, 0x74, 0x64, 0x69, 0x6e, 0x18, 0x03, 0x20, 0x01, 0x28, 0x08, 0x52, 0x0b, 0x72, 0x65, 0x6a, + 0x65, 0x63, 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 b7e7d195..ae9e5bf1 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 arguments provided by the InputMessage will be discarded. + // This should be set to true if the target command does not expect arguments. + bool reject_args = 2; + + // If true, the input provided by the InputMessage will be discarded. + // This should be set to true if the target command does not expect any input. + bool reject_stdin = 3; } diff --git a/fleetspeak/src/client/stdinservice/stdinservice.go b/fleetspeak/src/client/stdinservice/stdinservice.go index 3a06d143..5827251c 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.RejectArgs { + args = im.Args + } + cmd := exec.CommandContext(ctx, s.ssConf.Cmd, args...) + if !s.ssConf.RejectStdin { + 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 7d271b27..f6bda262 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,146 +59,80 @@ 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", } - - 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", } - - 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", + RejectArgs: 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{"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", + RejectStdin: true, } - - 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{