Free ports on shutdown (#1140)

* disconnect the port forwards when the context is cancelled
* make SSH connection context-aware

Signed-off-by: Ramiro Berrelleza <rberrelleza@gmail.com>

pkg/ssh/forward.go

@@ -59,14 +59,18 @@ func (f *forward) start() {
 	f.setConnected()
 
 	for {
-		localConn, err := localListener.Accept()
-		if err != nil {
-			log.Infof("%s -> failed to accept connection: %v", f.String(), err)
-			continue
+		select {
+		case <-f.ctx.Done():
+			return
+		default:
+			localConn, err := localListener.Accept()
+			if err != nil {
+				log.Infof("%s -> failed to accept connection: %v", f.String(), err)
+				continue
+			}
+
+			go f.handle(localConn)
 		}
-
-		go f.handle(localConn)
-
 	}
 }
 

pkg/ssh/manager.go

@@ -117,6 +117,7 @@ func (fm *ForwardManager) Start(devPod, namespace string) error {
 	for _, ff := range fm.forwards {
 		ff.pool = pool
 		go ff.start()
+
 	}
 
 	for _, rt := range fm.reverses {
@@ -129,7 +130,6 @@ func (fm *ForwardManager) Start(devPod, namespace string) error {
 
 // Stop sends a stop signal to all the connections
 func (fm *ForwardManager) Stop() {
-	// TODO stop forwards and reverses
 
 	if fm.pf != nil {
 		fm.pf.Stop()

pkg/ssh/pool.go

@@ -26,20 +26,16 @@ import (
 type pool struct {
 	ka     time.Duration
 	ctx    context.Context
-	cancel func()
-
 	client *ssh.Client
 }
 
 func startPool(ctx context.Context, serverAddr string, config *ssh.ClientConfig) (*pool, error) {
-	c, cancel := context.WithCancel(ctx)
 	p := &pool{
-		ka:     30 * time.Second,
-		ctx:    c,
-		cancel: cancel,
+		ka:  30 * time.Second,
+		ctx: ctx,
 	}
 
-	conn, err := getTCPConnection(serverAddr, p.ka)
+	conn, err := getTCPConnection(ctx, serverAddr, p.ka)
 	if err != nil {
 		return nil, fmt.Errorf("failed to establish a tcp connection for %s: %s", serverAddr, err)
 	}
@@ -90,8 +86,8 @@ func (p *pool) getListener(address string) (net.Listener, error) {
 	return l, nil
 }
 
-func getTCPConnection(serverAddr string, keepAlive time.Duration) (net.Conn, error) {
-	c, err := getConn(serverAddr, 3)
+func getTCPConnection(ctx context.Context, serverAddr string, keepAlive time.Duration) (net.Conn, error) {
+	c, err := getConn(ctx, serverAddr, 3)
 	if err != nil {
 		return nil, err
 	}
@@ -107,11 +103,12 @@ func getTCPConnection(serverAddr string, keepAlive time.Duration) (net.Conn, err
 	return c, nil
 }
 
-func getConn(serverAddr string, maxRetries int) (net.Conn, error) {
+func getConn(ctx context.Context, serverAddr string, maxRetries int) (net.Conn, error) {
 	var lastErr error
 	t := time.NewTicker(100 * time.Millisecond)
 	for i := 0; i < 3; i++ {
-		c, err := net.Dial("tcp", serverAddr)
+		d := net.Dialer{}
+		c, err := d.DialContext(ctx, "tcp", serverAddr)
 		if err == nil {
 			return c, nil
 		}

pkg/ssh/reverse.go

@@ -54,25 +54,27 @@ func (r *reverse) start() {
 	defer remoteListener.Close()
 
 	for {
-
-		r.setConnected()
-
-		remoteConn, err := remoteListener.Accept()
-		if err != nil {
-			log.Infof("%s -> failed to accept connection: %v", r.String(), err)
-			continue
+		select {
+		case <-r.ctx.Done():
+			return
+		default:
+			r.setConnected()
+			remoteConn, err := remoteListener.Accept()
+			if err != nil {
+				log.Infof("%s -> failed to accept connection: %v", r.String(), err)
+				continue
+			}
+
+			go r.handle(remoteConn)
 		}
-
-		go r.handle(remoteConn)
-
 	}
 }
 
 func (r *reverse) handle(remote net.Conn) {
 	defer remote.Close()
 
 	quit := make(chan struct{}, 1)
-	local, err := getConn(r.localAddress, 3)
+	local, err := getConn(r.ctx, r.localAddress, 3)
 	if err != nil {
 		log.Infof("%s -> failed to listen on local address: %v", r.String(), err)
 		return