fix: always relisten local address when ssh connect closed

ssh-kit · Jul 6, 2022 · ce5f9b2 · ce5f9b2
1 parent 978199b
commit ce5f9b2
Showing 1 changed file with 46 additions and 35 deletions.
diff --git a/ssh/ssh.go b/ssh/ssh.go
@@ -290,52 +290,63 @@ func (c *Client) run(ctx context.Context, conn *ssh.Client) {
 				break
 			}
 		} else {
-			for {
-				listen, err := net.Listen("tcp", rule.Local)
-				if err != nil {
-					tempDelay = c.getCurrentTempDelay(tempDelay)
+			go c.listenLocal(ctx, conn, rule)
+		}
+	}
+}
 
-					c.Logger.Error(err, "listen",
-						"status", "retry",
-						"host", c.Host,
-						"user", c.User,
-						"local", rule.Local,
-						"reverse", rule.Reverse,
-						"retry_in", tempDelay,
-					)
+func (c Client) listenLocal(ctx context.Context, conn *ssh.Client, rule Rule) {
+	var tempDelay time.Duration // how long to sleep on accept failure
+ListenLocal:
+	for {
+		listen, err := net.Listen("tcp", rule.Local)
+		if err != nil {
+			tempDelay = c.getCurrentTempDelay(tempDelay)
 
-					select {
-					case <-ctx.Done():
-						c.Logger.V(2).Info("listen",
-							"status", "canceled",
-							"host", c.Host,
-							"user", c.User,
-							"local", rule.Local,
-							"reverse", rule.Reverse,
-						)
-						return
-					case <-time.After(tempDelay):
-						continue
-					}
-				}
-				tempDelay = 0
+			c.Logger.Error(err, "listen",
+				"status", "retry",
+				"host", c.Host,
+				"user", c.User,
+				"local", rule.Local,
+				"reverse", rule.Reverse,
+				"retry_in", tempDelay,
+			)
 
-				c.Logger.V(1).Info("listen",
-					"status", "ok",
+			select {
+			case <-ctx.Done():
+				c.Logger.V(2).Info("listen",
+					"status", "canceled",
 					"host", c.Host,
 					"user", c.User,
 					"local", rule.Local,
 					"reverse", rule.Reverse,
 				)
-				// accept message and proxy
-				go c.proxy(ctx, listen, rule.Remote, rule, func(ctx context.Context, network, address string) (net.Conn, error) {
-					return conn.Dial(network, address)
-				})
-				break
+				return
+			case <-time.After(tempDelay):
+				continue
 			}
 		}
-	}
+		tempDelay = 0
 
+		c.Logger.V(1).Info("listen",
+			"status", "ok",
+			"host", c.Host,
+			"user", c.User,
+			"local", rule.Local,
+			"reverse", rule.Reverse,
+		)
+
+		// accept message and proxy
+		go c.proxy(ctx, listen, rule.Remote, rule, func(ctx context.Context, network, address string) (net.Conn, error) {
+			return conn.Dial(network, address)
+		})
+
+		select {
+		case <-ctx.Done():
+			listen.Close()
+			break ListenLocal
+		}
+	}
 }
 
 func (c *Client) proxy(ctx context.Context, source net.Listener, destination string, rule Rule, dialFunc func(ctx context.Context, network, address string) (net.Conn, error)) {