Skip to content

Commit

Permalink
stdinservice: Simplify command execution
Browse files Browse the repository at this point in the history
`exec.CommandContext()` respects context cancelation already, and that lets us avoid additional goroutines and channels.

Better to simplify this now before we grow dependencies on behaviour that is
out-of-line with the standard library.

PiperOrigin-RevId: 704272182
  • Loading branch information
gnoack authored and copybara-github committed Dec 9, 2024
1 parent d36785f commit e6e7df3
Showing 1 changed file with 6 additions and 35 deletions.
41 changes: 6 additions & 35 deletions fleetspeak/src/client/stdinservice/stdinservice.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@ import (
"context"
"fmt"
"os/exec"
"strings"

log "github.com/golang/glog"
anypb "google.golang.org/protobuf/types/known/anypb"
Expand Down Expand Up @@ -64,12 +63,14 @@ type StdinService struct {
sc service.Context
}

// Start starts the StdinService.
func (s *StdinService) Start(sc service.Context) error {
s.sc = sc

return nil
}

// ProcessMessage processes an incoming message from the server side.
func (s *StdinService) ProcessMessage(ctx context.Context, m *fspb.Message) error {
om := &sspb.OutputMessage{}

Expand All @@ -82,7 +83,7 @@ func (s *StdinService) ProcessMessage(ctx context.Context, m *fspb.Message) erro
return fmt.Errorf("error while unmarshalling common.Message.data: %v", err)
}

cmd := exec.Command(s.ssConf.Cmd, im.Args...)
cmd := exec.CommandContext(ctx, s.ssConf.Cmd, im.Args...)

inBuf := bytes.NewBuffer(im.Input)
var outBuf, errBuf bytes.Buffer
Expand All @@ -101,40 +102,9 @@ func (s *StdinService) ProcessMessage(ctx context.Context, m *fspb.Message) erro
log.Errorf("Failed to get resource usage for process: %v", ruErr)
}

waitChan := make(chan error)
go func() {
// Note that using cmd.Run() here triggers panics.
waitChan <- cmd.Wait()
}()

var err error
select {
case err = <-waitChan:
case <-ctx.Done():
err = fmt.Errorf("context done: %v", ctx.Err())

// The error message string literal is a copypaste from exec_unix.go .
if e := cmd.Process.Kill(); e != nil && e.Error() != "os: process already finished" {
err = fmt.Errorf("%v; also, an error occurred while killing the process: %v", err, e)
}

e, ok := <-waitChan

const (
killedMessage = "signal: killed"
killedMessageWin = "exit status " // + number representing the return code.
)

if !ok {
err = fmt.Errorf("%v; also, .Wait hasn't returned after killing the process", err)
} else if e != nil &&
e.Error() != killedMessage &&
!strings.HasPrefix(e.Error(), killedMessageWin) {
err = fmt.Errorf("%v; also, .Wait returned an error: %v", err, e)
}
}
err := cmd.Wait()
if err != nil {
return fmt.Errorf("error while running a process: %v", err)
return err
}

om.Stdout = outBuf.Bytes()
Expand All @@ -156,6 +126,7 @@ func (s *StdinService) ProcessMessage(ctx context.Context, m *fspb.Message) erro
return nil
}

// Stop stops the StdinService.
func (s *StdinService) Stop() error {
return nil
}
Expand Down

0 comments on commit e6e7df3

Please sign in to comment.