Merge branch 'antonsergeyev-context-support'

README.md

@@ -38,6 +38,7 @@ You can find the docs at [go docs](https://pkg.go.dev/github.com/melbahja/goph).
 - Supports connections with **ssh agent** (Unix systems only).
 - Supports adding new hosts to **known_hosts file**.
 - Supports **file system operations** like: `Open, Create, Chmod...`
+- Supports **context.Context** for command cancellation.
 
 ## 📄  Usage
 
@@ -119,6 +120,14 @@ err := client.Download("/path/to/remote/file", "/path/to/local/file")
 out, err := client.Run("bash -c 'printenv'")
 ```
 
+#### ☛ Execute Bash Command with timeout:
+```go
+context, cancel := context.WithTimeout(ctx, time.Second)
+defer cancel()
+// will send SIGINT and return error after 1 second
+out, err := client.RunContext(ctx, "sleep 5")
+```
+
 #### ☛ Execute Bash Command With Env Variables:
 ```go
 out, err := client.Run(`env MYVAR="MY VALUE" bash -c 'echo $MYVAR;'`)
@@ -132,6 +141,9 @@ out, err := client.Run(`env MYVAR="MY VALUE" bash -c 'echo $MYVAR;'`)
 // Get new `Goph.Cmd`
 cmd, err := client.Command("ls", "-alh", "/tmp")
 
+// or with context:
+// cmd, err := client.CommandContext(ctx, "ls", "-alh", "/tmp")
+
 if err != nil {
  // handle the error!
 }

client.go

@@ -4,6 +4,7 @@
 package goph
 
 import (
+ "context"
  "fmt"
  "io"
  "net"
@@ -106,6 +107,16 @@ func (c Client) Run(cmd string) ([]byte, error) {
  return sess.CombinedOutput(cmd)
 }
 
+// Run starts a new SSH session with context and runs the cmd. It returns CombinedOutput and err if any.
+func (c Client) RunContext(ctx context.Context, name string) ([]byte, error) {
+ cmd, err := c.CommandContext(ctx, name)
+ if err != nil {
+ return nil, err
+ }
+
+ return cmd.CombinedOutput()
+}
+
 // Command returns new Cmd and error if any.
 func (c Client) Command(name string, args ...string) (*Cmd, error) {
 
@@ -122,9 +133,22 @@ func (c Client) Command(name string, args ...string) (*Cmd, error) {
  Path: name,
  Args: args,
  Session: sess,
+ Context: context.Background(),
  }, nil
 }
 
+// Command returns new Cmd with context and error, if any.
+func (c Client) CommandContext(ctx context.Context, name string, args ...string) (*Cmd, error) {
+ cmd, err := c.Command(name, args...)
+ if err != nil {
+ return cmd, err
+ }
+
+ cmd.Context = ctx
+
+ return cmd, nil
+}
+
 // NewSftp returns new sftp client and error if any.
 func (c Client) NewSftp(opts ...sftp.ClientOption) (*sftp.Client, error) {
  return sftp.NewClient(c.Client, opts...)

cmd.go

@@ -4,11 +4,11 @@
 package goph
 
 import (
+ "context"
  "fmt"
- "strings"
-
  "github.com/pkg/errors"
  "golang.org/x/crypto/ssh"
+ "strings"
 )
 
 // Cmd it's like os/exec.Cmd but for ssh session.
@@ -25,30 +25,44 @@ type Cmd struct {
 
  // SSH session.
  *ssh.Session
+
+ // Context for cancellation
+ Context context.Context
 }
 
 // CombinedOutput runs cmd on the remote host and returns its combined stdout and stderr.
 func (c *Cmd) CombinedOutput() ([]byte, error) {
  if err := c.init(); err != nil {
  return nil, errors.Wrap(err, "cmd init")
  }
- return c.Session.CombinedOutput(c.String())
+
+ return c.runWithContext(func() ([]byte, error) {
+ return c.Session.CombinedOutput(c.String())
+ })
 }
 
 // Output runs cmd on the remote host and returns its stdout.
 func (c *Cmd) Output() ([]byte, error) {
  if err := c.init(); err != nil {
  return nil, errors.Wrap(err, "cmd init")
  }
- return c.Session.Output(c.String())
+
+ return c.runWithContext(func() ([]byte, error) {
+ return c.Session.Output(c.String())
+ })
 }
 
 // Run runs cmd on the remote host.
 func (c *Cmd) Run() error {
  if err := c.init(); err != nil {
  return errors.Wrap(err, "cmd init")
  }
- return c.Session.Run(c.String())
+
+ _, err := c.runWithContext(func() ([]byte, error) {
+ return nil, c.Session.Run(c.String())
+ })
+
+ return err
 }
 
 // Start runs the command on the remote host.
@@ -78,3 +92,30 @@ func (c *Cmd) init() (err error) {
 
  return nil
 }
+
+// Command with context output.
+type ctxCmdOutput struct {
+ output []byte
+ err error
+}
+
+// Executes the given callback within session. Sends SIGINT when the context is canceled.
+func (c *Cmd) runWithContext(callback func() ([]byte, error)) ([]byte, error) {
+ outputChan := make(chan ctxCmdOutput)
+ go func() {
+ output, err := callback()
+ outputChan <- ctxCmdOutput{
+ output: output,
+ err: err,
+ }
+ }()
+
+ select {
+ case <-c.Context.Done():
+ _ = c.Session.Signal(ssh.SIGINT)
+
+ return nil, c.Context.Err()
+ case result := <-outputChan:
+ return result.output, result.err
+ }
+}

examples/goph/main.go

@@ -2,13 +2,15 @@ package main
 
 import (
  "bufio"
+ "context"
  "errors"
  "flag"
  "fmt"
  "log"
  "net"
  "os"
  "strings"
+ "time"
 
  "github.com/melbahja/goph"
  "github.com/pkg/sftp"
@@ -28,8 +30,10 @@ import (
 // Run command and auth with private key and passphrase:
 // > go run main.go --ip 192.168.122.102 --passphrase --cmd ls
 //
+// Run a command and interrupt it after 1 second:
+// > go run main.go --ip 192.168.122.102 --cmd "sleep 10" --timeout=1s
 //
-// You can test with the interactive mode without passing --cmd falg.
+// You can test with the interactive mode without passing --cmd flag.
 //
 
 var (
@@ -43,6 +47,7 @@ var (
  cmd string
  pass bool
  passphrase bool
+ timeout time.Duration
  agent bool
  sftpc *sftp.Client
 )
@@ -57,6 +62,7 @@ func init() {
  flag.BoolVar(&pass, "pass", false, "ask for ssh password instead of private key.")
  flag.BoolVar(&agent, "agent", false, "use ssh agent for authentication (unix systems only).")
  flag.BoolVar(&passphrase, "passphrase", false, "ask for private key passphrase.")
+ flag.DurationVar(&timeout, "timeout", 0, "interrupt a command with SIGINT after a given timeout (0 means no timeout)")
 }
 
 func VerifyHost(host string, remote net.Addr, key ssh.PublicKey) error {
@@ -135,8 +141,15 @@ func main() {
 
  // If the cmd flag exists
  if cmd != "" {
+ ctx := context.Background()
+ // create a context with timeout, if supplied in the argumetns
+ if timeout > 0 {
+ var cancel context.CancelFunc
+ ctx, cancel = context.WithTimeout(ctx, timeout)
+ defer cancel()
+ }
 
- out, err := client.Run(cmd)
+ out, err := client.RunContext(ctx, cmd)
 
  fmt.Println(string(out), err)
  return

go.mod

@@ -4,8 +4,8 @@ go 1.13
 
 require (
  github.com/pkg/errors v0.9.1
- github.com/pkg/sftp v1.12.0
- golang.org/x/crypto v0.0.0-20201217014255-9d1352758620
- golang.org/x/sys v0.0.0-20201218084310-7d0127a74742 // indirect
- golang.org/x/term v0.0.0-20201210144234-2321bbc49cbf // indirect
+ github.com/pkg/sftp v1.13.4
+ golang.org/x/crypto v0.0.0-20210921155107-089bfa567519
+ golang.org/x/sys v0.0.0-20211031064116-611d5d643895 // indirect
+ golang.org/x/term v0.0.0-20210927222741-03fcf44c2211 // indirect
 )

go.sum

@@ -6,28 +6,44 @@ github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/sftp v1.12.0 h1:/f3b24xrDhkhddlaobPe2JgBqfdt+gC/NYl0QY9IOuI=
 github.com/pkg/sftp v1.12.0/go.mod h1:fUqqXB5vEgVCZ131L+9say31RAri6aF6KDViawhxKK8=
+github.com/pkg/sftp v1.13.4 h1:Lb0RYJCmgUcBgZosfoi9Y9sbl6+LJgOIgk/2Y4YjMFg=
+github.com/pkg/sftp v1.13.4/go.mod h1:LzqnAvaD5TWeNBsZpfKxSYn1MbjWwOsCIAFFJbpIsK8=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.6.1 h1:hDPOHmpOpP40lSULcqw7IrRb/u7w6RpDC9399XyoNd0=
 github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20200820211705-5c72a883971a/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/crypto v0.0.0-20201217014255-9d1352758620 h1:3wPMTskHO3+O6jqTEXyFcsnuxMQOqYSaHsDxcbUXpqA=
 golang.org/x/crypto v0.0.0-20201217014255-9d1352758620/go.mod h1:jdWPYTVW3xRLrWPugEBEK3UY2ZEsg3UU495nc5E+M+I=
+golang.org/x/crypto v0.0.0-20210421170649-83a5a9bb288b/go.mod h1:T9bdIzuCu7OtxOm1hfPfRQxPLYneinmdGuTeoZ9dtd4=
+golang.org/x/crypto v0.0.0-20210921155107-089bfa567519 h1:7I4JAnoQBe7ZtJcBaYHi5UtiO8tQHbUSXxL+pnGRANg=
+golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
 golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
+golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d h1:+R4KGOnez64A81RvjARKc4UT5/tI9ujCIVX+P5KiHuI=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20201218084310-7d0127a74742 h1:+CBz4km/0KPU3RGTwARGh/noP3bEwtHcq+0YcBQM2JQ=
 golang.org/x/sys v0.0.0-20201218084310-7d0127a74742/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20210423185535-09eb48e85fd7/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20211031064116-611d5d643895 h1:iaNpwpnrgL5jzWS0vCNnfa8HqzxveCFpFx3uC/X4Tps=
+golang.org/x/sys v0.0.0-20211031064116-611d5d643895/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201117132131-f5c789dd3221 h1:/ZHdbVpdR/jk3g30/d4yUL0JU9kksj8+F/bnQUVLGDM=
 golang.org/x/term v0.0.0-20201117132131-f5c789dd3221/go.mod h1:Nr5EML6q2oocZ2LXRh80K7BxOlk5/8JxuGnuhpl+muw=
+golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/term v0.0.0-20201210144234-2321bbc49cbf h1:MZ2shdL+ZM/XzY3ZGOnh4Nlpnxz5GSOhOmtHo3iPU6M=
 golang.org/x/term v0.0.0-20201210144234-2321bbc49cbf/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
+golang.org/x/term v0.0.0-20210927222741-03fcf44c2211 h1:JGgROgKl9N8DuW20oFS5gxc+lE67/N3FcwmBPMe7ArY=
+golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
+golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
+golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=