Upgrade to Go Mods and prevent leaks closing tunnels (#2)

Author: Enrique

go.mod

@@ -0,0 +1,5 @@
+module github.com/elliotchance/sshtunnel
+
+go 1.13
+
+require golang.org/x/crypto v0.0.0-20191122220453-ac88ee75c92c

go.sum

@@ -0,0 +1,8 @@
+golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
+golang.org/x/crypto v0.0.0-20191122220453-ac88ee75c92c h1:/nJuwDLoL/zrqY6gf57vxC+Pi+pZ8bfhpPkicO5H7W4=
+golang.org/x/crypto v0.0.0-20191122220453-ac88ee75c92c/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
+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-20190412213103-97732733099d h1:+R4KGOnez64A81RvjARKc4UT5/tI9ujCIVX+P5KiHuI=
+golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=

ssh_tunnel.go

@@ -5,6 +5,7 @@ import (
 	"io"
 	"log"
 	"net"
+	"sync"
 )
 
 type SSHTunnel struct {
@@ -13,6 +14,7 @@ type SSHTunnel struct {
 	Remote *Endpoint
 	Config *ssh.ClientConfig
 	Log    *log.Logger
+	close  chan interface{}
 }
 
 func (tunnel *SSHTunnel) logf(fmt string, args ...interface{}) {
@@ -26,47 +28,59 @@ func (tunnel *SSHTunnel) Start() error {
 	if err != nil {
 		return err
 	}
-	defer listener.Close()
-
 	tunnel.Local.Port = listener.Addr().(*net.TCPAddr).Port
-
 	for {
 		conn, err := listener.Accept()
 		if err != nil {
 			return err
 		}
-
 		tunnel.logf("accepted connection")
-		go tunnel.forward(conn)
+		var wg sync.WaitGroup
+		go tunnel.forward(conn, &wg)
+		wg.Wait()
+		tunnel.logf("tunnel closed")
+		break
 	}
+	err = listener.Close()
+	if err != nil {
+		return err
+	}
+	return nil
 }
 
-func (tunnel *SSHTunnel) forward(localConn net.Conn) {
+func (tunnel *SSHTunnel) forward(localConn net.Conn, wg *sync.WaitGroup) {
 	serverConn, err := ssh.Dial("tcp", tunnel.Server.String(), tunnel.Config)
 	if err != nil {
 		tunnel.logf("server dial error: %s", err)
 		return
 	}
-
 	tunnel.logf("connected to %s (1 of 2)\n", tunnel.Server.String())
-
 	remoteConn, err := serverConn.Dial("tcp", tunnel.Remote.String())
 	if err != nil {
 		tunnel.logf("remote dial error: %s", err)
 		return
 	}
-
 	tunnel.logf("connected to %s (2 of 2)\n", tunnel.Remote.String())
-
 	copyConn := func(writer, reader net.Conn) {
 		_, err := io.Copy(writer, reader)
 		if err != nil {
 			tunnel.logf("io.Copy error: %s", err)
 		}
 	}
-
 	go copyConn(localConn, remoteConn)
 	go copyConn(remoteConn, localConn)
+	<-tunnel.close
+	tunnel.logf("close signal received, closing...")
+	_ = localConn.Close()
+	_ = serverConn.Close()
+	_ = remoteConn.Close()
+	wg.Done()
+	return
+}
+
+func (tunnel *SSHTunnel) Close() {
+	tunnel.close <- struct{}{}
+	return
 }
 
 func NewSSHTunnel(tunnel string, auth ssh.AuthMethod, destination string) *SSHTunnel {
@@ -90,6 +104,7 @@ func NewSSHTunnel(tunnel string, auth ssh.AuthMethod, destination string) *SSHTu
 		Local:  localEndpoint,
 		Server: server,
 		Remote: NewEndpoint(destination),
+		close:  make(chan interface{}),
 	}
 
 	return sshTunnel