--- a/common-kex.c	Sun Apr 14 22:49:10 2013 +0800
+++ b/common-kex.c	Sun Apr 14 22:49:19 2013 +0800
@@ -80,9 +80,10 @@
 static const int DH_G_VAL = 2;
 
 static void kexinitialise();
-void gen_new_keys();
+static void gen_new_keys();
 #ifndef DISABLE_ZLIB
-static void gen_new_zstreams();
+static void gen_new_zstream_recv();
+static void gen_new_zstream_trans();
 #endif
 static void read_kex_algos();
 /* helper function for gen_new_keys */
@@ -131,8 +132,8 @@
 	/* languages_server_to_client */
 	buf_putstring(ses.writepayload, "", 0);
 
-	/* first_kex_packet_follows - unimplemented for now */
-	buf_putbyte(ses.writepayload, 0x00);
+	/* first_kex_packet_follows */
+	buf_putbyte(ses.writepayload, (ses.send_kex_first_guess != NULL));
 
 	/* reserved unit32 */
 	buf_putint(ses.writepayload, 0);
@@ -144,16 +145,56 @@
 	encrypt_packet();
 	ses.dataallowed = 0; /* don't send other packets during kex */
 
+	ses.kexstate.sentkexinit = 1;
+
+	ses.newkeys = (struct key_context*)m_malloc(sizeof(struct key_context));
+
+	if (ses.send_kex_first_guess) {
+		ses.newkeys->algo_kex = sshkex[0].val;
+		ses.newkeys->algo_hostkey = sshhostkey[0].val;
+		ses.send_kex_first_guess();
+	}
+
 	TRACE(("DATAALLOWED=0"))
 	TRACE(("-> KEXINIT"))
-	ses.kexstate.sentkexinit = 1;
+
 }
 
-/* *** NOTE regarding (send|recv)_msg_newkeys *** 
- * Changed by mihnea from the original kex.c to set dataallowed after a 
- * completed key exchange, no matter the order in which it was performed.
- * This enables client mode without affecting server functionality.
- */
+static void switch_keys() {
+	TRACE2(("enter switch_keys"))
+	if (!(ses.kexstate.sentkexinit && ses.kexstate.recvkexinit)) {
+		dropbear_exit("Unexpected newkeys message");
+	}
+
+	if (!ses.keys) {
+		ses.keys = m_malloc(sizeof(*ses.newkeys));
+	}
+	if (ses.kexstate.recvnewkeys && ses.newkeys->recv.valid) {
+		TRACE(("switch_keys recv"))
+		gen_new_zstream_recv();
+		ses.keys->recv = ses.newkeys->recv;
+		m_burn(&ses.newkeys->recv, sizeof(ses.newkeys->recv));
+		ses.newkeys->recv.valid = 0;
+	}
+	if (ses.kexstate.sentnewkeys && ses.newkeys->trans.valid) {
+		TRACE(("switch_keys trans"))
+		gen_new_zstream_trans();
+		ses.keys->trans = ses.newkeys->trans;
+		m_burn(&ses.newkeys->trans, sizeof(ses.newkeys->trans));
+		ses.newkeys->trans.valid = 0;
+	}
+	if (ses.kexstate.sentnewkeys && ses.kexstate.recvnewkeys)
+	{
+		TRACE(("switch_keys done"))
+		ses.keys->algo_kex = ses.newkeys->algo_kex;
+		ses.keys->algo_hostkey = ses.newkeys->algo_hostkey;
+		ses.keys->allow_compress = 0;
+		m_free(ses.newkeys);
+		ses.newkeys = NULL;
+		kexinitialise();
+	}
+	TRACE2(("leave switch_keys"))
+}
 
 /* Bring new keys into use after a key exchange, and let the client know*/
 void send_msg_newkeys() {
@@ -164,44 +205,25 @@
 	CHECKCLEARTOWRITE();
 	buf_putbyte(ses.writepayload, SSH_MSG_NEWKEYS);
 	encrypt_packet();
+
 	
-
 	/* set up our state */
-	if (ses.kexstate.recvnewkeys) {
-		TRACE(("while RECVNEWKEYS=1"))
-		gen_new_keys();
-		kexinitialise(); /* we've finished with this kex */
-		TRACE((" -> DATAALLOWED=1"))
-		ses.dataallowed = 1; /* we can send other packets again now */
-		ses.kexstate.donefirstkex = 1;
-	} else {
-		ses.kexstate.sentnewkeys = 1;
-		TRACE(("SENTNEWKEYS=1"))
-	}
+	ses.kexstate.sentnewkeys = 1;
+	ses.kexstate.donefirstkex = 1;
+	ses.dataallowed = 1; /* we can send other packets again now */
+	gen_new_keys();
+	switch_keys();
 
-	TRACE(("-> MSG_NEWKEYS"))
 	TRACE(("leave send_msg_newkeys"))
 }
 
 /* Bring the new keys into use after a key exchange */
 void recv_msg_newkeys() {
 
-	TRACE(("<- MSG_NEWKEYS"))
 	TRACE(("enter recv_msg_newkeys"))
 
-	/* simply check if we've sent SSH_MSG_NEWKEYS, and if so,
-	 * switch to the new keys */
-	if (ses.kexstate.sentnewkeys) {
-		TRACE(("while SENTNEWKEYS=1"))
-		gen_new_keys();
-		kexinitialise(); /* we've finished with this kex */
-	    TRACE((" -> DATAALLOWED=1"))
-	    ses.dataallowed = 1; /* we can send other packets again now */
-		ses.kexstate.donefirstkex = 1;
-	} else {
-		TRACE(("RECVNEWKEYS=1"))
-		ses.kexstate.recvnewkeys = 1;
-	}
+	ses.kexstate.recvnewkeys = 1;
+	switch_keys();
 	
 	TRACE(("leave recv_msg_newkeys"))
 }
@@ -236,11 +258,13 @@
 	ses.kexstate.sentnewkeys = 0;
 
 	/* first_packet_follows */
-	ses.kexstate.firstfollows = 0;
+	ses.kexstate.them_firstfollows = 0;
 
 	ses.kexstate.datatrans = 0;
 	ses.kexstate.datarecv = 0;
 
+	ses.kexstate.our_first_follows_matches = 0;
+
 	ses.kexstate.lastkextime = time(NULL);
 
 }
@@ -281,8 +305,7 @@
  * ses.newkeys is the new set of keys which are generated, these are only
  * taken into use after both sides have sent a newkeys message */
 
-/* Originally from kex.c, generalized for cli/svr mode --mihnea */
-void gen_new_keys() {
+static void gen_new_keys() {
 
 	unsigned char C2S_IV[MAX_IV_LEN];
 	unsigned char C2S_key[MAX_KEY_LEN];
@@ -366,15 +389,9 @@
 		ses.newkeys->recv.hash_index = find_hash(ses.newkeys->recv.algo_mac->hashdesc->name);
 	}
 
-#ifndef DISABLE_ZLIB
-	gen_new_zstreams();
-#endif
-	
-	/* Switch over to the new keys */
-	m_burn(ses.keys, sizeof(struct key_context));
-	m_free(ses.keys);
-	ses.keys = ses.newkeys;
-	ses.newkeys = NULL;
+	/* Ready to switch over */
+	ses.newkeys->trans.valid = 1;
+	ses.newkeys->recv.valid = 1;
 
 	m_burn(C2S_IV, sizeof(C2S_IV));
 	m_burn(C2S_key, sizeof(C2S_key));
@@ -400,7 +417,7 @@
 
 /* Set up new zlib compression streams, close the old ones. Only
  * called from gen_new_keys() */
-static void gen_new_zstreams() {
+static void gen_new_zstream_recv() {
 
 	/* create new zstreams */
 	if (ses.newkeys->recv.algo_comp == DROPBEAR_COMP_ZLIB
@@ -415,6 +432,17 @@
 	} else {
 		ses.newkeys->recv.zstream = NULL;
 	}
+	/* clean up old keys */
+	if (ses.keys->recv.zstream != NULL) {
+		if (inflateEnd(ses.keys->recv.zstream) == Z_STREAM_ERROR) {
+			/* Z_DATA_ERROR is ok, just means that stream isn't ended */
+			dropbear_exit("Crypto error");
+		}
+		m_free(ses.keys->recv.zstream);
+	}
+}
+
+static void gen_new_zstream_trans() {
 
 	if (ses.newkeys->trans.algo_comp == DROPBEAR_COMP_ZLIB
 			|| ses.newkeys->trans.algo_comp == DROPBEAR_COMP_ZLIB_DELAY) {
@@ -432,14 +460,6 @@
 		ses.newkeys->trans.zstream = NULL;
 	}
 
-	/* clean up old keys */
-	if (ses.keys->recv.zstream != NULL) {
-		if (inflateEnd(ses.keys->recv.zstream) == Z_STREAM_ERROR) {
-			/* Z_DATA_ERROR is ok, just means that stream isn't ended */
-			dropbear_exit("Crypto error");
-		}
-		m_free(ses.keys->recv.zstream);
-	}
 	if (ses.keys->trans.zstream != NULL) {
 		if (deflateEnd(ses.keys->trans.zstream) == Z_STREAM_ERROR) {
 			/* Z_DATA_ERROR is ok, just means that stream isn't ended */
@@ -555,7 +575,7 @@
 	DEF_MP_INT(dh_q);
 	DEF_MP_INT(dh_g);
 
-	TRACE(("enter send_msg_kexdh_reply"))
+	TRACE(("enter gen_kexdh_vals"))
 	
 	m_mp_init_multi(&dh_g, &dh_p, &dh_q, NULL);
 
@@ -678,20 +698,27 @@
 
 	buf_incrpos(ses.payload, 16); /* start after the cookie */
 
-	ses.newkeys = (struct key_context*)m_malloc(sizeof(struct key_context));
+	memset(ses.newkeys, 0x0, sizeof(*ses.newkeys));
+
+#ifdef USE_KEXGUESS2
+	enum kexguess2_used kexguess2 = KEXGUESS2_LOOK;
+#else
+	enum kexguess2_used kexguess2 = KEXGUESS2_NO;
+#endif
 
 	/* kex_algorithms */
-	algo = ses.buf_match_algo(ses.payload, sshkex, &goodguess);
+	algo = buf_match_algo(ses.payload, sshkex, &kexguess2, &goodguess);
 	allgood &= goodguess;
-	if (algo == NULL) {
+	if (algo == NULL || algo->val == KEXGUESS2_ALGO_ID) {
 		erralgo = "kex";
 		goto error;
 	}
+	TRACE(("kexguess2 %d", kexguess2))
 	TRACE(("kex algo %s", algo->name))
 	ses.newkeys->algo_kex = algo->val;
 
 	/* server_host_key_algorithms */
-	algo = ses.buf_match_algo(ses.payload, sshhostkey, &goodguess);
+	algo = buf_match_algo(ses.payload, sshhostkey, &kexguess2, &goodguess);
 	allgood &= goodguess;
 	if (algo == NULL) {
 		erralgo = "hostkey";
@@ -701,7 +728,7 @@
 	ses.newkeys->algo_hostkey = algo->val;
 
 	/* encryption_algorithms_client_to_server */
-	c2s_cipher_algo = ses.buf_match_algo(ses.payload, sshciphers, &goodguess);
+	c2s_cipher_algo = buf_match_algo(ses.payload, sshciphers, NULL, NULL);
 	if (c2s_cipher_algo == NULL) {
 		erralgo = "enc c->s";
 		goto error;
@@ -709,7 +736,7 @@
 	TRACE(("enc c2s is  %s", c2s_cipher_algo->name))
 
 	/* encryption_algorithms_server_to_client */
-	s2c_cipher_algo = ses.buf_match_algo(ses.payload, sshciphers, &goodguess);
+	s2c_cipher_algo = buf_match_algo(ses.payload, sshciphers, NULL, NULL);
 	if (s2c_cipher_algo == NULL) {
 		erralgo = "enc s->c";
 		goto error;
@@ -717,7 +744,7 @@
 	TRACE(("enc s2c is  %s", s2c_cipher_algo->name))
 
 	/* mac_algorithms_client_to_server */
-	c2s_hash_algo = ses.buf_match_algo(ses.payload, sshhashes, &goodguess);
+	c2s_hash_algo = buf_match_algo(ses.payload, sshhashes, NULL, NULL);
 	if (c2s_hash_algo == NULL) {
 		erralgo = "mac c->s";
 		goto error;
@@ -725,7 +752,7 @@
 	TRACE(("hash c2s is  %s", c2s_hash_algo->name))
 
 	/* mac_algorithms_server_to_client */
-	s2c_hash_algo = ses.buf_match_algo(ses.payload, sshhashes, &goodguess);
+	s2c_hash_algo = buf_match_algo(ses.payload, sshhashes, NULL, NULL);
 	if (s2c_hash_algo == NULL) {
 		erralgo = "mac s->c";
 		goto error;
@@ -733,7 +760,7 @@
 	TRACE(("hash s2c is  %s", s2c_hash_algo->name))
 
 	/* compression_algorithms_client_to_server */
-	c2s_comp_algo = ses.buf_match_algo(ses.payload, ses.compress_algos, &goodguess);
+	c2s_comp_algo = buf_match_algo(ses.payload, ses.compress_algos, NULL, NULL);
 	if (c2s_comp_algo == NULL) {
 		erralgo = "comp c->s";
 		goto error;
@@ -741,7 +768,7 @@
 	TRACE(("hash c2s is  %s", c2s_comp_algo->name))
 
 	/* compression_algorithms_server_to_client */
-	s2c_comp_algo = ses.buf_match_algo(ses.payload, ses.compress_algos, &goodguess);
+	s2c_comp_algo = buf_match_algo(ses.payload, ses.compress_algos, NULL, NULL);
 	if (s2c_comp_algo == NULL) {
 		erralgo = "comp s->c";
 		goto error;
@@ -754,9 +781,10 @@
 	/* languages_server_to_client */
 	buf_eatstring(ses.payload);
 
-	/* first_kex_packet_follows */
+	/* their first_kex_packet_follows */
 	if (buf_getbool(ses.payload)) {
-		ses.kexstate.firstfollows = 1;
+		TRACE(("them kex firstfollows. allgood %d", allgood))
+		ses.kexstate.them_firstfollows = 1;
 		/* if the guess wasn't good, we ignore the packet sent */
 		if (!allgood) {
 			ses.ignorenext = 1;
@@ -799,6 +827,11 @@
 
 	/* reserved for future extensions */
 	buf_getint(ses.payload);
+
+	if (ses.send_kex_first_guess && allgood) {
+		TRACE(("our_first_follows_matches 1"))
+		ses.kexstate.our_first_follows_matches = 1;
+	}
 	return;
 
 error: