fix stdio "cannot listen", fix udp goroutine leak

Author: jpillora

client/client.go

@@ -138,8 +138,8 @@ func NewClient(c *Config) (*Client, error) {
 			hasStdio = true
 		}
 		//confirm non-reverse tunnel is available
-		if !r.Reverse && !r.CanListen() {
-			return nil, fmt.Errorf("Remote %s cannot listen", r.String())
+		if !r.Reverse && !r.Stdio && !r.CanListen() {
+			return nil, fmt.Errorf("Client cannot listen on %s", r.String())
 		}
 		client.computed.Remotes = append(client.computed.Remotes, r)
 	}

client/client_connect.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	"net"
 	"strings"
 	"time"
@@ -29,8 +30,8 @@ func (c *Client) connectionLoop(ctx context.Context) error {
 		//connection error
 		attempt := int(b.Attempt())
 		maxAttempt := c.config.MaxRetryCount
-		if err != nil {
-			//show error and attempt counts
+		//show error message and attempt counts (excluding disconnects)
+		if err != nil && err != io.EOF && !strings.HasSuffix(err.Error(), "use of closed network connection") {
 			msg := fmt.Sprintf("Connection error: %s", err)
 			if attempt > 0 {
 				msg += fmt.Sprintf(" (Attempt: %d", attempt)

server/server_handler.go

@@ -128,7 +128,7 @@ func (s *Server) handleWebsocket(w http.ResponseWriter, req *http.Request) {
 		}
 		//confirm reverse tunnel is available
 		if r.Reverse && !r.CanListen() {
-			failed(s.Errorf("Remote %s cannot listen", r.String()))
+			failed(s.Errorf("Server cannot listen on %s", r.String()))
 			return
 		}
 	}

share/tunnel/tunnel_out_ssh.go

@@ -51,7 +51,8 @@ func (t *Tunnel) handleSSHChannel(ch ssh.NewChannel) {
 		t.Debugf("Failed to accept stream: %s", err)
 		return
 	}
-	stream := io.ReadWriteCloser(sshChan) //cnet.MeterRWC(t.Logger.Fork("sshchan"), sshChan)
+	stream := io.ReadWriteCloser(sshChan)
+	//cnet.MeterRWC(t.Logger.Fork("sshchan"), sshChan)
 	defer stream.Close()
 	go ssh.DiscardRequests(reqs)
 	l := t.Logger.Fork("conn#%d", t.connStats.New())

share/tunnel/tunnel_out_ssh_udp.go

@@ -12,6 +12,11 @@ import (
 )
 
 func (t *Tunnel) handleUDP(l *cio.Logger, rwc io.ReadWriteCloser, hostPort string) error {
+	conns := &udpConns{
+		Logger: l,
+		m:      map[string]*udpConn{},
+	}
+	defer conns.closeAll()
 	h := &udpHandler{
 		Logger:   l,
 		hostPort: hostPort,
@@ -20,10 +25,7 @@ func (t *Tunnel) handleUDP(l *cio.Logger, rwc io.ReadWriteCloser, hostPort strin
 			w: gob.NewEncoder(rwc),
 			c: rwc,
 		},
-		udpConns: &udpConns{
-			Logger: l,
-			m:      map[string]*udpConn{},
-		},
+		udpConns: conns,
 	}
 	for {
 		p := udpPacket{}
@@ -136,6 +138,15 @@ func (cs *udpConns) remove(id string) {
 	cs.Unlock()
 }
 
+func (cs *udpConns) closeAll() {
+	cs.Lock()
+	for id, conn := range cs.m {
+		conn.Close()
+		delete(cs.m, id)
+	}
+	cs.Unlock()
+}
+
 type udpConn struct {
 	id string
 	net.Conn

test/e2e/setup_test.go

@@ -6,6 +6,9 @@ import (
 	"log"
 	"net"
 	"net/http"
+	"os"
+	"runtime"
+	"runtime/pprof"
 	"strings"
 	"testing"
 	"time"
@@ -27,7 +30,7 @@ type testLayout struct {
 
 func (tl *testLayout) setup(t *testing.T) (server *chserver.Server, client *chclient.Client, teardown func()) {
 	//start of the world
-	// goroutines := runtime.NumGoroutine()
+	goroutines := runtime.NumGoroutine()
 	//root cancel
 	ctx, cancel := context.WithCancel(context.Background())
 	//fileserver (fake endpoint)
@@ -97,13 +100,13 @@ func (tl *testLayout) setup(t *testing.T) (server *chserver.Server, client *chcl
 		server.Wait()
 		client.Wait()
 		//confirm goroutines have been cleaned up
-		// time.Sleep(1 * time.Second)
-		//TODO remove sleep
-		// d := runtime.NumGoroutine() - goroutines
-		// if d != 0 {
-		// 	pprof.Lookup("goroutine").WriteTo(os.Stdout, 1)
-		// 	t.Fatalf("goroutines left %d", d)
-		// }
+		time.Sleep(500 * time.Millisecond)
+		// TODO remove sleep
+		d := runtime.NumGoroutine() - goroutines
+		if d != 0 {
+			pprof.Lookup("goroutine").WriteTo(os.Stdout, 1)
+			t.Fatalf("goroutines left %d", d)
+		}
 	}
 	//wait a bit...
 	//TODO: client signal API, similar to os.Notify(signal)