Merge pull request #17 from kruton/extended-hostkey

Add extended server hostkey verification API

3 files changed, 87 insertions, 2 deletions

diff --git a/sshlib/src/main/java/com/trilead/ssh2/ExtendedServerHostKeyVerifier.java b/sshlib/src/main/java/com/trilead/ssh2/ExtendedServerHostKeyVerifier.java
new file mode 100644
index 0000000..f757aa6
--- /dev/null
+++ b/sshlib/src/main/java/com/trilead/ssh2/ExtendedServerHostKeyVerifier.java
@@ -0,0 +1,47 @@
+package com.trilead.ssh2;
+
+import java.util.List;
+
+/**
+ * This extends the {@link ServerHostKeyVerifier} interface by allowing the remote server to indicate it has multiple
+ * server key algorithms available. After authentication, the {@link #getKnownKeyAlgorithmsForHost(String, int)} method
+ * may be called and compared against the list of server-controller keys. If a key algorithm has been added then
+ * {@link #addServerHostKey(String, int, String, byte[])} will be called. If a key algorithm has been removed, then
+ * {@link #removeServerHostKey(String, int, String, byte[])} will be called.
+ *
+ * @author Kenny Root
+ */
+public abstract class ExtendedServerHostKeyVerifier implements ServerHostKeyVerifier {
+	/**
+	 * Called during connection to determine which keys are known for this host.
+	 *
+	 * @param hostname the hostname used to create the {@link Connection} object
+	 * @param port the server's remote TCP port
+	 * @return list of hostkey algorithms for the given <code>hostname</code> and <code>port</code> combination
+	 * 			or {@code null} if none are known.
+	 */
+	public abstract List<String> getKnownKeyAlgorithmsForHost(String hostname, int port);
+
+	/**
+	 * After authentication, if the server indicates it no longer uses this key, this method will be called
+	 * for the app to remove its record of it.
+	 *
+	 * @param hostname the hostname used to create the {@link Connection} object
+	 * @param port the server's remote TCP port
+	 * @param serverHostKeyAlgorithm key algorithm of removed key
+	 * @param serverHostKey key data of removed key
+	 */
+	public abstract void removeServerHostKey(String hostname, int port, String serverHostKeyAlgorithm,
+			byte[] serverHostKey);
+
+	/**
+	 * After authentication, if the server indicates it has another <code>keyAlgorithm</code>, this method will be
+	 * called for the app to add it to its record of known keys for this <code>hostname</code>.
+	 *
+	 * @param hostname the hostname used to create the {@link Connection} object
+	 * @param port the server's remote TCP port
+	 * @param keyAlgorithm SSH standard name for the key to be added
+	 * @param serverHostKey SSH encoding of the key data for the key to be added
+	 */
+	public abstract void addServerHostKey(String hostname, int port, String keyAlgorithm, byte[] serverHostKey);
+}
diff --git a/sshlib/src/main/java/com/trilead/ssh2/ServerHostKeyVerifier.java b/sshlib/src/main/java/com/trilead/ssh2/ServerHostKeyVerifier.java
index ac65955..6e18b3e 100644
--- a/sshlib/src/main/java/com/trilead/ssh2/ServerHostKeyVerifier.java
+++ b/sshlib/src/main/java/com/trilead/ssh2/ServerHostKeyVerifier.java
@@ -26,6 +26,6 @@ public interface ServerHostKeyVerifier
 	 *         connection will be closed.
 	 * @throws Exception Will be wrapped with an IOException, extended version of returning false =)
 	 */
-	public boolean verifyServerHostKey(String hostname, int port, String serverHostKeyAlgorithm, byte[] serverHostKey)
+	boolean verifyServerHostKey(String hostname, int port, String serverHostKeyAlgorithm, byte[] serverHostKey)
 			throws Exception;
 }
diff --git a/sshlib/src/main/java/com/trilead/ssh2/transport/KexManager.java b/sshlib/src/main/java/com/trilead/ssh2/transport/KexManager.java
index ab6d0b6..3b7db3e 100644
--- a/sshlib/src/main/java/com/trilead/ssh2/transport/KexManager.java
+++ b/sshlib/src/main/java/com/trilead/ssh2/transport/KexManager.java
@@ -8,12 +8,14 @@ import java.security.SecureRandom;
 import java.security.interfaces.DSAPublicKey;
 import java.security.interfaces.ECPublicKey;
 import java.security.interfaces.RSAPublicKey;
+import java.util.ArrayList;
 import java.util.LinkedHashSet;
+import java.util.List;
 import java.util.Set;
-import java.util.TreeSet;
 
 import com.trilead.ssh2.ConnectionInfo;
 import com.trilead.ssh2.DHGexParameters;
+import com.trilead.ssh2.ExtendedServerHostKeyVerifier;
 import com.trilead.ssh2.ServerHostKeyVerifier;
 import com.trilead.ssh2.compression.CompressionFactory;
 import com.trilead.ssh2.compression.ICompressor;
@@ -282,6 +284,8 @@ public class KexManager
 	public synchronized void initiateKEX(CryptoWishList cwl, DHGexParameters dhgex) throws IOException
 	{
 		nextKEXcryptoWishList = cwl;
+		filterHostKeyTypes(nextKEXcryptoWishList);
+
 		nextKEXdhgexParameters = dhgex;
 
 		if (kxs == null)
@@ -295,6 +299,40 @@ public class KexManager
 		}
 	}
 
+	/**
+	 * If the verifier can indicate which algorithms it knows about for this host, then
+	 * filter out our crypto wish list to only include those algorithms. Otherwise we'll
+	 * negotiate a host key we have not previously confirmed.
+	 *
+	 * @param cwl crypto wish list to filter
+	 */
+	private void filterHostKeyTypes(CryptoWishList cwl) {
+		if (verifier instanceof ExtendedServerHostKeyVerifier) {
+			ExtendedServerHostKeyVerifier extendedVerifier = (ExtendedServerHostKeyVerifier) verifier;
+
+			List<String> knownAlgorithms = extendedVerifier.getKnownKeyAlgorithmsForHost(hostname, port);
+			if (knownAlgorithms != null && knownAlgorithms.size() > 0) {
+				ArrayList<String> filteredAlgorithms = new ArrayList<String>(knownAlgorithms.size());
+
+				/*
+				 * Look at our current wish list and adjust it based on what the client already knows, but
+				 * be careful to keep it in the order desired by the wish list.
+				 */
+				for (String capableAlgo : cwl.serverHostKeyAlgorithms) {
+					for (String knownAlgo : knownAlgorithms) {
+						if (capableAlgo.equals(knownAlgo)) {
+							filteredAlgorithms.add(knownAlgo);
+						}
+					}
+				}
+
+				if (filteredAlgorithms.size() > 0) {
+					cwl.serverHostKeyAlgorithms = filteredAlgorithms.toArray(new String[filteredAlgorithms.size()]);
+				}
+			}
+		}
+	}
+
 	private boolean establishKeyMaterial()
 	{
 		try