Harden StdinService configs

Make StdinService configs explicitly opt into accepting args or stdin by the incoming InputMessages.

Because there are no known users of StdinService, this should not break any existing configs.

PiperOrigin-RevId: 704641242

Author: torsm

fleetspeak/src/client/stdinservice/proto/fleetspeak_stdinservice/config.pb.go

@@ -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;
 }

fleetspeak/src/client/stdinservice/proto/fleetspeak_stdinservice/config.proto

@@ -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
 

fleetspeak/src/client/stdinservice/stdinservice.go

@@ -18,6 +18,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"runtime"
 	"strings"
 	"testing"
 
@@ -28,12 +29,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 +50,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 +60,99 @@ 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)
 	}
+	return om
+}
 
-	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)
+func cmpStdout(got []byte, want string) bool {
+	if runtime.GOOS == "windows" {
+		want = strings.ReplaceAll(want, "\n", "\r\n")
 	}
+	return bytes.Equal(got, []byte(want))
 }
 
-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 := "foo bar\n"
+	if !cmpStdout(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 := "foo bar"
+	if !cmpStdout(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 := "\n"
+	if !cmpStdout(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 := ""
+	if !cmpStdout(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 {