From d59efaa5f98877d51ec6351736ebaf18a6f4fa46 Mon Sep 17 00:00:00 2001
From: hierynomus <jeroen@javadude.nl>
Date: Tue, 20 Jan 2015 17:13:47 +0100
Subject: [PATCH] Fixed KeepAliveRunner for when service not yet set

---
 .../net/schmizz/sshj/examples/KeepAlive.java  | 55 +++++++++++++++++++
 .../schmizz/keepalive/KeepAliveRunner.java    |  9 ++-
 2 files changed, 61 insertions(+), 3 deletions(-)
 create mode 100644 examples/src/main/java/net/schmizz/sshj/examples/KeepAlive.java

diff --git a/examples/src/main/java/net/schmizz/sshj/examples/KeepAlive.java b/examples/src/main/java/net/schmizz/sshj/examples/KeepAlive.java
new file mode 100644
index 00000000..8d6bdcf8
--- /dev/null
+++ b/examples/src/main/java/net/schmizz/sshj/examples/KeepAlive.java
@@ -0,0 +1,55 @@
+/**
+ * Copyright 2009 sshj contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package net.schmizz.sshj.examples;
+
+import net.schmizz.keepalive.KeepAliveProvider;
+import net.schmizz.sshj.DefaultConfig;
+import net.schmizz.sshj.SSHClient;
+import net.schmizz.sshj.common.IOUtils;
+import net.schmizz.sshj.connection.channel.direct.Session;
+import net.schmizz.sshj.connection.channel.direct.Session.Command;
+import net.schmizz.sshj.transport.verification.PromiscuousVerifier;
+
+import java.io.IOException;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+/** This examples demonstrates how to setup keep-alive to detect connection dropping. */
+public class KeepAlive {
+
+    public static void main(String... args)
+            throws IOException, InterruptedException {
+        DefaultConfig defaultConfig = new DefaultConfig();
+        defaultConfig.setKeepAliveProvider(KeepAliveProvider.KEEP_ALIVE);
+        final SSHClient ssh = new SSHClient(defaultConfig);
+        try {
+            ssh.addHostKeyVerifier(new PromiscuousVerifier());
+            ssh.connect(args[0]);
+            ssh.getConnection().getKeepAlive().setKeepAliveInterval(5); //every 60sec
+            ssh.authPassword(args[1], args[2]);
+            Session session = ssh.startSession();
+            session.allocateDefaultPTY();
+            new CountDownLatch(1).await();
+            try {
+                session.allocateDefaultPTY();
+            } finally {
+                session.close();
+            }
+        } finally {
+            ssh.disconnect();
+        }
+    }
+}
diff --git a/src/main/java/net/schmizz/keepalive/KeepAliveRunner.java b/src/main/java/net/schmizz/keepalive/KeepAliveRunner.java
index 5b82f634..15b98302 100644
--- a/src/main/java/net/schmizz/keepalive/KeepAliveRunner.java
+++ b/src/main/java/net/schmizz/keepalive/KeepAliveRunner.java
@@ -50,9 +50,12 @@ public class KeepAliveRunner extends KeepAlive {
 
     @Override
     protected void doKeepAlive() throws TransportException, ConnectionException {
-        emptyQueue(queue);
-        checkMaxReached(queue);
-        queue.add(conn.sendGlobalRequest("keepalive@openssh.com", true, new byte[0]));
+        // Ensure the service is set... This means that the key exchange is done and the connection is up.
+        if (conn.equals(conn.getTransport().getService())) {
+            emptyQueue(queue);
+            checkMaxReached(queue);
+            queue.add(conn.sendGlobalRequest("keepalive@openssh.com", true, new byte[0]));
+        }
     }
 
     private void checkMaxReached(Queue<Promise<SSHPacket, ConnectionException>> queue) throws ConnectionException {