grim/hgkeeper

Move the command running to a helper function

2019-05-03, Gary Kramlich
7d7cae4b3989
Parents f7953a3da60c
Children a98b0f17942a
Move the command running to a helper function
2 files changed, 62 insertions(+), 42 deletions(-)
  • +59 -0
    ssh/command.go
  • +3 -42
    ssh/server.go
    • --- a/ssh/command.go Fri May 03 07:19:45 2019 -0500
    +++ b/ssh/command.go Fri May 03 08:51:59 2019 -0500
    @@ -2,11 +2,15 @@
    import (
    "fmt"
    + "io"
    "os/exec"
    "path/filepath"
    + "sync"
    "github.com/alecthomas/kong"
    "github.com/kballard/go-shellquote"
    + log "github.com/sirupsen/logrus"
    + "golang.org/x/crypto/ssh"
    "bitbucket.org/rw_grim/hgkeeper/hg"
    )
    @@ -35,6 +39,61 @@
    }
    }
    +func (c *Command) run(conn ssh.Channel, req *ssh.Request) error {
    + teardown := func() {
    + conn.Close()
    + if err := c.cmd.Wait(); err != nil {
    + log.Warnf("failed running command : %v", err)
    + req.Reply(false, nil)
    + return
    + }
    +
    + log.Info("finished command")
    + req.Reply(true, nil)
    + }
    +
    + stdinWriter, err := c.cmd.StdinPipe()
    + if err != nil {
    + return err
    + }
    +
    + stdoutReader, err := c.cmd.StdoutPipe()
    + if err != nil {
    + return err
    + }
    +
    + stderrReader, err := c.cmd.StderrPipe()
    + if err != nil {
    + return err
    + }
    +
    + var once sync.Once
    +
    + // now wire up stdin/stdout/stderr
    + go func() {
    + io.Copy(stdinWriter, conn)
    + once.Do(teardown)
    + }()
    +
    + go func() {
    + io.Copy(conn, stdoutReader)
    + once.Do(teardown)
    + }()
    +
    + go func() {
    + io.Copy(conn.Stderr(), stderrReader)
    + once.Do(teardown)
    + }()
    +
    + if err := c.cmd.Start(); err != nil {
    + req.Reply(false, nil)
    + once.Do(teardown)
    + return err
    + }
    +
    + return nil
    +}
    +
    func parseCommand(cmd string) (cli, string, error) {
    values := cli{}
    --- a/ssh/server.go Fri May 03 07:19:45 2019 -0500
    +++ b/ssh/server.go Fri May 03 08:51:59 2019 -0500
    @@ -2,9 +2,7 @@
    import (
    "fmt"
    - "io"
    "net"
    - "sync"
    log "github.com/sirupsen/logrus"
    "golang.org/x/crypto/ssh"
    @@ -120,48 +118,11 @@
    req.Reply(false, nil)
    }
    - teardown := func() {
    - conn.Close()
    - if err := cmd.cmd.Wait(); err != nil {
    - log.Warnf("failed running command : %v", err)
    - req.Reply(false, nil)
    - return
    - }
    -
    - log.Info("finished command")
    - req.Reply(true, nil)
    + if err := cmd.run(conn, req); err != nil {
    + log.Warnf("failed to run command: %v", err)
    + req.Reply(false, nil)
    }
    - stdinWriter, _ := cmd.cmd.StdinPipe()
    - stdoutReader, _ := cmd.cmd.StdoutPipe()
    - stderrReader, _ := cmd.cmd.StderrPipe()
    -
    - var once sync.Once
    -
    - // now wire up stdin/stdout/stderr
    - go func() {
    - io.Copy(stdinWriter, conn)
    - once.Do(teardown)
    - }()
    -
    - go func() {
    - io.Copy(conn, stdoutReader)
    - once.Do(teardown)
    - }()
    -
    - go func() {
    - io.Copy(conn.Stderr(), stderrReader)
    - once.Do(teardown)
    - }()
    -
    - log.Warnf("running %#v", cmd.cmd.Args)
    - if err := cmd.cmd.Start(); err != nil {
    - log.Warnf("failed to start command: %v", err)
    - req.Reply(false, nil)
    - continue
    - }
    - log.Info("command started")
    -
    default:
    log.Debugf("unsupported request: %s", req.Type)
    if req.WantReply {