Skip to content

Commit

Permalink
Harden StdinService configs
Browse files Browse the repository at this point in the history
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
  • Loading branch information
torsm authored and copybara-github committed Dec 10, 2024
1 parent ceb2de0 commit 697afc0
Show file tree
Hide file tree
Showing 4 changed files with 103 additions and 131 deletions.

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
10 changes: 8 additions & 2 deletions fleetspeak/src/client/stdinservice/stdinservice.go
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
165 changes: 48 additions & 117 deletions fleetspeak/src/client/stdinservice/stdinservice_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
Expand All @@ -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{
Expand Down

0 comments on commit 697afc0

Please sign in to comment.