From 3100eb823812c76d63589823805e58fb4cf3cadd Mon Sep 17 00:00:00 2001
From: Armin Novak <armin.novak@thincast.com>
Date: Mon, 10 Oct 2022 11:22:20 +0200
Subject: [PATCH] Add return values to TPDU functions

---
 libfreerdp/core/tpdu.c | 26 ++++++++++++++++----------
 libfreerdp/core/tpdu.h |  8 ++++----
 2 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/libfreerdp/core/tpdu.c b/libfreerdp/core/tpdu.c
index d29ac8456..6b3174fe1 100644
--- a/libfreerdp/core/tpdu.c
+++ b/libfreerdp/core/tpdu.c
@@ -63,7 +63,7 @@
  * |         ...        |
  */
 
-static void tpdu_write_header(wStream* s, UINT16 length, BYTE code);
+static BOOL tpdu_write_header(wStream* s, UINT16 length, BYTE code);
 
 /**
  * Read TPDU header.
@@ -114,8 +114,11 @@ BOOL tpdu_read_header(wStream* s, BYTE* code, BYTE* li, UINT16 tpktlength)
  * @param code TPDU code
  */
 
-void tpdu_write_header(wStream* s, UINT16 length, BYTE code)
+BOOL tpdu_write_header(wStream* s, UINT16 length, BYTE code)
 {
+	if (Stream_GetRemainingCapacity(s) < 3)
+		return FALSE;
+
 	Stream_Write_UINT8(s, length); /* LI */
 	Stream_Write_UINT8(s, code);   /* code */
 
@@ -125,10 +128,13 @@ void tpdu_write_header(wStream* s, UINT16 length, BYTE code)
 	}
 	else
 	{
+		if (Stream_GetRemainingCapacity(s) < 5)
+			return FALSE;
 		Stream_Write_UINT16(s, 0); /* DST-REF */
 		Stream_Write_UINT16(s, 0); /* SRC-REF */
 		Stream_Write_UINT8(s, 0);  /* Class 0 */
 	}
+	return TRUE;
 }
 
 /**
@@ -159,9 +165,9 @@ BOOL tpdu_read_connection_request(wStream* s, BYTE* li, UINT16 tpktlength)
  * @param length TPDU length
  */
 
-void tpdu_write_connection_request(wStream* s, UINT16 length)
+BOOL tpdu_write_connection_request(wStream* s, UINT16 length)
 {
-	tpdu_write_header(s, length, X224_TPDU_CONNECTION_REQUEST);
+	return tpdu_write_header(s, length, X224_TPDU_CONNECTION_REQUEST);
 }
 
 /**
@@ -207,9 +213,9 @@ BOOL tpdu_read_connection_confirm(wStream* s, BYTE* li, UINT16 tpktlength)
  * @param length TPDU length
  */
 
-void tpdu_write_connection_confirm(wStream* s, UINT16 length)
+BOOL tpdu_write_connection_confirm(wStream* s, UINT16 length)
 {
-	tpdu_write_header(s, length, X224_TPDU_CONNECTION_CONFIRM);
+	return tpdu_write_header(s, length, X224_TPDU_CONNECTION_CONFIRM);
 }
 
 /**
@@ -218,9 +224,9 @@ void tpdu_write_connection_confirm(wStream* s, UINT16 length)
  * @param length TPDU length
  */
 
-void tpdu_write_disconnect_request(wStream* s, UINT16 length)
+BOOL tpdu_write_disconnect_request(wStream* s, UINT16 length)
 {
-	tpdu_write_header(s, length, X224_TPDU_DISCONNECT_REQUEST);
+	return tpdu_write_header(s, length, X224_TPDU_DISCONNECT_REQUEST);
 }
 
 /**
@@ -228,9 +234,9 @@ void tpdu_write_disconnect_request(wStream* s, UINT16 length)
  * @param s stream
  */
 
-void tpdu_write_data(wStream* s)
+BOOL tpdu_write_data(wStream* s)
 {
-	tpdu_write_header(s, 2, X224_TPDU_DATA);
+	return tpdu_write_header(s, 2, X224_TPDU_DATA);
 }
 
 /**
diff --git a/libfreerdp/core/tpdu.h b/libfreerdp/core/tpdu.h
index daccbeff9..f6298e228 100644
--- a/libfreerdp/core/tpdu.h
+++ b/libfreerdp/core/tpdu.h
@@ -47,11 +47,11 @@ enum X224_TPDU_TYPE
 const char* tpdu_type_to_string(int type);
 FREERDP_LOCAL BOOL tpdu_read_header(wStream* s, BYTE* code, BYTE* li, UINT16 tpktlength);
 FREERDP_LOCAL BOOL tpdu_read_connection_request(wStream* s, BYTE* li, UINT16 tpktlength);
-FREERDP_LOCAL void tpdu_write_connection_request(wStream* s, UINT16 length);
+FREERDP_LOCAL BOOL tpdu_write_connection_request(wStream* s, UINT16 length);
 FREERDP_LOCAL BOOL tpdu_read_connection_confirm(wStream* s, BYTE* li, UINT16 tpktlength);
-FREERDP_LOCAL void tpdu_write_connection_confirm(wStream* s, UINT16 length);
-FREERDP_LOCAL void tpdu_write_disconnect_request(wStream* s, UINT16 length);
+FREERDP_LOCAL BOOL tpdu_write_connection_confirm(wStream* s, UINT16 length);
+FREERDP_LOCAL BOOL tpdu_write_disconnect_request(wStream* s, UINT16 length);
 FREERDP_LOCAL BOOL tpdu_read_data(wStream* s, UINT16* li, UINT16 tpktlength);
-FREERDP_LOCAL void tpdu_write_data(wStream* s);
+FREERDP_LOCAL BOOL tpdu_write_data(wStream* s);
 
 #endif /* FREERDP_LIB_CORE_TPDU_H */