1 files changed, 61 insertions, 17 deletions

diff --git a/src/com/trilead/ssh2/channel/AuthAgentForwardThread.java b/src/com/trilead/ssh2/channel/AuthAgentForwardThread.java
index 57b9a5e..d3f10a3 100644
--- a/src/com/trilead/ssh2/channel/AuthAgentForwardThread.java
+++ b/src/com/trilead/ssh2/channel/AuthAgentForwardThread.java
@@ -21,7 +21,20 @@ import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.math.BigInteger;
+import java.security.KeyFactory;
+import java.security.KeyPair;
+import java.security.NoSuchAlgorithmException;
+import java.security.PrivateKey;
+import java.security.PublicKey;
 import java.security.SecureRandom;
+import java.security.interfaces.DSAPrivateKey;
+import java.security.interfaces.RSAPrivateKey;
+import java.security.spec.DSAPrivateKeySpec;
+import java.security.spec.DSAPublicKeySpec;
+import java.security.spec.InvalidKeySpecException;
+import java.security.spec.KeySpec;
+import java.security.spec.RSAPrivateKeySpec;
+import java.security.spec.RSAPublicKeySpec;
 import java.util.Map;
 import java.util.Map.Entry;
 
@@ -29,12 +42,8 @@ import com.trilead.ssh2.AuthAgentCallback;
 import com.trilead.ssh2.log.Logger;
 import com.trilead.ssh2.packets.TypesReader;
 import com.trilead.ssh2.packets.TypesWriter;
-import com.trilead.ssh2.signature.DSAPrivateKey;
 import com.trilead.ssh2.signature.DSASHA1Verify;
-import com.trilead.ssh2.signature.DSASignature;
-import com.trilead.ssh2.signature.RSAPrivateKey;
 import com.trilead.ssh2.signature.RSASHA1Verify;
-import com.trilead.ssh2.signature.RSASignature;
 
 /**
  * AuthAgentForwardThread.
@@ -268,7 +277,7 @@ public class AuthAgentForwardThread extends Thread implements IChannelWorkerThre
 
 			String type = tr.readString();
 
-			Object key;
+			KeyPair pair;
 			String comment;
 
 			if (type.equals("ssh-rsa")) {
@@ -280,7 +289,24 @@ public class AuthAgentForwardThread extends Thread implements IChannelWorkerThre
 				tr.readMPINT(); // q
 				comment = tr.readString();
 
-				key = new RSAPrivateKey(d, e, n);
+				KeySpec pubSpec = new RSAPublicKeySpec(n, e);
+				KeySpec privSpec = new RSAPrivateKeySpec(n, d);
+
+				PublicKey pubKey;
+				PrivateKey privKey;
+				try {
+					KeyFactory kf = KeyFactory.getInstance("RSA");
+					pubKey = kf.generatePublic(pubSpec);
+					privKey = kf.generatePrivate(privSpec);
+				} catch (NoSuchAlgorithmException ex) {
+					// TODO: log error
+					return;
+				} catch (InvalidKeySpecException ex) {
+					// TODO: log error
+					return;
+				}
+
+				pair = new KeyPair(pubKey, privKey);
 			} else if (type.equals("ssh-dss")) {
 				BigInteger p = tr.readMPINT();
 				BigInteger q = tr.readMPINT();
@@ -289,7 +315,24 @@ public class AuthAgentForwardThread extends Thread implements IChannelWorkerThre
 				BigInteger x = tr.readMPINT();
 				comment = tr.readString();
 
-				key = new DSAPrivateKey(p, q, g, y, x);
+				KeySpec pubSpec = new DSAPublicKeySpec(y, p, q, g);
+				KeySpec privSpec = new DSAPrivateKeySpec(x, p, q, g);
+
+				PublicKey pubKey;
+				PrivateKey privKey;
+				try {
+					KeyFactory kf = KeyFactory.getInstance("DSA");
+					pubKey = kf.generatePublic(pubSpec);
+					privKey = kf.generatePrivate(privSpec);
+				} catch (NoSuchAlgorithmException ex) {
+					// TODO: log error
+					return;
+				} catch (InvalidKeySpecException ex) {
+					// TODO: log error
+					return;
+				}
+
+				pair = new KeyPair(pubKey, privKey);
 			} else {
 				os.write(SSH_AGENT_FAILURE);
 				return;
@@ -313,7 +356,7 @@ public class AuthAgentForwardThread extends Thread implements IChannelWorkerThre
 				}
 			}
 
-			if (authAgent.addIdentity(key, comment, confirmUse, lifetime))
+			if (authAgent.addIdentity(pair, comment, confirmUse, lifetime))
 				os.write(SSH_AGENT_SUCCESS);
 			else
 				os.write(SSH_AGENT_FAILURE);
@@ -390,7 +433,7 @@ public class AuthAgentForwardThread extends Thread implements IChannelWorkerThre
 			if (failWhenLocked())
 				return;
 
-			byte[] publicKey = tr.readByteString();
+			byte[] publicKeyBytes = tr.readByteString();
 			byte[] challenge = tr.readByteString();
 
 			int flags = tr.readUINT32();
@@ -401,22 +444,23 @@ public class AuthAgentForwardThread extends Thread implements IChannelWorkerThre
 				return;
 			}
 
-			Object trileadKey = authAgent.getPrivateKey(publicKey);
+			KeyPair pair = authAgent.getKeyPair(publicKeyBytes);
 
-			if (trileadKey == null) {
+			if (pair == null) {
 				os.write(SSH_AGENT_FAILURE);
 				return;
 			}
 
 			byte[] response;
 
-			if (trileadKey instanceof RSAPrivateKey) {
-				RSASignature signature = RSASHA1Verify.generateSignature(challenge,
-						(RSAPrivateKey) trileadKey);
+			PrivateKey privKey = pair.getPrivate();
+			if (privKey instanceof RSAPrivateKey) {
+				byte[] signature = RSASHA1Verify.generateSignature(challenge,
+						(RSAPrivateKey) privKey);
 				response = RSASHA1Verify.encodeSSHRSASignature(signature);
-			} else if (trileadKey instanceof DSAPrivateKey) {
-				DSASignature signature = DSASHA1Verify.generateSignature(challenge,
-						(DSAPrivateKey) trileadKey, new SecureRandom());
+			} else if (privKey instanceof DSAPrivateKey) {
+				byte[] signature = DSASHA1Verify.generateSignature(challenge,
+						(DSAPrivateKey) privKey, new SecureRandom());
 				response = DSASHA1Verify.encodeSSHDSASignature(signature);
 			} else {
 				os.write(SSH_AGENT_FAILURE);