Reimplemented approach to blocking reader (#12)

* initial implementation of runUntilPrompt goroutine

* removed stale comment

* remove extra loop

* fixed readertimeout test

* fixed linter errors

* removed note about foked crypto/ssh

* modified reader

* added WithDialTimeout

README.md

@@ -84,20 +84,3 @@ Credits
 Netrasp was created by [Patrick Ogenstad](https://github.com/ogenstad). Special
 thanks to [David Barroso](https://github.com/dbarrosop) for providing feedback
 and recommendations of the structure and code.
-
-SSH Crypto package
-------------------
-
-When creating Netrasp, I encountered an issue with Golang’s crypto package for
-SSH. The problem is that the Read call when retrieving information from a device
-is blocking, so it wasn’t possible to cancel that read using a Go context. As
-I’m fairly new to Go, it could be just that I’m not aware of the correct way
-to make the read cancelable without starting any goroutines that might remain
-in the background. I only needed one line to be changed within the crypto/ssh
-library, so I made a [fork](https://github.com/ogenstad/crypto) (referenced in
-the go.mod file for Netrasp). In the file
-https://github.com/golang/crypto/blob/master/ssh/buffer.go, I needed to change
-the Read() method. Instead of hanging at `b.Cond.Wait()`, I added `return 0, nil`.
-Edit your local go.mod file if you want a similar replacement to avoid blocking
-reads:
-`replace golang.org/x/crypto => github.com/ogenstad/crypto v0.0.0-20210308070823-6d211c1ce3d7`

go.mod

@@ -6,5 +6,3 @@ require (
 	github.com/magefile/mage v1.11.0
 	golang.org/x/crypto v0.0.0-20201203163018-be400aefbc4c
 )
-
-replace golang.org/x/crypto => github.com/ogenstad/crypto v0.0.0-20210308070823-6d211c1ce3d7

go.sum

@@ -1,8 +1,10 @@
 github.com/magefile/mage v1.11.0 h1:C/55Ywp9BpgVVclD3lRnSYCwXTYxmSppIgLeDYlNuls=
 github.com/magefile/mage v1.11.0/go.mod h1:z5UZb/iS3GoOSn0JgWuiw7dxlurVYTu+/jHXqQg881A=
-github.com/ogenstad/crypto v0.0.0-20210308070823-6d211c1ce3d7 h1:eFSQKFBE3IO3CUDnc6mNnQEO/Ky2HAOShDCUpJPQhq0=
-github.com/ogenstad/crypto v0.0.0-20210308070823-6d211c1ce3d7/go.mod h1:jdWPYTVW3xRLrWPugEBEK3UY2ZEsg3UU495nc5E+M+I=
+golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
+golang.org/x/crypto v0.0.0-20201203163018-be400aefbc4c h1:9HhBz5L/UjnK9XLtiZhYAdue5BVKep3PMmS2LuPDt8k=
+golang.org/x/crypto v0.0.0-20201203163018-be400aefbc4c/go.mod h1:jdWPYTVW3xRLrWPugEBEK3UY2ZEsg3UU495nc5E+M+I=
 golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
+golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20191026070338-33540a1f6037 h1:YyJpGZS1sBuBCzLAR1VEpK193GlqGZbnPFnPV/5Rsb4=
 golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/term v0.0.0-20201117132131-f5c789dd3221 h1:/ZHdbVpdR/jk3g30/d4yUL0JU9kksj8+F/bnQUVLGDM=

pkg/netrasp/netrasp.go

@@ -7,6 +7,7 @@ package netrasp
 import (
 	"errors"
 	"fmt"
+	"time"
 
 	"golang.org/x/crypto/ssh"
 )
@@ -135,3 +136,10 @@ func WithSSHKeyExchange(name string) ConfigOpt {
 		c.SSHConfig.KeyExchanges = append(c.SSHConfig.KeyExchanges, name)
 	})
 }
+
+// WithDialTimeout allows you to configure timeout for dialing SSH server
+func WithDialTimeout(t time.Duration) ConfigOpt {
+	return newFuncConfigOpt(func(c *config) {
+		c.SSHConfig.Timeout = t
+	})
+}

pkg/netrasp/reader.go

@@ -2,27 +2,51 @@ package netrasp
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"regexp"
 	"strings"
-	"time"
 )
 
+var errRead = errors.New("reader error")
+
 type contextReader struct {
 	ctx context.Context
 	r   io.Reader
 }
 
+type readResult struct {
+	n   int
+	err error
+}
+
 func (c *contextReader) Read(p []byte) (n int, err error) {
+	ctx, cancel := context.WithCancel(c.ctx)
+	defer cancel()
+	rrCh := make(chan *readResult)
+
+	go func() {
+		select {
+		case rrCh <- c.read(p):
+		case <-ctx.Done():
+		}
+	}()
+
 	select {
-	case <-c.ctx.Done():
-		return 0, c.ctx.Err()
-	default:
-		return c.r.Read(p)
+	case <-ctx.Done():
+		return 0, ctx.Err()
+	case rr := <-rrCh:
+		return rr.n, rr.err
 	}
 }
 
+func (c *contextReader) read(p []byte) *readResult {
+	n, err := c.r.Read(p)
+
+	return &readResult{n, err}
+}
+
 func newContextReader(ctx context.Context, r io.Reader) io.Reader {
 	return &contextReader{
 		ctx: ctx,
@@ -33,14 +57,13 @@ func newContextReader(ctx context.Context, r io.Reader) io.Reader {
 // readUntilPrompt reads until the specified prompt is found and returns the read data.
 func readUntilPrompt(ctx context.Context, r io.Reader, prompt *regexp.Regexp) (string, error) {
 	var output string
-	reader := newContextReader(ctx, r)
-
+	r = newContextReader(ctx, r)
 	for {
 		buffer := make([]byte, 10000)
-		time.Sleep(time.Millisecond * 10)
-		bytes, err := reader.Read(buffer)
+
+		bytes, err := r.Read(buffer)
 		if err != nil {
-			return "", fmt.Errorf("error reading output from device: %w", err)
+			return "", fmt.Errorf("error reading output from device %w: %v", errRead, err)
 		}
 		latestOutput := string(buffer[:bytes])
 

pkg/netrasp/reader_test.go

@@ -26,10 +26,10 @@ func TestReader(t *testing.T) {
 func TestReaderTimeout(t *testing.T) {
 	r := strings.NewReader(testReaderBasic)
 
-	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Millisecond)
+	ctx, cancel := context.WithTimeout(context.Background(), 0*time.Second)
 	defer cancel()
-	_, err := readUntilPrompt(ctx, r, generalPrompt)
+	out, err := readUntilPrompt(ctx, r, generalPrompt)
 	if err == nil {
-		t.Fatalf("expected read to fail")
+		t.Fatalf("expected reader to timeout, but got %v", out)
 	}
 }