1 files changed, 15 insertions, 17 deletions

diff --git a/src/main/java/com/gitblit/transport/ssh/MemoryKeyManager.java b/src/main/java/com/gitblit/transport/ssh/MemoryKeyManager.java
index 26bd021a..18f9a4e1 100644
--- a/src/main/java/com/gitblit/transport/ssh/MemoryKeyManager.java
+++ b/src/main/java/com/gitblit/transport/ssh/MemoryKeyManager.java
@@ -15,7 +15,6 @@
  */
 package com.gitblit.transport.ssh;
 
-import java.security.PublicKey;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -29,10 +28,10 @@ import java.util.Map;
  */
 public class MemoryKeyManager extends IPublicKeyManager {
 
-	Map<String, List<PublicKey>> keys;
+	Map<String, List<SshKey>> keys;
 
 	public MemoryKeyManager() {
-		keys = new HashMap<String, List<PublicKey>>();
+		keys = new HashMap<String, List<SshKey>>();
 	}
 
 	@Override
@@ -62,7 +61,7 @@ public class MemoryKeyManager extends IPublicKeyManager {
 	}
 
 	@Override
-	protected List<PublicKey> getKeysImpl(String username) {
+	protected List<SshKey> getKeysImpl(String username) {
 		String id = username.toLowerCase();
 		if (keys.containsKey(id)) {
 			return keys.get(id);
@@ -71,13 +70,21 @@ public class MemoryKeyManager extends IPublicKeyManager {
 	}
 
 	@Override
-	public boolean addKey(String username, String data) {
-		return false;
+	public boolean addKey(String username, SshKey key) {
+		String id = username.toLowerCase();
+		if (!keys.containsKey(id)) {
+			keys.put(id, new ArrayList<SshKey>());
+		}
+		return keys.get(id).add(key);
 	}
 
 	@Override
-	public boolean removeKey(String username, String data) {
-		return false;
+	public boolean removeKey(String username, SshKey key) {
+		String id = username.toLowerCase();
+		if (!keys.containsKey(id)) {
+			return false;
+		}
+		return keys.get(id).remove(key);
 	}
 
 	@Override
@@ -86,13 +93,4 @@ public class MemoryKeyManager extends IPublicKeyManager {
 		keys.remove(id.toLowerCase());
 		return true;
 	}
-
-	/* Test method for populating the memory key manager */
-	public void addKey(String username, PublicKey key) {
-		String id = username.toLowerCase();
-		if (!keys.containsKey(id)) {
-			keys.put(id, new ArrayList<PublicKey>());
-		}
-		keys.get(id).add(key);
-	}
 }