From adc0451b3facf7c909208027fdfe0031a2d506de Mon Sep 17 00:00:00 2001
From: Jeroen van Erp <jeroen@hierynomus.com>
Date: Tue, 10 Jul 2018 16:29:09 +0200
Subject: [PATCH] Cleanup OpenSSHKeyFile and add Disconnection test

---
 .../userauth/keyprovider/OpenSSHKeyFile.java  | 32 +++++++++----------
 .../com/hierynomus/sshj/test/SshFixture.java  |  2 +-
 .../sshj/transport/DisconnectionTest.java     | 14 ++++++++
 3 files changed, 31 insertions(+), 17 deletions(-)

diff --git a/src/main/java/net/schmizz/sshj/userauth/keyprovider/OpenSSHKeyFile.java b/src/main/java/net/schmizz/sshj/userauth/keyprovider/OpenSSHKeyFile.java
index 39a9ec0c..da379e63 100644
--- a/src/main/java/net/schmizz/sshj/userauth/keyprovider/OpenSSHKeyFile.java
+++ b/src/main/java/net/schmizz/sshj/userauth/keyprovider/OpenSSHKeyFile.java
@@ -75,7 +75,12 @@ public class OpenSSHKeyFile
     @Override
     public void init(String privateKey, String publicKey) {
         if (publicKey != null) {
-            initPubKey(new StringReader(publicKey));
+            try {
+                initPubKey(new StringReader(publicKey));
+            } catch (IOException e) {
+                // let super provide both public & private key
+                log.warn("Error reading public key: {}", e.toString());
+            }
         }
         super.init(privateKey, null);
     }
@@ -85,23 +90,18 @@ public class OpenSSHKeyFile
      *
      * @param publicKey Public key accessible through a {@code Reader}
      */
-    private void initPubKey(Reader publicKey) {
+    private void initPubKey(Reader publicKey) throws IOException {
+        final BufferedReader br = new BufferedReader(publicKey);
         try {
-            final BufferedReader br = new BufferedReader(publicKey);
-            try {
-                final String keydata = br.readLine();
-                if (keydata != null) {
-                    String[] parts = keydata.trim().split(" ");
-                    assert parts.length >= 2;
-                    type = KeyType.fromString(parts[0]);
-                    pubKey = new Buffer.PlainBuffer(Base64.decode(parts[1])).readPublicKey();
-                }
-            } finally {
-                br.close();
+            final String keydata = br.readLine();
+            if (keydata != null) {
+                String[] parts = keydata.trim().split(" ");
+                assert parts.length >= 2;
+                type = KeyType.fromString(parts[0]);
+                pubKey = new Buffer.PlainBuffer(Base64.decode(parts[1])).readPublicKey();
             }
-        } catch (IOException e) {
-            // let super provide both public & private key
-            log.warn("Error reading public key: {}", e.toString());
+        } finally {
+            br.close();
         }
     }
 }
diff --git a/src/test/java/com/hierynomus/sshj/test/SshFixture.java b/src/test/java/com/hierynomus/sshj/test/SshFixture.java
index 93cbb3fb..ae2a0e2b 100644
--- a/src/test/java/com/hierynomus/sshj/test/SshFixture.java
+++ b/src/test/java/com/hierynomus/sshj/test/SshFixture.java
@@ -127,7 +127,7 @@ public class SshFixture extends ExternalResource {
             }
         });
         sshServer.setCommandFactory(commandFactory);
-
+        sshServer.setShellFactory(new ProcessShellFactory("ls"));
         return sshServer;
     }
 
diff --git a/src/test/java/com/hierynomus/sshj/transport/DisconnectionTest.java b/src/test/java/com/hierynomus/sshj/transport/DisconnectionTest.java
index 0d2d41d1..fce39036 100644
--- a/src/test/java/com/hierynomus/sshj/transport/DisconnectionTest.java
+++ b/src/test/java/com/hierynomus/sshj/transport/DisconnectionTest.java
@@ -16,13 +16,17 @@
 package com.hierynomus.sshj.transport;
 
 import com.hierynomus.sshj.test.SshFixture;
+import net.schmizz.sshj.DefaultConfig;
 import net.schmizz.sshj.SSHClient;
 import net.schmizz.sshj.common.DisconnectReason;
+import net.schmizz.sshj.connection.channel.direct.Session;
 import net.schmizz.sshj.transport.DisconnectListener;
 import net.schmizz.sshj.transport.TransportException;
+import net.schmizz.sshj.transport.verification.PromiscuousVerifier;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
+import org.junit.experimental.theories.suppliers.TestedOn;
 
 import java.io.IOException;
 import java.util.concurrent.TimeUnit;
@@ -89,4 +93,14 @@ public class DisconnectionTest {
         assertFalse(joinToClientTransport(2));
     }
 
+    @Test
+    public void shouldNotThrowTimeoutOnDisconnect() throws IOException {
+        fixture.getClient().authPassword("u", "u");
+        Session session = fixture.getClient().startSession();
+        session.allocateDefaultPTY();
+        Session.Shell shell = session.startShell();
+
+        session.close();
+        fixture.getClient().disconnect();
+    }
 }