From d7d7c8a352c6b54cc78f2bd930c5e0bcb8545195 Mon Sep 17 00:00:00 2001
From: Vic Lee <llyzs@163.com>
Date: Fri, 16 Sep 2011 23:16:16 +0800
Subject: [PATCH] libfreerdp-core: add more checking in nego phase.

---
 libfreerdp-core/connection.c |  3 +-
 libfreerdp-core/nego.c       | 64 +++++++++++++++++++++++++++---------
 libfreerdp-core/nego.h       |  8 ++---
 3 files changed, 54 insertions(+), 21 deletions(-)

diff --git a/libfreerdp-core/connection.c b/libfreerdp-core/connection.c
index 6552e7e6a..511b46658 100644
--- a/libfreerdp-core/connection.c
+++ b/libfreerdp-core/connection.c
@@ -430,7 +430,8 @@ boolean rdp_server_accept_nego(rdpRdp* rdp, STREAM* s)
 	}
 	printf("\n");
 
-	nego_send_negotiation_response(rdp->nego);
+	if (!nego_send_negotiation_response(rdp->nego))
+		return False;
 
 	ret = False;
 	if (rdp->nego->selected_protocol & PROTOCOL_NLA)
diff --git a/libfreerdp-core/nego.c b/libfreerdp-core/nego.c
index 5c5168d64..452205cad 100644
--- a/libfreerdp-core/nego.c
+++ b/libfreerdp-core/nego.c
@@ -145,9 +145,17 @@ void nego_attempt_nla(rdpNego* nego)
 		return;
 	}
 
-	nego_send_negotiation_request(nego);
+	if (!nego_send_negotiation_request(nego))
+	{
+		nego->state = NEGO_STATE_FAIL;
+		return;
+	}
 
-	nego_recv_response(nego);
+	if (!nego_recv_response(nego))
+	{
+		nego->state = NEGO_STATE_FAIL;
+		return;
+	}
 
 	if (nego->state != NEGO_STATE_FINAL)
 	{
@@ -179,9 +187,17 @@ void nego_attempt_tls(rdpNego* nego)
 		return;
 	}
 
-	nego_send_negotiation_request(nego);
+	if (!nego_send_negotiation_request(nego))
+	{
+		nego->state = NEGO_STATE_FAIL;
+		return;
+	}
 
-	nego_recv_response(nego);
+	if (!nego_recv_response(nego))
+	{
+		nego->state = NEGO_STATE_FAIL;
+		return;
+	}
 
 	if (nego->state != NEGO_STATE_FINAL)
 	{
@@ -211,9 +227,17 @@ void nego_attempt_rdp(rdpNego* nego)
 		return;
 	}
 
-	nego_send_negotiation_request(nego);
+	if (!nego_send_negotiation_request(nego))
+	{
+		nego->state = NEGO_STATE_FAIL;
+		return;
+	}
 
-	nego_recv_response(nego);
+	if (!nego_recv_response(nego))
+	{
+		nego->state = NEGO_STATE_FAIL;
+		return;
+	}
 }
 
 /**
@@ -221,11 +245,12 @@ void nego_attempt_rdp(rdpNego* nego)
  * @param nego
  */
 
-void nego_recv_response(rdpNego* nego)
+boolean nego_recv_response(rdpNego* nego)
 {
 	STREAM* s = transport_recv_stream_init(nego->transport, 1024);
-	transport_read(nego->transport, s);
-	nego_recv(nego->transport, s, nego->transport->recv_extra);
+	if (transport_read(nego->transport, s) < 0)
+		return False;
+	return nego_recv(nego->transport, s, nego->transport->recv_extra);
 }
 
 /**
@@ -236,13 +261,14 @@ void nego_recv_response(rdpNego* nego)
  * @param extra nego pointer
  */
 
-int nego_recv(rdpTransport* transport, STREAM* s, void* extra)
+boolean nego_recv(rdpTransport* transport, STREAM* s, void* extra)
 {
 	uint8 li;
 	uint8 type;
 	rdpNego* nego = (rdpNego*) extra;
 
-	tpkt_read_header(s);
+	if (tpkt_read_header(s) == 0)
+		return False;
 	li = tpdu_read_connection_confirm(s);
 
 	if (li > 6)
@@ -267,7 +293,7 @@ int nego_recv(rdpTransport* transport, STREAM* s, void* extra)
 		nego->state = NEGO_STATE_FINAL;
 	}
 
-	return 0;
+	return True;
 }
 
 /**
@@ -348,7 +374,7 @@ void nego_send(rdpNego* nego)
  * @param nego
  */
 
-void nego_send_negotiation_request(rdpNego* nego)
+boolean nego_send_negotiation_request(rdpNego* nego)
 {
 	STREAM* s;
 	int length;
@@ -390,7 +416,10 @@ void nego_send_negotiation_request(rdpNego* nego)
 	tpdu_write_connection_request(s, length - 5);
 	stream_set_mark(s, em);
 
-	transport_write(nego->transport, s);
+	if (transport_write(nego->transport, s) < 0)
+		return False;
+
+	return True;
 }
 
 /**
@@ -481,7 +510,7 @@ void nego_process_negotiation_failure(rdpNego* nego, STREAM* s)
  * @param nego
  */
 
-void nego_send_negotiation_response(rdpNego* nego)
+boolean nego_send_negotiation_response(rdpNego* nego)
 {
 	STREAM* s;
 	int length;
@@ -508,11 +537,14 @@ void nego_send_negotiation_response(rdpNego* nego)
 	tpdu_write_connection_confirm(s, length - 5);
 	stream_set_mark(s, em);
 
-	transport_write(nego->transport, s);
+	if (transport_write(nego->transport, s) < 0)
+		return False;
 
 	/* update settings with negotiated protocol security */
 	nego->transport->settings->requested_protocols = nego->requested_protocols;
 	nego->transport->settings->selected_protocol = nego->selected_protocol;
+
+	return True;
 }
 
 /**
diff --git a/libfreerdp-core/nego.h b/libfreerdp-core/nego.h
index c1708c14c..005d3fb2e 100644
--- a/libfreerdp-core/nego.h
+++ b/libfreerdp-core/nego.h
@@ -93,15 +93,15 @@ void nego_attempt_tls(rdpNego* nego);
 void nego_attempt_rdp(rdpNego* nego);
 
 void nego_send(rdpNego* nego);
-int nego_recv(rdpTransport* transport, STREAM* s, void* extra);
-void nego_recv_response(rdpNego* nego);
+boolean nego_recv(rdpTransport* transport, STREAM* s, void* extra);
+boolean nego_recv_response(rdpNego* nego);
 boolean nego_read_request(rdpNego* nego, STREAM* s);
 
-void nego_send_negotiation_request(rdpNego* nego);
+boolean nego_send_negotiation_request(rdpNego* nego);
 void nego_process_negotiation_request(rdpNego* nego, STREAM* s);
 void nego_process_negotiation_response(rdpNego* nego, STREAM* s);
 void nego_process_negotiation_failure(rdpNego* nego, STREAM* s);
-void nego_send_negotiation_response(rdpNego* nego);
+boolean nego_send_negotiation_response(rdpNego* nego);
 
 rdpNego* nego_new(struct rdp_transport * transport);
 void nego_free(rdpNego* nego);