From ba347f927d34ad25d709161bfeffe24e78cca3f7 Mon Sep 17 00:00:00 2001
From: David Solin <solind@Gallifrey.local>
Date: Mon, 15 Aug 2016 00:17:48 -0500
Subject: [PATCH] Close any LocalPortForwarders with the SSHClient that
 produced them (assuming they're still open).

---
 src/main/java/net/schmizz/sshj/SSHClient.java | 18 +++++++++-
 .../channel/direct/LocalPortForwarder.java    | 33 ++++++++++++++++---
 2 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/src/main/java/net/schmizz/sshj/SSHClient.java b/src/main/java/net/schmizz/sshj/SSHClient.java
index b391a866..a24cf4ae 100644
--- a/src/main/java/net/schmizz/sshj/SSHClient.java
+++ b/src/main/java/net/schmizz/sshj/SSHClient.java
@@ -139,6 +139,8 @@ public class SSHClient
     /** {@code ssh-connection} service */
     protected final Connection conn;
 
+    private List<LocalPortForwarder> forwarders;
+
     /** Default constructor. Initializes this object using {@link DefaultConfig}. */
     public SSHClient() {
         this(new DefaultConfig());
@@ -431,6 +433,15 @@ public class SSHClient
     @Override
     public void disconnect()
             throws IOException {
+        if (forwarders != null) {
+            for (LocalPortForwarder forwarder : forwarders) {
+                try {
+                    forwarder.close();
+                } catch (IOException e) {
+                    log.warn("Error closing forwarder", e);
+                }
+            }
+        }
         trans.disconnect();
         super.disconnect();
     }
@@ -648,7 +659,12 @@ public class SSHClient
      */
     public LocalPortForwarder newLocalPortForwarder(LocalPortForwarder.Parameters parameters,
                                                     ServerSocket serverSocket) {
-        return new LocalPortForwarder(conn, parameters, serverSocket);
+        if (forwarders == null) {
+            forwarders = new ArrayList<LocalPortForwarder>();
+        }
+        LocalPortForwarder forwarder = new LocalPortForwarder(conn, parameters, serverSocket);
+        forwarders.add(forwarder);
+        return forwarder;
     }
 
     /**
diff --git a/src/main/java/net/schmizz/sshj/connection/channel/direct/LocalPortForwarder.java b/src/main/java/net/schmizz/sshj/connection/channel/direct/LocalPortForwarder.java
index a3511cac..1721e5a8 100644
--- a/src/main/java/net/schmizz/sshj/connection/channel/direct/LocalPortForwarder.java
+++ b/src/main/java/net/schmizz/sshj/connection/channel/direct/LocalPortForwarder.java
@@ -28,6 +28,7 @@ import org.slf4j.LoggerFactory;
 import java.io.IOException;
 import java.net.ServerSocket;
 import java.net.Socket;
+import java.net.SocketException;
 import java.util.concurrent.TimeUnit;
 
 import static com.hierynomus.sshj.backport.Sockets.asCloseable;
@@ -134,11 +135,33 @@ public class LocalPortForwarder {
             throws IOException {
         log.info("Listening on {}", serverSocket.getLocalSocketAddress());
         while (!Thread.currentThread().isInterrupted()) {
-            final Socket socket = serverSocket.accept();
-            log.debug("Got connection from {}", socket.getRemoteSocketAddress());
-            startChannel(socket);
+            try {
+                final Socket socket = serverSocket.accept();
+                log.debug("Got connection from {}", socket.getRemoteSocketAddress());
+                startChannel(socket);
+            } catch (SocketException e) {
+                if (!serverSocket.isClosed()) {
+                    throw e;
+                }
+            }
+        }
+        if (serverSocket.isClosed()) {
+            log.debug("LocalPortForwarder closed");
+        } else {
+            log.debug("LocalPortForwarder interrupted!");
         }
-        log.debug("Interrupted!");
     }
 
-}
\ No newline at end of file
+    /**
+     * Close the ServerSocket that's listening for connections to forward.
+     *
+     * @throws IOException
+     */
+    public void close() throws IOException {
+        if (!serverSocket.isClosed()) {
+            log.info("Closing listener on {}", serverSocket.getLocalSocketAddress());
+            serverSocket.close();
+        }
+    }
+
+}