Merge pull request pkg#279 from wutzx/client-conn-wait

eikenb · web-flow · commit a713b07e6d90 · 2019-02-05T12:56:04.000-08:00
Add Wait method to detect underlying SFTP connection closed
diff --git a/client.go b/client.go
@@ -122,6 +122,7 @@ func NewClientPipe(rd io.Reader, wr io.WriteCloser, opts ...ClientOption) (*Clie
 				WriteCloser: wr,
 			},
 			inflight: make(map[uint32]chan<- result),
+			closed:   make(chan struct{}),
 		},
 		maxPacket:             1 << 15,
 		maxConcurrentRequests: 64,
diff --git a/conn.go b/conn.go
@@ -36,6 +36,17 @@ type clientConn struct {
 	wg         sync.WaitGroup
 	sync.Mutex                          // protects inflight
 	inflight   map[uint32]chan<- result // outstanding requests
+
+	closed chan struct{}
+	err    error
+}
+
+// Wait blocks until the conn has shut down, and return the error
+// causing the shutdown. It can be called concurrently from multiple
+// goroutines.
+func (c *clientConn) Wait() error {
+	<-c.closed
+	return c.err
 }
 
 // Close closes the SFTP session.
@@ -122,6 +133,8 @@ func (c *clientConn) broadcastErr(err error) {
 	for _, ch := range listeners {
 		ch <- result{err: err}
 	}
+	c.err = err
+	close(c.closed)
 }
 
 type serverConn struct {
