Merge branch 'release_24_0_1' of https://github.com/dbeaver/dbeaver i…

…nto release_24_0_1
dbeaver · Mar 22, 2024 · e3d108b · e3d108b
2 parents e5282fb + dfc412c
commit e3d108b
Showing 1 changed file with 49 additions and 31 deletions.
diff --git a/plugins/org.jkiss.dbeaver.net.ssh.sshj/src/org/jkiss/dbeaver/model/net/ssh/SSHJSession.java b/plugins/org.jkiss.dbeaver.net.ssh.sshj/src/org/jkiss/dbeaver/model/net/ssh/SSHJSession.java
@@ -36,7 +36,9 @@
 import java.net.InetAddress;
 import java.net.ServerSocket;
 import java.util.Map;
+import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
 
 public class SSHJSession extends AbstractSession {
     private static final Log log = Log.getLog(SSHJSession.class);
@@ -65,15 +67,11 @@ public void disconnect(
         @NotNull DBWHandlerConfiguration configuration,
         long timeout
     ) throws DBException {
-        // FIXME: timeout is not used
-
-        for (LocalPortListener listener : listeners.values()) {
-            listener.disconnect();
-        }
-
+        // Listeners will be closed by the client itself
         listeners.clear();
 
         try {
+            // FIXME: timeout is not used
             client.disconnect();
         } catch (IOException e) {
             throw new DBException("Error disconnecting SSH session", e);
@@ -82,25 +80,15 @@ public void disconnect(
 
     @NotNull
     @Override
-    public SSHPortForwardConfiguration setupPortForward(@NotNull SSHPortForwardConfiguration cfg) throws DBException {
+    public SSHPortForwardConfiguration setupPortForward(@NotNull SSHPortForwardConfiguration config) throws DBException {
         try {
-            final ServerSocket ss = new ServerSocket(cfg.localPort(), 0, InetAddress.getByName(cfg.localHost()));
-            final Parameters parameters = new Parameters(cfg.localHost(), ss.getLocalPort(), cfg.remoteHost(), cfg.remotePort());
-            final LocalPortForwarder forwarder = client.newLocalPortForwarder(parameters, ss);
-            final LocalPortListener listener = new LocalPortListener(forwarder, parameters);
-
-            final SSHPortForwardConfiguration resolved = new SSHPortForwardConfiguration(
-                cfg.localHost(),
-                ss.getLocalPort(),
-                cfg.remoteHost(),
-                cfg.remotePort()
-            );
+            final LocalPortListener listener = LocalPortListener.setup(client, config);
+            final SSHPortForwardConfiguration resolved = Objects.requireNonNull(listener.resolved);
 
-            listener.start();
             listeners.put(resolved, listener);
 
             return resolved;
-        } catch (IOException e) {
+        } catch (Exception e) {
             throw new DBException("Error setting up port forwarding", e);
         }
     }
@@ -186,21 +174,43 @@ private SFTPClient openSftpClient() throws IOException {
     }
 
     private static class LocalPortListener extends Thread {
-        private final LocalPortForwarder forwarder;
+        private final SSHClient client;
+        private final SSHPortForwardConfiguration config;
+        private final CountDownLatch started = new CountDownLatch(1);
+
+        private volatile LocalPortForwarder forwarder;
+        private volatile SSHPortForwardConfiguration resolved;
+
+        public LocalPortListener(@NotNull SSHClient client, @NotNull SSHPortForwardConfiguration config) {
+            this.client = client;
+            this.config = config;
+        }
+
+        @NotNull
+        public static LocalPortListener setup(
+            @NotNull SSHClient client,
+            @NotNull SSHPortForwardConfiguration config
+        ) throws InterruptedException {
+            final LocalPortListener listener = new LocalPortListener(client, config);
 
-        public LocalPortListener(@NotNull LocalPortForwarder forwarder, @NotNull Parameters parameters) {
-            this.forwarder = forwarder;
+            listener.start();
+            listener.await();
 
-            setName(String.format(
-                "Port forwarder listener (%s:%d -> %s:%d)",
-                parameters.getLocalHost(), parameters.getLocalPort(),
-                parameters.getRemoteHost(), parameters.getRemotePort()
-            ));
+            return listener;
         }
 
         @Override
         public void run() {
             try {
+                final ServerSocket socket = new ServerSocket(config.localPort(), 0, InetAddress.getByName(config.localHost()));
+                final Parameters parameters = new Parameters(config.localHost(), socket.getLocalPort(), config.remoteHost(), config.remotePort());
+
+                forwarder = client.newLocalPortForwarder(parameters, socket);
+                resolved = new SSHPortForwardConfiguration(config.localHost(), socket.getLocalPort(), config.remoteHost(), config.remotePort());
+
+                setName("Port forwarder listener (" + resolved + ")");
+
+                started.countDown();
                 forwarder.listen();
             } catch (IOException e) {
                 log.error("Error while listening on the port forwarder", e);
@@ -209,12 +219,20 @@ public void run() {
 
         public void disconnect() {
             try {
-                if (forwarder.isRunning()) {
-                    forwarder.close();
-                }
+                forwarder.close();
+                forwarder = null;
             } catch (Exception e) {
                 log.error("Error while stopping port forwarding", e);
             }
         }
+
+        private void await() throws InterruptedException {
+            started.await();
+
+            while (!forwarder.isRunning()) {
+                // Wait for the forwarder to actually start
+                Thread.yield();
+            }
+        }
     }
 }