5 files changed, 62 insertions, 23 deletions

diff --git a/src/main/java/com/gitblit/transport/ssh/PublicKeyAuthenticator.java b/src/main/java/com/gitblit/transport/ssh/CachingPublicKeyAuthenticator.java
index 84e7afa5..ee1de591 100644
--- a/src/main/java/com/gitblit/transport/ssh/PublicKeyAuthenticator.java
+++ b/src/main/java/com/gitblit/transport/ssh/CachingPublicKeyAuthenticator.java
@@ -16,9 +16,14 @@
 package com.gitblit.transport.ssh;
 
 import java.security.PublicKey;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Locale;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
+import org.apache.sshd.common.Session;
+import org.apache.sshd.common.SessionListener;
 import org.apache.sshd.server.PublickeyAuthenticator;
 import org.apache.sshd.server.session.ServerSession;
 import org.slf4j.Logger;
@@ -26,13 +31,15 @@ import org.slf4j.LoggerFactory;
 
 import com.gitblit.manager.IAuthenticationManager;
 import com.gitblit.models.UserModel;
+import com.google.common.base.Preconditions;
 
 /**
- *
+ * 
  * @author Eric Myrhe
- *
+ * 
  */
-public class PublicKeyAuthenticator implements PublickeyAuthenticator {
+public class CachingPublicKeyAuthenticator implements PublickeyAuthenticator,
+		SessionListener {
 
 	protected final Logger log = LoggerFactory.getLogger(getClass());
 
@@ -40,26 +47,41 @@ public class PublicKeyAuthenticator implements PublickeyAuthenticator {
 
 	protected final IAuthenticationManager authManager;
 
-	public PublicKeyAuthenticator(IKeyManager keyManager, IAuthenticationManager authManager) {
+	private final Map<ServerSession, Map<PublicKey, Boolean>> cache =
+			new ConcurrentHashMap<ServerSession, Map<PublicKey, Boolean>>();
+
+	public CachingPublicKeyAuthenticator(IKeyManager keyManager,
+			IAuthenticationManager authManager) {
 		this.keyManager = keyManager;
 		this.authManager = authManager;
 	}
 
 	@Override
-	public boolean authenticate(String username, final PublicKey suppliedKey,
-			final ServerSession session) {
-		final SshDaemonClient client = session.getAttribute(SshDaemonClient.KEY);
-
-		if (client.getUser() != null) {
-			// TODO why do we re-authenticate?
-			log.info("{} has already authenticated!", username);
-			return true;
+	public boolean authenticate(String username, PublicKey key,
+			ServerSession session) {
+		Map<PublicKey, Boolean> map = cache.get(session);
+		if (map == null) {
+			map = new HashMap<PublicKey, Boolean>();
+			cache.put(session, map);
+			session.addListener(this);
 		}
+		if (map.containsKey(key)) {
+			return map.get(key);
+		}
+		boolean result = doAuthenticate(username, key, session);
+		map.put(key, result);
+		return result;
+	}
 
+	private boolean doAuthenticate(String username, PublicKey suppliedKey,
+			ServerSession session) {
+		SshDaemonClient client = session.getAttribute(SshDaemonClient.KEY);
+		Preconditions.checkState(client.getUser() == null);
 		username = username.toLowerCase(Locale.US);
 		List<PublicKey> keys = keyManager.getKeys(username);
 		if (keys == null || keys.isEmpty()) {
-			log.info("{} has not added any public keys for ssh authentication", username);
+			log.info("{} has not added any public keys for ssh authentication",
+					username);
 			return false;
 		}
 
@@ -73,11 +95,23 @@ public class PublicKeyAuthenticator implements PublickeyAuthenticator {
 			}
 		}
 
-		log.warn("could not authenticate {} for SSH using the supplied public key", username);
+		log.warn(
+				"could not authenticate {} for SSH using the supplied public key",
+				username);
 		return false;
 	}
 
 	public IKeyManager getKeyManager() {
 		return keyManager;
 	}
+
+	public void sessionCreated(Session session) {
+	}
+
+	public void sessionEvent(Session sesssion, Event event) {
+	}
+
+	public void sessionClosed(Session session) {
+		cache.remove(session);
+	}
 }
diff --git a/src/main/java/com/gitblit/transport/ssh/SshCommandFactory.java b/src/main/java/com/gitblit/transport/ssh/SshCommandFactory.java
index da57f76e..48e8869a 100644
--- a/src/main/java/com/gitblit/transport/ssh/SshCommandFactory.java
+++ b/src/main/java/com/gitblit/transport/ssh/SshCommandFactory.java
@@ -61,10 +61,12 @@ public class SshCommandFactory implements CommandFactory {
 	private static final Logger logger = LoggerFactory.getLogger(SshCommandFactory.class);
 
 	private final IGitblit gitblit;
-	private final PublicKeyAuthenticator keyAuthenticator;
+	private final CachingPublicKeyAuthenticator keyAuthenticator;
 	private final ScheduledExecutorService startExecutor;
 
-	public SshCommandFactory(IGitblit gitblit, PublicKeyAuthenticator keyAuthenticator, IdGenerator idGenerator) {
+	public SshCommandFactory(IGitblit gitblit,
+			CachingPublicKeyAuthenticator keyAuthenticator,
+			IdGenerator idGenerator) {
 		this.gitblit = gitblit;
 		this.keyAuthenticator = keyAuthenticator;
 
@@ -252,6 +254,7 @@ public class SshCommandFactory implements CommandFactory {
 			}
 		}
 
+		@SuppressWarnings("unused")
 		private void onDestroy() {
 			synchronized (this) {
 				if (cmd != null) {
diff --git a/src/main/java/com/gitblit/transport/ssh/SshDaemon.java b/src/main/java/com/gitblit/transport/ssh/SshDaemon.java
index c3d48600..c954b347 100644
--- a/src/main/java/com/gitblit/transport/ssh/SshDaemon.java
+++ b/src/main/java/com/gitblit/transport/ssh/SshDaemon.java
@@ -104,7 +104,8 @@ public class SshDaemon {
 			addr = new InetSocketAddress(bindInterface, port);
 		}
 
-		PublicKeyAuthenticator keyAuthenticator = new PublicKeyAuthenticator(keyManager, gitblit);
+		CachingPublicKeyAuthenticator keyAuthenticator = 
+				new CachingPublicKeyAuthenticator(keyManager, gitblit);
 
 		sshd = SshServer.setUpDefaultServer();
 		sshd.setPort(addr.getPort());
@@ -176,6 +177,7 @@ public class SshDaemon {
 		}
 	}
 
+	@SuppressWarnings("unchecked")
 	protected IKeyManager getKeyManager() {
 		IKeyManager keyManager = null;
 		IStoredSettings settings = gitblit.getSettings();
diff --git a/src/main/java/com/gitblit/transport/ssh/commands/BaseKeyCommand.java b/src/main/java/com/gitblit/transport/ssh/commands/BaseKeyCommand.java
index 36475244..f92ea6f9 100644
--- a/src/main/java/com/gitblit/transport/ssh/commands/BaseKeyCommand.java
+++ b/src/main/java/com/gitblit/transport/ssh/commands/BaseKeyCommand.java
@@ -21,7 +21,7 @@ import java.io.InputStreamReader;
 import java.io.UnsupportedEncodingException;
 import java.util.List;
 
-import com.gitblit.transport.ssh.PublicKeyAuthenticator;
+import com.gitblit.transport.ssh.CachingPublicKeyAuthenticator;
 import com.google.common.base.Charsets;
 
 /**
@@ -51,8 +51,8 @@ public abstract class BaseKeyCommand extends SshCommand {
 		return sshKeys;
 	}
 
-	protected PublicKeyAuthenticator authenticator;
-	public void setAuthenticator(PublicKeyAuthenticator authenticator) {
+	protected CachingPublicKeyAuthenticator authenticator;
+	public void setAuthenticator(CachingPublicKeyAuthenticator authenticator) {
 		this.authenticator = authenticator;
 	}
 }
diff --git a/src/main/java/com/gitblit/transport/ssh/commands/DispatchCommand.java b/src/main/java/com/gitblit/transport/ssh/commands/DispatchCommand.java
index 1e43e2f2..3c041af6 100644
--- a/src/main/java/com/gitblit/transport/ssh/commands/DispatchCommand.java
+++ b/src/main/java/com/gitblit/transport/ssh/commands/DispatchCommand.java
@@ -34,7 +34,7 @@ import com.gitblit.git.GitblitUploadPackFactory;
 import com.gitblit.git.RepositoryResolver;
 import com.gitblit.models.UserModel;
 import com.gitblit.transport.ssh.CommandMetaData;
-import com.gitblit.transport.ssh.PublicKeyAuthenticator;
+import com.gitblit.transport.ssh.CachingPublicKeyAuthenticator;
 import com.gitblit.transport.ssh.SshDaemonClient;
 import com.gitblit.utils.cli.SubcommandHandler;
 import com.google.common.base.Charsets;
@@ -237,9 +237,9 @@ public class DispatchCommand extends BaseCommand {
 		this.gitblitReceivePackFactory = gitblitReceivePackFactory;
 	}
 
-	private PublicKeyAuthenticator authenticator;
+	private CachingPublicKeyAuthenticator authenticator;
 
-	public void setAuthenticator(PublicKeyAuthenticator authenticator) {
+	public void setAuthenticator(CachingPublicKeyAuthenticator authenticator) {
 		this.authenticator = authenticator;
 	}
 }