Skip to content

Commit 697afc0

Browse files
torsmcopybara-github
authored andcommitted
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
1 parent ceb2de0 commit 697afc0

File tree

4 files changed

+103
-131
lines changed

4 files changed

+103
-131
lines changed

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

Lines changed: 36 additions & 10 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

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

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,17 @@ package fleetspeak.stdinservice;
44

55
option go_package = "github.com/google/fleetspeak/fleetspeak/src/client/stdinservice/proto/fleetspeak_stdinservice";
66

7-
// The configuration information expected by stdinservice.Factory
8-
// in ClientServiceConfig.config.
7+
// The configuration information expected by stdinservice.Factory in
8+
// ClientServiceConfig.config.
99
message Config {
10+
// The path of the executable to run.
1011
string cmd = 1;
12+
13+
// If true, the arguments provided by the InputMessage will be discarded.
14+
// This should be set to true if the target command does not expect arguments.
15+
bool reject_args = 2;
16+
17+
// If true, the input provided by the InputMessage will be discarded.
18+
// This should be set to true if the target command does not expect any input.
19+
bool reject_stdin = 3;
1120
}

fleetspeak/src/client/stdinservice/stdinservice.go

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -81,8 +81,14 @@ func (s *StdinService) ProcessMessage(ctx context.Context, m *fspb.Message) erro
8181

8282
var stdout, stderr bytes.Buffer
8383

84-
cmd := exec.CommandContext(ctx, s.ssConf.Cmd, im.Args...)
85-
cmd.Stdin = bytes.NewBuffer(im.Input)
84+
var args []string
85+
if !s.ssConf.RejectArgs {
86+
args = im.Args
87+
}
88+
cmd := exec.CommandContext(ctx, s.ssConf.Cmd, args...)
89+
if !s.ssConf.RejectStdin {
90+
cmd.Stdin = bytes.NewBuffer(im.Input)
91+
}
8692
cmd.Stdout = &stdout
8793
cmd.Stderr = &stderr
8894

fleetspeak/src/client/stdinservice/stdinservice_test.go

Lines changed: 48 additions & 117 deletions
Original file line numberDiff line numberDiff line change
@@ -28,12 +28,11 @@ import (
2828
fspb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak"
2929
)
3030

31-
func TestStdinServiceWithEcho(t *testing.T) {
31+
func mustProcessStdinService(t *testing.T, conf *sspb.Config, im *sspb.InputMessage) *sspb.OutputMessage {
32+
t.Helper()
3233
s, err := Factory(&fspb.ClientServiceConfig{
33-
Name: "EchoService",
34-
Config: anypbtest.New(t, &sspb.Config{
35-
Cmd: "python",
36-
}),
34+
Name: "TestService",
35+
Config: anypbtest.New(t, conf),
3736
})
3837
if err != nil {
3938
t.Fatal(err)
@@ -50,9 +49,7 @@ func TestStdinServiceWithEcho(t *testing.T) {
5049
err = s.ProcessMessage(context.Background(),
5150
&fspb.Message{
5251
MessageType: "StdinServiceInputMessage",
53-
Data: anypbtest.New(t, &sspb.InputMessage{
54-
Args: []string{"-c", `print("foo bar")`},
55-
}),
52+
Data: anypbtest.New(t, im),
5653
})
5754
if err != nil {
5855
t.Fatal(err)
@@ -62,146 +59,80 @@ func TestStdinServiceWithEcho(t *testing.T) {
6259
select {
6360
case output = <-outChan:
6461
default:
65-
t.Fatal(".ProcessMessage (/bin/echo foo bar) expected to produce message, but none found")
62+
t.Fatal("ProcessMessage() expected to produce message, but none found")
6663
}
6764

6865
om := &sspb.OutputMessage{}
6966
if err := output.Data.UnmarshalTo(om); err != nil {
7067
t.Fatal(err)
7168
}
72-
73-
wantStdout := []byte("foo bar\n")
74-
wantStdoutWin := []byte("foo bar\r\n")
75-
if !bytes.Equal(om.Stdout, wantStdout) &&
76-
!bytes.Equal(om.Stdout, wantStdoutWin) {
77-
t.Fatalf("unexpected output; got %q, want %q", om.Stdout, wantStdout)
78-
}
69+
return om
7970
}
8071

81-
func TestStdinServiceWithCat(t *testing.T) {
82-
s, err := Factory(&fspb.ClientServiceConfig{
83-
Name: "CatService",
84-
Config: anypbtest.New(t, &sspb.Config{
85-
Cmd: "python",
86-
}),
87-
})
88-
if err != nil {
89-
t.Fatal(err)
72+
func TestStdinService_AcceptsArgs(t *testing.T) {
73+
conf := &sspb.Config{
74+
Cmd: "echo",
9075
}
91-
92-
outChan := make(chan *fspb.Message, 1)
93-
err = s.Start(&clitesting.MockServiceContext{
94-
OutChan: outChan,
95-
})
96-
if err != nil {
97-
t.Fatal(err)
76+
im := &sspb.InputMessage{
77+
Args: []string{"foo bar"},
9878
}
9979

100-
err = s.ProcessMessage(context.Background(),
101-
&fspb.Message{
102-
MessageType: "StdinServiceInputMessage",
103-
Data: anypbtest.New(t, &sspb.InputMessage{
104-
Args: []string{"-c", `
105-
try:
106-
my_input = raw_input # Python2 compat
107-
except NameError:
108-
my_input = input
109-
110-
try:
111-
while True:
112-
print(my_input())
113-
except EOFError:
114-
pass
115-
`},
116-
Input: []byte("foo bar"),
117-
}),
118-
})
119-
if err != nil {
120-
t.Fatalf("s.ProcessMessage(...) = %q, want success", err)
80+
om := mustProcessStdinService(t, conf, im)
81+
wantStdout := []byte("foo bar\n")
82+
if !bytes.Equal(om.Stdout, wantStdout) {
83+
t.Fatalf("unexpected output; got %q, want %q", om.Stdout, wantStdout)
12184
}
85+
}
12286

123-
var output *fspb.Message
124-
select {
125-
case output = <-outChan:
126-
default:
127-
t.Fatal(".ProcessMessage (/bin/cat <<< 'foo bar') expected to produce message, but none found")
87+
func TestStdinService_AcceptsStdin(t *testing.T) {
88+
conf := &sspb.Config{
89+
Cmd: "cat",
12890
}
129-
130-
om := &sspb.OutputMessage{}
131-
if err := output.Data.UnmarshalTo(om); err != nil {
132-
t.Fatal(err)
91+
im := &sspb.InputMessage{
92+
Input: []byte("foo bar"),
13393
}
13494

135-
wantStdout := []byte("foo bar\n")
136-
wantStdoutWin := []byte("foo bar\r\n")
137-
if !bytes.Equal(om.Stdout, wantStdout) &&
138-
!bytes.Equal(om.Stdout, wantStdoutWin) {
95+
om := mustProcessStdinService(t, conf, im)
96+
97+
wantStdout := []byte("foo bar")
98+
if !bytes.Equal(om.Stdout, wantStdout) {
13999
t.Fatalf("unexpected output; got %q, want %q", om.Stdout, wantStdout)
140100
}
141101
}
142102

143-
func TestStdinServiceReportsResourceUsage(t *testing.T) {
144-
s, err := Factory(&fspb.ClientServiceConfig{
145-
Name: "BashService",
146-
Config: anypbtest.New(t, &sspb.Config{
147-
Cmd: "python",
148-
}),
149-
})
150-
if err != nil {
151-
t.Fatal(err)
103+
func TestStdinService_RejectsArgs(t *testing.T) {
104+
conf := &sspb.Config{
105+
Cmd: "echo",
106+
RejectArgs: true,
152107
}
153-
154-
outChan := make(chan *fspb.Message, 1)
155-
err = s.Start(&clitesting.MockServiceContext{
156-
OutChan: outChan,
157-
})
158-
if err != nil {
159-
t.Fatal(err)
108+
im := &sspb.InputMessage{
109+
Args: []string{"don't print this"},
160110
}
111+
om := mustProcessStdinService(t, conf, im)
161112

162-
err = s.ProcessMessage(context.Background(),
163-
&fspb.Message{
164-
MessageType: "StdinServiceInputMessage",
165-
Data: anypbtest.New(t, &sspb.InputMessage{
166-
// Generate some system (os.listdir) and user (everything else) execution time...
167-
Args: []string{"-c", `
168-
import os
169-
import time
170-
171-
t0 = time.time()
172-
while time.time() - t0 < 1.:
173-
os.listdir(".")
174-
`},
175-
}),
176-
})
177-
if err != nil {
178-
t.Fatal(err)
113+
wantStdout := []byte("\n")
114+
if !bytes.Equal(om.Stdout, wantStdout) {
115+
t.Fatalf("unexpected output; got %q, want %q", om.Stdout, wantStdout)
179116
}
117+
}
180118

181-
var output *fspb.Message
182-
select {
183-
case output = <-outChan:
184-
default:
185-
t.Fatal(".ProcessMessage (/bin/bash ...) expected to produce message, but none found")
119+
func TestStdinService_RejectsStdin(t *testing.T) {
120+
conf := &sspb.Config{
121+
Cmd: "cat",
122+
RejectStdin: true,
186123
}
187-
188-
om := &sspb.OutputMessage{}
189-
if err := output.Data.UnmarshalTo(om); err != nil {
190-
t.Fatal(err)
124+
im := &sspb.InputMessage{
125+
Input: []byte("don't print this"),
191126
}
127+
om := mustProcessStdinService(t, conf, im)
192128

193-
// We don't test for ResourceUsage.MeanResidentMemory because memory is currently not being
194-
// queried after the process has terminated. It's only queried right after launching the command
195-
// in which case it can be recorded as "0" which would be indistinguishable from it not being set
196-
// at all, resulting in a flaky test case. The fact that the other resource usage metrics have
197-
// been set here is good enough for now.
198-
199-
if om.Timestamp.Seconds <= 0 {
200-
t.Fatalf("unexpected output; StdinServiceOutputMessage.timestamp.seconds not set: %q", om)
129+
wantStdout := []byte("")
130+
if !bytes.Equal(om.Stdout, wantStdout) {
131+
t.Fatalf("unexpected output; got %q, want %q", om.Stdout, wantStdout)
201132
}
202133
}
203134

204-
func TestStdinServiceCancellation(t *testing.T) {
135+
func TestStdinService_Cancelation(t *testing.T) {
205136
s, err := Factory(&fspb.ClientServiceConfig{
206137
Name: "SleepService",
207138
Config: anypbtest.New(t, &sspb.Config{

0 commit comments

Comments
 (0)