From 3e340f2f200861dafea480e6bf2d96b1e44f0202 Mon Sep 17 00:00:00 2001
From: Pierre Ossman <ossman@cendio.se>
Date: Tue, 9 Apr 2019 12:20:27 +0200
Subject: [PATCH] Add explicit STREAM allocation function

Avoids mistakes by making sure everyone allocates these the same
way.

The smart card code still has manual allocation because it has it's
own magical memory management.
---
 cssp.c              |  7 +------
 mcs.c               |  5 +----
 rdp5.c              |  6 +-----
 scard.c             |  1 +
 secure.c            | 11 +++++------
 stream.c            | 12 ++++++++++++
 stream.h            |  5 +++++
 tests/asn_test.c    | 24 +++++++-----------------
 tests/mcs_test.c    | 13 ++-----------
 tests/resize_test.c | 11 -----------
 10 files changed, 35 insertions(+), 60 deletions(-)

diff --git a/cssp.c b/cssp.c
index 7e4f611..98237d8 100644
--- a/cssp.c
+++ b/cssp.c
@@ -36,12 +36,7 @@ ber_wrap_hdr_data(int tagval, STREAM in)
 	STREAM out;
 	int size = s_length(in) + 16;
 
-	out = xmalloc(sizeof(struct stream));
-	memset(out, 0, sizeof(struct stream));
-	out->data = xmalloc(size);
-	out->size = size;
-	out->p = out->data;
-
+	out = s_alloc(size);
 	ber_out_header(out, tagval, s_length(in));
 	out_uint8p(out, in->data, s_length(in));
 	s_mark_end(out);
diff --git a/mcs.c b/mcs.c
index ed466a0..67ec178 100644
--- a/mcs.c
+++ b/mcs.c
@@ -273,10 +273,7 @@ mcs_send_dpu(unsigned short reason)
 
 	logger(Protocol, Debug, "mcs_send_dpu(), reason=%d", reason);
 
-	contents = malloc(sizeof(struct stream));
-	memset(contents, 0, sizeof(struct stream));
-	s_realloc(contents, 6);
-	s_reset(contents);
+	contents = s_alloc(6);
 	ber_out_integer(contents, reason);	/* Reason */
 	ber_out_sequence(contents, NULL);	/* SEQUENCE OF NonStandradParameters OPTIONAL */
 	s_mark_end(contents);
diff --git a/rdp5.c b/rdp5.c
index 5fce5f3..9e8790b 100644
--- a/rdp5.c
+++ b/rdp5.c
@@ -132,11 +132,7 @@ process_ts_fp_updates(STREAM s)
 		{
 			if (assembled[code] == NULL)
 			{
-				assembled[code] = xmalloc(sizeof(struct stream));
-				memset(assembled[code], 0, sizeof(struct stream));
-				s_realloc(assembled[code],
-					  RDESKTOP_FASTPATH_MULTIFRAGMENT_MAX_SIZE);
-				s_reset(assembled[code]);
+				assembled[code] = s_alloc(RDESKTOP_FASTPATH_MULTIFRAGMENT_MAX_SIZE);
 			}
 
 			if (frag == FASTPATH_FRAGMENT_FIRST)
diff --git a/scard.c b/scard.c
index 6178d6e..52d94b9 100644
--- a/scard.c
+++ b/scard.c
@@ -2438,6 +2438,7 @@ scard_device_control(RD_NTHANDLE handle, uint32 request, STREAM in, STREAM out,
 static STREAM
 duplicateStream(PMEM_HANDLE * handle, STREAM s, uint32 buffer_size, RD_BOOL isInputStream)
 {
+	// FIXME: Shouldn't be allocating streams manually
 	STREAM d = SC_xmalloc(handle, sizeof(struct stream));
 	if (d != NULL)
 	{
diff --git a/secure.c b/secure.c
index 16755ff..980d5b2 100644
--- a/secure.c
+++ b/secure.c
@@ -976,25 +976,24 @@ RD_BOOL
 sec_connect(char *server, char *username, char *domain, char *password, RD_BOOL reconnect)
 {
 	uint32 selected_proto;
-	struct stream mcs_data;
+	STREAM mcs_data;
 
 	/* Start a MCS connect sequence */
 	if (!mcs_connect_start(server, username, domain, password, reconnect, &selected_proto))
 		return False;
 
 	/* We exchange some RDP data during the MCS-Connect */
-	mcs_data.size = 512;
-	mcs_data.p = mcs_data.data = (uint8 *) xmalloc(mcs_data.size);
-	sec_out_mcs_connect_initial_pdu(&mcs_data, selected_proto);
+	mcs_data = s_alloc(512);
+	sec_out_mcs_connect_initial_pdu(mcs_data, selected_proto);
 
 	/* finalize the MCS connect sequence */
-	if (!mcs_connect_finalize(&mcs_data))
+	if (!mcs_connect_finalize(mcs_data))
 		return False;
 
 	/* sec_process_mcs_data(&mcs_data); */
 	if (g_encryption)
 		sec_establish_key();
-	xfree(mcs_data.data);
+	s_free(mcs_data);
 	return True;
 }
 
diff --git a/stream.c b/stream.c
index f66158f..8400446 100644
--- a/stream.c
+++ b/stream.c
@@ -25,6 +25,18 @@
 
 extern char g_codepage[16];
 
+STREAM
+s_alloc(unsigned int size)
+{
+	STREAM s;
+
+	s = xmalloc(sizeof(struct stream));
+	memset(s, 0, sizeof(struct stream));
+	s_realloc(s, size);
+
+	return s;
+}
+
 void
 s_realloc(STREAM s, unsigned int size)
 {
diff --git a/stream.h b/stream.h
index 228b648..ec61f0e 100644
--- a/stream.h
+++ b/stream.h
@@ -41,8 +41,13 @@ typedef struct stream
 }
  *STREAM;
 
+/* Return a newly allocated STREAM object of the specified size */
+STREAM s_alloc(unsigned int size);
+/* Resize an existing STREAM object, keeping all data and offsets intact */
 void s_realloc(STREAM s, unsigned int size);
+/* Free STREAM object and its associated buffer */
 void s_free(STREAM s);
+/* Reset all internal offsets, but keep the allocated size */
 void s_reset(STREAM s);
 
 void out_utf16s(STREAM s, const char *string);
diff --git a/tests/asn_test.c b/tests/asn_test.c
index ff47d93..9546e59 100644
--- a/tests/asn_test.c
+++ b/tests/asn_test.c
@@ -46,23 +46,13 @@ xfree(void *mem)
 	free(mem);
 }
 
-static struct stream  *stream_new(size_t size) {
-  struct stream  *s;
-  s = malloc(sizeof(struct stream));
-  memset(s, 0, sizeof(struct stream));
-  s_realloc(s, size);
-  s_reset(s);
-  return(s);
-}
-
-
 Ensure(ASN1, can_create_empty_sequence)
 {
   struct stream *s, *empty;
   uint8_t expected_data[] = {0x30, 0x00};
 
-  s = stream_new(100);
-  empty = stream_new(100);
+  s = s_alloc(100);
+  empty = s_alloc(100);
 
   ber_out_sequence(s, empty);
   s_mark_end(s);
@@ -76,7 +66,7 @@ Ensure(ASN1, can_create_empty_sequence_using_null)
   struct stream *s;
   uint8_t expected_data[] = {0x30, 0x00};
 
-  s = stream_new(100);
+  s = s_alloc(100);
 
   ber_out_sequence(s, NULL);
   s_mark_end(s);
@@ -90,8 +80,8 @@ Ensure(ASN1, can_create_sequence_of_two_integers)
   struct stream *s, *content;
   uint8_t expected_data[] = {0x30, 0x08, 0x02, 0x02, 0x00, 0xbe, 0x02, 0x02, 0x00, 0xef};
 
-  s = stream_new(100);
-  content = stream_new(100);
+  s = s_alloc(100);
+  content = s_alloc(100);
 
   ber_out_integer(content, 0xbe);
   ber_out_integer(content, 0xef);
@@ -109,8 +99,8 @@ Ensure(ASN1, can_create_sequence_of_one_integer)
   struct stream *s, *content;
   uint8_t expected_data[] = {0x30, 0x04, 0x02, 0x02, 0x00, 0xbe};
 
-  s = stream_new(100);
-  content = stream_new(100);
+  s = s_alloc(100);
+  content = s_alloc(100);
 
   ber_out_integer(content, 0xbe);
   s_mark_end(content);
diff --git a/tests/mcs_test.c b/tests/mcs_test.c
index 37d6ff5..bad86be 100644
--- a/tests/mcs_test.c
+++ b/tests/mcs_test.c
@@ -52,15 +52,6 @@ xfree(void *mem)
 	free(mem);
 }
 
-static struct stream  *stream_new(size_t size) {
-  struct stream  *s;
-  s = malloc(sizeof(struct stream));
-  memset(s, 0, sizeof(struct stream));
-  s_realloc(s, size);
-  s_reset(s);
-  return(s);
-}
-
 
 /* Test function */
 Ensure(MCS, should_produce_valid_packet_for_McsSendCJrq)
@@ -69,7 +60,7 @@ Ensure(MCS, should_produce_valid_packet_for_McsSendCJrq)
   uint8_t content[] = {0x38, 0x00, 0x2A, 0x00, 0x0D};
 
   struct stream *s;
-  s = stream_new(5);
+  s = s_alloc(5);
 
   chan_id = 13;
   g_mcs_userid = 42;
@@ -89,7 +80,7 @@ Ensure(MCS, should_produce_valid_packet_for_McsSendDPU)
   struct stream *s;
   uint8_t content[] = {0x30, 0x06, 0x02, 0x02, 0x00, reason, 0x30, 0x00};
 
-  s = stream_new(8);
+  s = s_alloc(8);
 
   expect(logger);
   expect(iso_init, will_return(s));
diff --git a/tests/resize_test.c b/tests/resize_test.c
index 973d1fc..880e4c4 100644
--- a/tests/resize_test.c
+++ b/tests/resize_test.c
@@ -412,17 +412,6 @@ Ensure(Resize, UsingRDPEDISPHonoursServerSessionWidthConstraintMustBeEven)
 	free(s.data);
 }
 
-/* FIXME: promote to actual function in stream.c */
-STREAM s_alloc(size_t capacity)
-{
-	STREAM s;
-	s = xmalloc(sizeof(struct stream));
-	memset(s, 0, sizeof(struct stream));
-	s_realloc(s, capacity);
-	s_reset(s);
-	return s;
-}
-
 void get_width_and_height_from_mcs_connect_initial(int *width, int *height)
 {
 	STREAM s;