stream: Stream_Ensure*Capacity: change return type

Change the return type of Stream_Ensure*Capacity from void to BOOL to be
able to detect realloc problems easily. Otherwise the only way to detect
this was to check if the capacity after the call was >= the required
size.
In case Stream_Ensure*Capacity fails the old memory is still available
and need to freed outside.

This commit also adds checks to most calls of Stream_Ensure*Capacity to
check if the call was successful.

channels/audin/client/audin_main.c

@@ -176,7 +176,14 @@ static int audin_process_formats(IWTSVirtualChannelCallback* pChannelCallback, w
 			/* Store the agreed format in the corresponding index */
 			callback->formats[callback->formats_count++] = format;
 			/* Put the format to output buffer */
-			Stream_EnsureRemainingCapacity(out, 18 + format.cbSize);
+			if(!Stream_EnsureRemainingCapacity(out, 18 + format.cbSize))
+			{
+				free(callback->formats);
+				Stream_Free(out, TRUE);
+				return 1;
+			}
+
+
 			Stream_Write(out, fm, 18 + format.cbSize);
 		}
 	}

channels/audin/server/audin.c

@@ -118,7 +118,8 @@ static void audin_server_send_formats(audin_server* audin, wStream* s)
 			audin->context.server_formats[i].nChannels *
 			audin->context.server_formats[i].wBitsPerSample / 8;
 
-		Stream_EnsureRemainingCapacity(s, 18);
+		if(!Stream_EnsureRemainingCapacity(s, 18))
+			return;
 
 		Stream_Write_UINT16(s, audin->context.server_formats[i].wFormatTag);
 		Stream_Write_UINT16(s, audin->context.server_formats[i].nChannels);
@@ -130,7 +131,9 @@ static void audin_server_send_formats(audin_server* audin, wStream* s)
 
 		if (audin->context.server_formats[i].cbSize)
 		{
-			Stream_EnsureRemainingCapacity(s, audin->context.server_formats[i].cbSize);
+			if(!Stream_EnsureRemainingCapacity(s, audin->context.server_formats[i].cbSize))
+				return;
+
 			Stream_Write(s, audin->context.server_formats[i].data,
 				audin->context.server_formats[i].cbSize);
 		}
@@ -350,7 +353,8 @@ static void* audin_server_thread_func(void* arg)
 		WTSVirtualChannelRead(audin->audin_channel, 0, NULL, 0, &BytesReturned);
 		if (BytesReturned < 1)
 			continue;
-		Stream_EnsureRemainingCapacity(s, BytesReturned);
+		if (!Stream_EnsureRemainingCapacity(s, BytesReturned))
+			goto out_capacity;
 		if (WTSVirtualChannelRead(audin->audin_channel, 0, (PCHAR) Stream_Buffer(s),
 			Stream_Capacity(s), &BytesReturned) == FALSE)
 		{
@@ -392,8 +396,8 @@ static void* audin_server_thread_func(void* arg)
 		}
 	}
 
+out_capacity:
 	Stream_Free(s, TRUE);
-
 out:
 	WTSVirtualChannelClose(audin->audin_channel);
 	audin->audin_channel = NULL;

channels/cliprdr/client/cliprdr_main.c

@@ -793,7 +793,13 @@ static void cliprdr_virtual_channel_event_data_received(cliprdrPlugin* cliprdr,
 	}
 
 	data_in = cliprdr->data_in;
-	Stream_EnsureRemainingCapacity(data_in, (int) dataLength);
+	if(!Stream_EnsureRemainingCapacity(data_in, (int) dataLength))
+	{
+		Stream_Free(cliprdr->data_in, TRUE);
+		cliprdr->data_in = NULL;
+		return;
+	}
+
 	Stream_Write(data_in, pData, dataLength);
 
 	if (dataFlags & CHANNEL_FLAG_LAST)

channels/cliprdr/server/cliprdr_main.c

@@ -925,7 +925,8 @@ int cliprdr_server_read(CliprdrServerContext* context)
 		Stream_Read_UINT16(s, header.msgFlags); /* msgFlags (2 bytes) */
 		Stream_Read_UINT32(s, header.dataLen); /* dataLen (4 bytes) */
 
-		Stream_EnsureCapacity(s, (header.dataLen + CLIPRDR_HEADER_LENGTH));
+		if(!Stream_EnsureCapacity(s, (header.dataLen + CLIPRDR_HEADER_LENGTH)))
+			return -1;
 		Stream_SetPosition(s, position);
 
 		if (Stream_GetPosition(s) < (header.dataLen + CLIPRDR_HEADER_LENGTH))

channels/drdynvc/client/drdynvc_main.c

@@ -465,6 +465,8 @@ int dvcman_receive_channel_data_first(IWTSVirtualChannelManager* pChannelMgr, UI
 		Stream_Release(channel->dvc_data);
 
 	channel->dvc_data = StreamPool_Take(channel->dvcman->pool, length);
+	if (!channel->dvc_data)
+		return 1;
 	channel->dvc_data_length = length;
 
 	return 0;
@@ -968,7 +970,13 @@ static void drdynvc_virtual_channel_event_data_received(drdynvcPlugin* drdynvc,
 	}
 
 	data_in = drdynvc->data_in;
-	Stream_EnsureRemainingCapacity(data_in, (int) dataLength);
+	if(!Stream_EnsureRemainingCapacity(data_in, (int) dataLength))
+	{
+		Stream_Free(drdynvc->data_in, TRUE);
+		drdynvc->data_in = NULL;
+		return;
+	}
+
 	Stream_Write(data_in, pData, dataLength);
 
 	if (dataFlags & CHANNEL_FLAG_LAST)

channels/drdynvc/server/drdynvc_main.c

@@ -70,7 +70,8 @@ static void* drdynvc_server_thread(void* arg)
 		WTSVirtualChannelRead(context->priv->ChannelHandle, 0, NULL, 0, &BytesReturned);
 		if (BytesReturned < 1)
 			continue;
-		Stream_EnsureRemainingCapacity(s, BytesReturned);
+		if (!Stream_EnsureRemainingCapacity(s, BytesReturned))
+			break;
 		if (!WTSVirtualChannelRead(context->priv->ChannelHandle, 0,
 			(PCHAR) Stream_Buffer(s), Stream_Capacity(s), &BytesReturned))
 		{

channels/drive/client/drive_file.c

@@ -419,8 +419,9 @@ BOOL drive_file_query_information(DRIVE_FILE* file, UINT32 FsInformationClass, w
 	{
 		case FileBasicInformation:
 			/* http://msdn.microsoft.com/en-us/library/cc232094.aspx */
+			if(!Stream_EnsureRemainingCapacity(output, 4 + 36))
+				goto out_fail;
 			Stream_Write_UINT32(output, 36); /* Length */
-			Stream_EnsureRemainingCapacity(output, 36);
 			Stream_Write_UINT64(output, FILE_TIME_SYSTEM_TO_RDP(st.st_mtime)); /* CreationTime */
 			Stream_Write_UINT64(output, FILE_TIME_SYSTEM_TO_RDP(st.st_atime)); /* LastAccessTime */
 			Stream_Write_UINT64(output, FILE_TIME_SYSTEM_TO_RDP(st.st_mtime)); /* LastWriteTime */
@@ -431,8 +432,9 @@ BOOL drive_file_query_information(DRIVE_FILE* file, UINT32 FsInformationClass, w
 
 		case FileStandardInformation:
 			/*  http://msdn.microsoft.com/en-us/library/cc232088.aspx */
+			if(!Stream_EnsureRemainingCapacity(output, 4 + 22))
+				goto out_fail;
 			Stream_Write_UINT32(output, 22); /* Length */
-			Stream_EnsureRemainingCapacity(output, 22);
 			Stream_Write_UINT64(output, st.st_size); /* AllocationSize */
 			Stream_Write_UINT64(output, st.st_size); /* EndOfFile */
 			Stream_Write_UINT32(output, st.st_nlink); /* NumberOfLinks */
@@ -443,17 +445,23 @@ BOOL drive_file_query_information(DRIVE_FILE* file, UINT32 FsInformationClass, w
 
 		case FileAttributeTagInformation:
 			/* http://msdn.microsoft.com/en-us/library/cc232093.aspx */
+			if(!Stream_EnsureRemainingCapacity(output, 4 + 8))
+				goto out_fail;
 			Stream_Write_UINT32(output, 8); /* Length */
-			Stream_EnsureRemainingCapacity(output, 8);
 			Stream_Write_UINT32(output, FILE_ATTR_SYSTEM_TO_RDP(file, st)); /* FileAttributes */
 			Stream_Write_UINT32(output, 0); /* ReparseTag */
 			break;
 
 		default:
+			/* Unhandled FsInformationClass */
 			Stream_Write_UINT32(output, 0); /* Length */
 			return FALSE;
 	}
 	return TRUE;
+
+out_fail:
+	Stream_Write_UINT32(output, 0); /* Length */
+	return FALSE;
 }
 
 int dir_empty(const char *path)
@@ -680,8 +688,9 @@ BOOL drive_file_query_directory(DRIVE_FILE* file, UINT32 FsInformationClass, BYT
 	{
 		case FileDirectoryInformation:
 			/* http://msdn.microsoft.com/en-us/library/cc232097.aspx */
+			if(!Stream_EnsureRemainingCapacity(output, 4 + 64 + length))
+				goto out_fail;
 			Stream_Write_UINT32(output, 64 + length); /* Length */
-			Stream_EnsureRemainingCapacity(output, 64 + length);
 			Stream_Write_UINT32(output, 0); /* NextEntryOffset */
 			Stream_Write_UINT32(output, 0); /* FileIndex */
 			Stream_Write_UINT64(output, FILE_TIME_SYSTEM_TO_RDP(st.st_mtime)); /* CreationTime */
@@ -697,8 +706,9 @@ BOOL drive_file_query_directory(DRIVE_FILE* file, UINT32 FsInformationClass, BYT
 
 		case FileFullDirectoryInformation:
 			/* http://msdn.microsoft.com/en-us/library/cc232068.aspx */
+			if(!Stream_EnsureRemainingCapacity(output, 4 + 68 + length))
+				goto out_fail;
 			Stream_Write_UINT32(output, 68 + length); /* Length */
-			Stream_EnsureRemainingCapacity(output, 68 + length);
 			Stream_Write_UINT32(output, 0); /* NextEntryOffset */
 			Stream_Write_UINT32(output, 0); /* FileIndex */
 			Stream_Write_UINT64(output, FILE_TIME_SYSTEM_TO_RDP(st.st_mtime)); /* CreationTime */
@@ -715,8 +725,9 @@ BOOL drive_file_query_directory(DRIVE_FILE* file, UINT32 FsInformationClass, BYT
 
 		case FileBothDirectoryInformation:
 			/* http://msdn.microsoft.com/en-us/library/cc232095.aspx */
+			if(!Stream_EnsureRemainingCapacity(output, 4 + 93 + length))
+				goto out_fail;
 			Stream_Write_UINT32(output, 93 + length); /* Length */
-			Stream_EnsureRemainingCapacity(output, 93 + length);
 			Stream_Write_UINT32(output, 0); /* NextEntryOffset */
 			Stream_Write_UINT32(output, 0); /* FileIndex */
 			Stream_Write_UINT64(output, FILE_TIME_SYSTEM_TO_RDP(st.st_mtime)); /* CreationTime */
@@ -736,8 +747,9 @@ BOOL drive_file_query_directory(DRIVE_FILE* file, UINT32 FsInformationClass, BYT
 
 		case FileNamesInformation:
 			/* http://msdn.microsoft.com/en-us/library/cc232077.aspx */
+			if(!Stream_EnsureRemainingCapacity(output, 4 + 12 + length))
+				goto out_fail;
 			Stream_Write_UINT32(output, 12 + length); /* Length */
-			Stream_EnsureRemainingCapacity(output, 12 + length);
 			Stream_Write_UINT32(output, 0); /* NextEntryOffset */
 			Stream_Write_UINT32(output, 0); /* FileIndex */
 			Stream_Write_UINT32(output, length); /* FileNameLength */
@@ -745,6 +757,7 @@ BOOL drive_file_query_directory(DRIVE_FILE* file, UINT32 FsInformationClass, BYT
 			break;
 
 		default:
+			/* Unhandled FsInformationClass */
 			Stream_Write_UINT32(output, 0); /* Length */
 			Stream_Write_UINT8(output, 0); /* Padding */
 			ret = FALSE;
@@ -752,8 +765,13 @@ BOOL drive_file_query_directory(DRIVE_FILE* file, UINT32 FsInformationClass, BYT
 	}
 
 	free(ent_path);
-
 	return ret;
+
+out_fail:
+	free(ent_path);
+	Stream_Write_UINT32(output, 0); /* Length */
+	Stream_Write_UINT8(output, 0); /* Padding */
+	return FALSE;
 }
 
 #ifdef _WIN32

channels/rdpdr/server/rdpdr_main.c

@@ -146,8 +146,11 @@ static int rdpdr_server_write_general_capability_set(RdpdrServerContext* context
 	UINT32 extraFlags1;
 	UINT32 SpecialTypeDeviceCap;
 	RDPDR_CAPABILITY_HEADER header;
-	header.CapabilityType = CAP_GENERAL_TYPE;
 	header.CapabilityLength = RDPDR_CAPABILITY_HEADER_LENGTH + 36;
+	if(!Stream_EnsureRemainingCapacity(s, header.CapabilityLength))
+		return -1;
+
+	header.CapabilityType = CAP_GENERAL_TYPE;
 	header.Version = GENERAL_CAPABILITY_VERSION_02;
 	ioCode1 = 0;
 	ioCode1 |= RDPDR_IRP_MJ_CREATE; /* always set */
@@ -176,7 +179,6 @@ static int rdpdr_server_write_general_capability_set(RdpdrServerContext* context
 	extraFlags1 = 0;
 	extraFlags1 |= ENABLE_ASYNCIO; /* optional */
 	SpecialTypeDeviceCap = 0;
-	Stream_EnsureRemainingCapacity(s, header.CapabilityLength);
 	rdpdr_server_write_capability_set_header(s, &header);
 	Stream_Write_UINT32(s, 0); /* osType (4 bytes), ignored on receipt */
 	Stream_Write_UINT32(s, 0); /* osVersion (4 bytes), unused and must be set to zero */
@@ -202,7 +204,8 @@ static int rdpdr_server_write_printer_capability_set(RdpdrServerContext* context
 	header.CapabilityType = CAP_PRINTER_TYPE;
 	header.CapabilityLength = RDPDR_CAPABILITY_HEADER_LENGTH;
 	header.Version = PRINT_CAPABILITY_VERSION_01;
-	Stream_EnsureRemainingCapacity(s, header.CapabilityLength);
+	if(!Stream_EnsureRemainingCapacity(s, header.CapabilityLength))
+		return -1;
 	rdpdr_server_write_capability_set_header(s, &header);
 	return 0;
 }
@@ -218,7 +221,8 @@ static int rdpdr_server_write_port_capability_set(RdpdrServerContext* context, w
 	header.CapabilityType = CAP_PORT_TYPE;
 	header.CapabilityLength = RDPDR_CAPABILITY_HEADER_LENGTH;
 	header.Version = PORT_CAPABILITY_VERSION_01;
-	Stream_EnsureRemainingCapacity(s, header.CapabilityLength);
+	if(!Stream_EnsureRemainingCapacity(s, header.CapabilityLength))
+		return -1;
 	rdpdr_server_write_capability_set_header(s, &header);
 	return 0;
 }
@@ -234,7 +238,8 @@ static int rdpdr_server_write_drive_capability_set(RdpdrServerContext* context,
 	header.CapabilityType = CAP_DRIVE_TYPE;
 	header.CapabilityLength = RDPDR_CAPABILITY_HEADER_LENGTH;
 	header.Version = DRIVE_CAPABILITY_VERSION_02;
-	Stream_EnsureRemainingCapacity(s, header.CapabilityLength);
+	if(!Stream_EnsureRemainingCapacity(s, header.CapabilityLength))
+		return -1;
 	rdpdr_server_write_capability_set_header(s, &header);
 	return 0;
 }
@@ -250,7 +255,8 @@ static int rdpdr_server_write_smartcard_capability_set(RdpdrServerContext* conte
 	header.CapabilityType = CAP_SMARTCARD_TYPE;
 	header.CapabilityLength = RDPDR_CAPABILITY_HEADER_LENGTH;
 	header.Version = SMARTCARD_CAPABILITY_VERSION_01;
-	Stream_EnsureRemainingCapacity(s, header.CapabilityLength);
+	if(!Stream_EnsureRemainingCapacity(s, header.CapabilityLength))
+		return -1;
 	rdpdr_server_write_capability_set_header(s, &header);
 	return 0;
 }
@@ -267,19 +273,31 @@ static int rdpdr_server_send_core_capability_request(RdpdrServerContext* context
 	header.PacketId = PAKID_CORE_SERVER_CAPABILITY;
 	numCapabilities = 5;
 	s = Stream_New(NULL, RDPDR_HEADER_LENGTH + 512);
+	if (!s)
+		return -1;
 	Stream_Write_UINT16(s, header.Component); /* Component (2 bytes) */
 	Stream_Write_UINT16(s, header.PacketId); /* PacketId (2 bytes) */
 	Stream_Write_UINT16(s, numCapabilities); /* numCapabilities (2 bytes) */
 	Stream_Write_UINT16(s, 0); /* Padding (2 bytes) */
-	rdpdr_server_write_general_capability_set(context, s);
+	if(rdpdr_server_write_general_capability_set(context, s) != 0)
-	rdpdr_server_write_printer_capability_set(context, s);
+		goto out_error;
-	rdpdr_server_write_port_capability_set(context, s);
+	if(rdpdr_server_write_printer_capability_set(context, s) != 0)
-	rdpdr_server_write_drive_capability_set(context, s);
+		goto out_error;
-	rdpdr_server_write_smartcard_capability_set(context, s);
+	if(rdpdr_server_write_port_capability_set(context, s) != 0)
+		goto out_error;
+	if(rdpdr_server_write_drive_capability_set(context, s) != 0)
+		goto out_error;
+	if(rdpdr_server_write_smartcard_capability_set(context, s) != 0)
+		goto out_error;
 	Stream_SealLength(s);
 	status = WTSVirtualChannelWrite(context->priv->ChannelHandle, (PCHAR) Stream_Buffer(s), Stream_Length(s), &written);
 	Stream_Free(s, TRUE);
-	return 0;
+	return status ? 0 : -1;
+
+out_error:
+	Stream_Free(s, TRUE);
+	return -1;
+
 }
 
 static int rdpdr_server_receive_core_capability_response(RdpdrServerContext* context, wStream* s, RDPDR_HEADER* header)

channels/rdpei/server/rdpei_main.c

@@ -320,7 +320,11 @@ int rdpei_server_handle_messages(RdpeiServerContext *context) {
 		Stream_SetPosition(s, 0);
 		if (priv->expectedBytes)
 		{
-			Stream_EnsureCapacity(s, priv->expectedBytes);
+			if (!Stream_EnsureCapacity(s, priv->expectedBytes))
+			{
+				fprintf(stderr, "%s: couldn't allocate memory for stream\n", __FUNCTION__);
+				return -1;
+			}
 			return 1;
 		}
 	}
@@ -376,7 +380,12 @@ int rdpei_server_send_sc_ready(RdpeiServerContext *context, UINT32 version)
 	}
 
 	Stream_SetPosition(priv->outputStream, 0);
-	Stream_EnsureCapacity(priv->outputStream, RDPINPUT_HEADER_LENGTH + 4);
+
+	if (!Stream_EnsureCapacity(priv->outputStream, RDPINPUT_HEADER_LENGTH + 4))
+	{
+		fprintf(stderr, "%s: could not re-allocate memory for stream\n", __FUNCTION__);
+		return -1;
+	}
 
 	Stream_Write_UINT16(priv->outputStream, EVENTID_SC_READY);
 	Stream_Write_UINT32(priv->outputStream, RDPINPUT_HEADER_LENGTH + 4);
@@ -411,7 +420,11 @@ int rdpei_server_suspend(RdpeiServerContext *context)
 	}
 
 	Stream_SetPosition(priv->outputStream, 0);
-	Stream_EnsureCapacity(priv->outputStream, RDPINPUT_HEADER_LENGTH);
+	if (!Stream_EnsureCapacity(priv->outputStream, RDPINPUT_HEADER_LENGTH))
+	{
+		fprintf(stderr, "%s: could not re-allocate memory for stream\n", __FUNCTION__);
+		return -1;
+	}
 
 	Stream_Write_UINT16(priv->outputStream, EVENTID_SUSPEND_TOUCH);
 	Stream_Write_UINT32(priv->outputStream, RDPINPUT_HEADER_LENGTH);
@@ -446,7 +459,11 @@ int rdpei_server_resume(RdpeiServerContext *context)
 	}
 
 	Stream_SetPosition(priv->outputStream, 0);
-	Stream_EnsureCapacity(priv->outputStream, RDPINPUT_HEADER_LENGTH);
+	if(!Stream_EnsureCapacity(priv->outputStream, RDPINPUT_HEADER_LENGTH))
+	{
+		fprintf(stderr, "%s: couldn't allocate memory for stream\n", __FUNCTION__);
+		return -1;
+	}
 
 	Stream_Write_UINT16(priv->outputStream, EVENTID_RESUME_TOUCH);
 	Stream_Write_UINT32(priv->outputStream, RDPINPUT_HEADER_LENGTH);

channels/rdpsnd/client/pulse/rdpsnd_pulse.c

@@ -515,7 +515,8 @@ static BYTE* rdpsnd_pulse_convert_audio(rdpsndDevicePlugin* device, BYTE* data,
 				inSize -= 32;
 			}
 
-			Stream_EnsureRemainingCapacity(pulse->gsmBuffer, 160 * 2);
+			if(!Stream_EnsureRemainingCapacity(pulse->gsmBuffer, 160 * 2))
+				return NULL;
 			Stream_Write(pulse->gsmBuffer, (void*) gsmBlockBuffer, 160 * 2);
 		}
 
@@ -544,6 +545,8 @@ static void rdpsnd_pulse_play(rdpsndDevicePlugin* device, BYTE* data, int size)
 		return;
 
 	pcmData = rdpsnd_pulse_convert_audio(device, data, &size);
+	if (pcmData)
+		return;
 
 	pa_threaded_mainloop_lock(pulse->mainloop);
 

channels/rdpsnd/server/rdpsnd_main.c

@@ -364,7 +364,11 @@ static BOOL rdpsnd_server_send_audio_pdu(RdpsndServerContext* context, UINT16 wT
 	Stream_SetPosition(s, 0);
 
 	/* Wave PDU */
-	Stream_EnsureRemainingCapacity(s, size + fill_size);
+	if (!Stream_EnsureRemainingCapacity(s, size + fill_size))
+	{
+		status = -1;
+		goto out;
+	}
 	Stream_Write_UINT32(s, 0); /* bPad */
 	Stream_Write(s, src + 4, size - 4);
 
@@ -662,7 +666,8 @@ int rdpsnd_server_handle_messages(RdpsndServerContext *context)
 		Stream_SetPosition(s, 0);
 		if (priv->expectedBytes)
 		{
-			Stream_EnsureCapacity(s, priv->expectedBytes);
+			if(!Stream_EnsureCapacity(s, priv->expectedBytes))
+				return -1;
 			return 1;
 		}
 	}

channels/tsmf/client/tsmf_ifman.c

@@ -45,7 +45,8 @@ int tsmf_ifman_rim_exchange_capability_request(TSMF_IFMAN* ifman)
 
 	DEBUG_TSMF("server CapabilityValue %d", CapabilityValue);
 
-	Stream_EnsureRemainingCapacity(ifman->output, 8);
+	if(!Stream_EnsureRemainingCapacity(ifman->output, 8))
+		return -1;
 	Stream_Write_UINT32(ifman->output, 1); /* CapabilityValue */
 	Stream_Write_UINT32(ifman->output, 0); /* Result */
 
@@ -61,8 +62,9 @@ int tsmf_ifman_exchange_capability_request(TSMF_IFMAN* ifman)
 	UINT32 cbCapabilityLength;
 	UINT32 numHostCapabilities;
 
+	if(!Stream_EnsureRemainingCapacity(ifman->output, ifman->input_size + 4))
+		return -1;
 	pos = Stream_GetPosition(ifman->output);
-	Stream_EnsureRemainingCapacity(ifman->output, ifman->input_size + 4);
 	Stream_Copy(ifman->output, ifman->input, ifman->input_size);
 	Stream_SetPosition(ifman->output, pos);
 	Stream_Read_UINT32(ifman->output, numHostCapabilities);
@@ -116,7 +118,8 @@ int tsmf_ifman_check_format_support_request(TSMF_IFMAN* ifman)
 	if (FormatSupported)
 		DEBUG_TSMF("format ok.");
 
-	Stream_EnsureRemainingCapacity(ifman->output, 12);
+	if(!Stream_EnsureRemainingCapacity(ifman->output, 12))
+		return -1;
 	Stream_Write_UINT32(ifman->output, FormatSupported);
 	Stream_Write_UINT32(ifman->output, PlatformCookie);
 	Stream_Write_UINT32(ifman->output, 0); /* Result */
@@ -184,7 +187,8 @@ int tsmf_ifman_add_stream(TSMF_IFMAN* ifman)
 int tsmf_ifman_set_topology_request(TSMF_IFMAN* ifman)
 {
 	DEBUG_TSMF("");
-	Stream_EnsureRemainingCapacity(ifman->output, 8);
+	if(!Stream_EnsureRemainingCapacity(ifman->output, 8))
+		return -1;
 	Stream_Write_UINT32(ifman->output, 1); /* TopologyReady */
 	Stream_Write_UINT32(ifman->output, 0); /* Result */
 	ifman->output_interface_id = TSMF_INTERFACE_DEFAULT | STREAM_ID_STUB;
@@ -278,7 +282,8 @@ int tsmf_ifman_shutdown_presentation(TSMF_IFMAN* ifman)
 	else
 		WLog_ERR(TAG, "unknown presentation id");
 
-	Stream_EnsureRemainingCapacity(ifman->output, 4);
+	if(!Stream_EnsureRemainingCapacity(ifman->output, 4))
+		return -1;
 	Stream_Write_UINT32(ifman->output, 0); /* Result */
 	ifman->output_interface_id = TSMF_INTERFACE_DEFAULT | STREAM_ID_STUB;
 
@@ -516,7 +521,8 @@ int tsmf_ifman_on_end_of_stream(TSMF_IFMAN* ifman)
 	}
 
 	DEBUG_TSMF("StreamId %d", StreamId);
-	Stream_EnsureRemainingCapacity(ifman->output, 16);
+	if(!Stream_EnsureRemainingCapacity(ifman->output, 16))
+		return -1;
 	Stream_Write_UINT32(ifman->output, CLIENT_EVENT_NOTIFICATION); /* FunctionId */
 	Stream_Write_UINT32(ifman->output, StreamId); /* StreamId */
 	Stream_Write_UINT32(ifman->output, TSMM_CLIENT_EVENT_ENDOFSTREAM); /* EventId */
@@ -539,7 +545,8 @@ int tsmf_ifman_on_playback_started(TSMF_IFMAN* ifman)
 	else
 		WLog_ERR(TAG, "unknown presentation id");
 
-	Stream_EnsureRemainingCapacity(ifman->output, 16);
+	if(!Stream_EnsureRemainingCapacity(ifman->output, 16))
+		return -1;
 	Stream_Write_UINT32(ifman->output, CLIENT_EVENT_NOTIFICATION); /* FunctionId */
 	Stream_Write_UINT32(ifman->output, 0); /* StreamId */
 	Stream_Write_UINT32(ifman->output, TSMM_CLIENT_EVENT_START_COMPLETED); /* EventId */
@@ -598,7 +605,8 @@ int tsmf_ifman_on_playback_stopped(TSMF_IFMAN* ifman)
 	else
 		WLog_ERR(TAG, "unknown presentation id");
 
-	Stream_EnsureRemainingCapacity(ifman->output, 16);
+	if(!Stream_EnsureRemainingCapacity(ifman->output, 16))
+		return -1;
 	Stream_Write_UINT32(ifman->output, CLIENT_EVENT_NOTIFICATION); /* FunctionId */
 	Stream_Write_UINT32(ifman->output, 0); /* StreamId */
 	Stream_Write_UINT32(ifman->output, TSMM_CLIENT_EVENT_STOP_COMPLETED); /* EventId */
@@ -611,7 +619,8 @@ int tsmf_ifman_on_playback_stopped(TSMF_IFMAN* ifman)
 int tsmf_ifman_on_playback_rate_changed(TSMF_IFMAN* ifman)
 {
 	DEBUG_TSMF("");
-	Stream_EnsureRemainingCapacity(ifman->output, 16);
+	if(!Stream_EnsureRemainingCapacity(ifman->output, 16))
+		return -1;
 	Stream_Write_UINT32(ifman->output, CLIENT_EVENT_NOTIFICATION); /* FunctionId */
 	Stream_Write_UINT32(ifman->output, 0); /* StreamId */
 	Stream_Write_UINT32(ifman->output, TSMM_CLIENT_EVENT_MONITORCHANGED); /* EventId */

libfreerdp/codec/nsc_encode.c

@@ -493,7 +493,8 @@ int nsc_write_message(NSC_CONTEXT* context, wStream* s, NSC_MESSAGE* message)
 	totalPlaneByteCount = message->LumaPlaneByteCount + message->OrangeChromaPlaneByteCount +
 			message->GreenChromaPlaneByteCount + message->AlphaPlaneByteCount;
 
-	Stream_EnsureRemainingCapacity(s, 20 + totalPlaneByteCount);
+	if(!Stream_EnsureRemainingCapacity(s, 20 + totalPlaneByteCount))
+		return -1;
 	Stream_Write_UINT32(s, message->LumaPlaneByteCount); /* LumaPlaneByteCount (4 bytes) */
 	Stream_Write_UINT32(s, message->OrangeChromaPlaneByteCount); /* OrangeChromaPlaneByteCount (4 bytes) */
 	Stream_Write_UINT32(s, message->GreenChromaPlaneByteCount); /* GreenChromaPlaneByteCount (4 bytes) */

libfreerdp/codec/rfx.c

@@ -1110,7 +1110,8 @@ static void rfx_write_message_context(RFX_CONTEXT* context, wStream* s)
 
 void rfx_compose_message_header(RFX_CONTEXT* context, wStream* s)
 {
-	Stream_EnsureRemainingCapacity(s, 12 + 10 + 12 + 13);
+	if(!Stream_EnsureRemainingCapacity(s, 12 + 10 + 12 + 13))
+		return;
 
 	rfx_write_message_sync(context, s);
 	rfx_write_message_context(context, s);
@@ -1128,7 +1129,8 @@ static void rfx_write_tile(RFX_CONTEXT* context, wStream* s, RFX_TILE* tile)
 	UINT32 blockLen;
 
 	blockLen = rfx_tile_length(tile);
-	Stream_EnsureRemainingCapacity(s, blockLen);
+	if(!Stream_EnsureRemainingCapacity(s, blockLen))
+		return;
 
 	Stream_Write_UINT16(s, CBT_TILE); /* BlockT.blockType (2 bytes) */
 	Stream_Write_UINT32(s, blockLen); /* BlockT.blockLen (4 bytes) */
@@ -1514,7 +1516,8 @@ static void rfx_write_message_tileset(RFX_CONTEXT* context, wStream* s, RFX_MESS
 	UINT32* quantVals;
 
 	blockLen = 22 + (message->numQuant * 5) + message->tilesDataSize;
-	Stream_EnsureRemainingCapacity(s, blockLen);
+	if(!Stream_EnsureRemainingCapacity(s, blockLen))
+		return;
 
 	Stream_Write_UINT16(s, WBT_EXTENSION); /* CodecChannelT.blockType (2 bytes) */
 	Stream_Write_UINT32(s, blockLen); /* set CodecChannelT.blockLen (4 bytes) */
@@ -1548,7 +1551,8 @@ static void rfx_write_message_tileset(RFX_CONTEXT* context, wStream* s, RFX_MESS
 
 void rfx_write_message_frame_begin(RFX_CONTEXT* context, wStream* s, RFX_MESSAGE* message)
 {
-	Stream_EnsureRemainingCapacity(s, 14);
+	if(!Stream_EnsureRemainingCapacity(s, 14))
+		return;
 
 	Stream_Write_UINT16(s, WBT_FRAME_BEGIN); /* CodecChannelT.blockType */
 	Stream_Write_UINT32(s, 14); /* CodecChannelT.blockLen */
@@ -1564,7 +1568,8 @@ void rfx_write_message_region(RFX_CONTEXT* context, wStream* s, RFX_MESSAGE* mes
 	UINT32 blockLen;
 
 	blockLen = 15 + (message->numRects * 8);
-	Stream_EnsureRemainingCapacity(s, blockLen);
+	if(!Stream_EnsureRemainingCapacity(s, blockLen))
+		return;
 
 	Stream_Write_UINT16(s, WBT_REGION); /* CodecChannelT.blockType (2 bytes) */
 	Stream_Write_UINT32(s, blockLen); /* set CodecChannelT.blockLen (4 bytes) */
@@ -1589,7 +1594,8 @@ void rfx_write_message_region(RFX_CONTEXT* context, wStream* s, RFX_MESSAGE* mes
 
 void rfx_write_message_frame_end(RFX_CONTEXT* context, wStream* s, RFX_MESSAGE* message)
 {
-	Stream_EnsureRemainingCapacity(s, 8);
+	if(!Stream_EnsureRemainingCapacity(s, 8))
+		return;
 
 	Stream_Write_UINT16(s, WBT_FRAME_END); /* CodecChannelT.blockType */
 	Stream_Write_UINT32(s, 8); /* CodecChannelT.blockLen */

libfreerdp/core/autodetect.c

@@ -153,12 +153,16 @@ BOOL autodetect_send_bandwidth_measure_payload(rdpContext* context, UINT16 paylo
 	/* 4-bytes aligned */
 	payloadLength &= ~3;
 
+	if(!Stream_EnsureRemainingCapacity(s, 8 + payloadLength))
+	{
+		Stream_Release(s);
+		return FALSE;
+	}
 	Stream_Write_UINT8(s, 0x08); /* headerLength (1 byte) */
 	Stream_Write_UINT8(s, TYPE_ID_AUTODETECT_REQUEST); /* headerTypeId (1 byte) */
 	Stream_Write_UINT16(s, sequenceNumber); /* sequenceNumber (2 bytes) */
 	Stream_Write_UINT16(s, RDP_BW_PAYLOAD_REQUEST_TYPE); /* requestType (2 bytes) */
 	Stream_Write_UINT16(s, payloadLength); /* payloadLength (2 bytes) */
-	Stream_EnsureRemainingCapacity(s, payloadLength);
 	/* Random data (better measurement in case the line is compressed) */
 	for (i = 0; i < payloadLength / 4; i++)
 	{
@@ -192,7 +196,11 @@ static BOOL autodetect_send_bandwidth_measure_stop(rdpContext* context, UINT16 p
 		Stream_Write_UINT16(s, payloadLength); /* payloadLength (2 bytes) */
 		if (payloadLength > 0)
 		{
-			Stream_EnsureRemainingCapacity(s, payloadLength);
+			if(!Stream_EnsureRemainingCapacity(s, payloadLength))
+			{
+				Stream_Release(s);
+				return FALSE;
+			}
 			/* Random data (better measurement in case the line is compressed) */
 			for (i = 0; i < payloadLength / 4; i++)
 			{

libfreerdp/core/channels.c

@@ -80,6 +80,8 @@ BOOL freerdp_channel_send(rdpRdp* rdp, UINT16 channelId, BYTE* data, int size)
 	while (left > 0)
 	{
 		s = rdp_send_stream_init(rdp);
+		if (!s)
+			return FALSE;
 
 		if (left > (int) rdp->settings->VirtualChannelChunkSize)
 		{
@@ -98,10 +100,18 @@ BOOL freerdp_channel_send(rdpRdp* rdp, UINT16 channelId, BYTE* data, int size)
 
 		Stream_Write_UINT32(s, size);
 		Stream_Write_UINT32(s, flags);
-		Stream_EnsureCapacity(s, chunkSize);
+		if(!Stream_EnsureCapacity(s, chunkSize))
+		{
+			Stream_Release(s);
+			return FALSE;
+		}
 		Stream_Write(s, data, chunkSize);
 
-		rdp_send(rdp, s, channelId);
+		if(!rdp_send(rdp, s, channelId))
+		{
+			Stream_Release(s);
+			return FALSE;
+		}
 
 		data += chunkSize;
 		left -= chunkSize;

libfreerdp/core/fastpath.c

@@ -460,7 +460,11 @@ static int fastpath_recv_update_data(rdpFastPath* fastpath, wStream* s)
 				return -1;
 			}
 
-			Stream_EnsureCapacity(fastpath->updateData, totalSize);
+			if(!Stream_EnsureCapacity(fastpath->updateData, totalSize))
+			{
+				WLog_ERR(TAG,  "Couldn't re-allocate memory for stream");
+				return -1;
+			}
 
 			Stream_Copy(fastpath->updateData, cs, size);
 		}
@@ -484,7 +488,11 @@ static int fastpath_recv_update_data(rdpFastPath* fastpath, wStream* s)
 				return -1;
 			}
 
-			Stream_EnsureCapacity(fastpath->updateData, totalSize);
+			if(!Stream_EnsureCapacity(fastpath->updateData, totalSize))
+			{
+				WLog_ERR(TAG,  "Couldn't re-allocate memory for stream");
+				return -1;
+			}
 
 			Stream_Copy(fastpath->updateData, cs, size);
 

libfreerdp/core/gateway/http.c

@@ -676,7 +676,8 @@ HttpResponse* http_response_recv(rdpTls* tls)
 
 			if (Stream_GetRemainingLength(s) < 1024)
 			{
-				Stream_EnsureRemainingCapacity(s, 1024);
+				if(!Stream_EnsureRemainingCapacity(s, 1024))
+					goto out_error;
 				buffer = (char*) Stream_Buffer(s);
 				payload = &buffer[payloadOffset];
 			}
@@ -785,7 +786,8 @@ HttpResponse* http_response_recv(rdpTls* tls)
 
 		if (Stream_GetRemainingLength(s) < 1024)
 		{
-			Stream_EnsureRemainingCapacity(s, 1024);
+			if(!Stream_EnsureRemainingCapacity(s, 1024))
+				goto out_error;
 			buffer = (char*) Stream_Buffer(s);
 			payload = &buffer[payloadOffset];
 		}

libfreerdp/core/gateway/rdg.c

@@ -104,7 +104,11 @@ wStream* rdg_receive_packet(rdpRdg* rdg)
 
 	if (Stream_Capacity(s) < packet->packetLength)
 	{
-		Stream_EnsureCapacity(s, packet->packetLength);
+		if(!Stream_EnsureCapacity(s, packet->packetLength))
+		{
+			Stream_Free(s, TRUE);
+			return NULL;
+		}
 		packet = (RdgPacketHeader*) Stream_Buffer(s);
 	}
 

libfreerdp/core/gateway/rpc_client.c

@@ -365,7 +365,8 @@ int rpc_client_recv_fragment(rdpRpc* rpc, wStream* fragment)
 
 		if (call->OpNum != TsProxySetupReceivePipeOpnum)
 		{
-			Stream_EnsureCapacity(pdu->s, header->response.alloc_hint);
+			if(!Stream_EnsureCapacity(pdu->s, header->response.alloc_hint))
+				return -1;
 			Stream_Write(pdu->s, &buffer[StubOffset], StubLength);
 			rpc->StubFragCount++;
 
@@ -402,7 +403,8 @@ int rpc_client_recv_fragment(rdpRpc* rpc, wStream* fragment)
 			pdu->Flags = 0;
 			pdu->Type = header->common.ptype;
 			pdu->CallId = header->common.call_id;
-			Stream_EnsureCapacity(pdu->s, Stream_Length(fragment));
+			if(!Stream_EnsureCapacity(pdu->s, Stream_Length(fragment)))
+				return -1;
 			Stream_Write(pdu->s, buffer, Stream_Length(fragment));
 			Stream_SealLength(pdu->s);
 			rpc_client_recv_pdu(rpc, pdu);
@@ -423,7 +425,8 @@ int rpc_client_recv_fragment(rdpRpc* rpc, wStream* fragment)
 		pdu->Flags = 0;
 		pdu->Type = header->common.ptype;
 		pdu->CallId = header->common.call_id;
-		Stream_EnsureCapacity(pdu->s, Stream_Length(fragment));
+		if(Stream_EnsureCapacity(pdu->s, Stream_Length(fragment)))
+			return -1;
 		Stream_Write(pdu->s, buffer, Stream_Length(fragment));
 		Stream_SealLength(pdu->s);
 		rpc_client_recv_pdu(rpc, pdu);

libfreerdp/core/orders.c

@@ -843,9 +843,11 @@ int update_approximate_dstblt_order(ORDER_INFO* orderInfo, DSTBLT_ORDER* dstblt)
 
 BOOL update_write_dstblt_order(wStream* s, ORDER_INFO* orderInfo, DSTBLT_ORDER* dstblt)
 {
-	orderInfo->fieldFlags = 0;
 
-	Stream_EnsureRemainingCapacity(s, update_approximate_dstblt_order(orderInfo, dstblt));
+	if(!Stream_EnsureRemainingCapacity(s, update_approximate_dstblt_order(orderInfo, dstblt)))
+		return FALSE;
+
+	orderInfo->fieldFlags = 0;
 
 	orderInfo->fieldFlags |= ORDER_FIELD_01;
 	update_write_coord(s, dstblt->nLeftRect);
@@ -884,9 +886,11 @@ int update_approximate_patblt_order(ORDER_INFO* orderInfo, PATBLT_ORDER* patblt)
 
 BOOL update_write_patblt_order(wStream* s, ORDER_INFO* orderInfo, PATBLT_ORDER* patblt)
 {
-	orderInfo->fieldFlags = 0;
 
-	Stream_EnsureRemainingCapacity(s, update_approximate_patblt_order(orderInfo, patblt));
+	if(!Stream_EnsureRemainingCapacity(s, update_approximate_patblt_order(orderInfo, patblt)))
+		return FALSE;
+
+	orderInfo->fieldFlags = 0;
 
 	orderInfo->fieldFlags |= ORDER_FIELD_01;
 	update_write_coord(s, patblt->nLeftRect);
@@ -939,9 +943,11 @@ int update_approximate_scrblt_order(ORDER_INFO* orderInfo, SCRBLT_ORDER* scrblt)
 
 BOOL update_write_scrblt_order(wStream* s, ORDER_INFO* orderInfo, SCRBLT_ORDER* scrblt)
 {
-	orderInfo->fieldFlags = 0;
 
-	Stream_EnsureRemainingCapacity(s, update_approximate_scrblt_order(orderInfo, scrblt));
+	if(!Stream_EnsureRemainingCapacity(s, update_approximate_scrblt_order(orderInfo, scrblt)));
+		return FALSE;
+
+	orderInfo->fieldFlags = 0;
 
 	orderInfo->fieldFlags |= ORDER_FIELD_01;
 	update_write_coord(s, scrblt->nLeftRect);
@@ -1015,7 +1021,8 @@ BOOL update_write_opaque_rect_order(wStream* s, ORDER_INFO* orderInfo, OPAQUE_RE
 {
 	BYTE byte;
 
-	Stream_EnsureRemainingCapacity(s, update_approximate_opaque_rect_order(orderInfo, opaque_rect));
+	if(!Stream_EnsureRemainingCapacity(s, update_approximate_opaque_rect_order(orderInfo, opaque_rect)))
+		return FALSE;
 
 	orderInfo->fieldFlags = 0;
 
@@ -1274,7 +1281,8 @@ int update_approximate_line_to_order(ORDER_INFO* orderInfo, LINE_TO_ORDER* line_
 
 BOOL update_write_line_to_order(wStream* s, ORDER_INFO* orderInfo, LINE_TO_ORDER* line_to)
 {
-	Stream_EnsureRemainingCapacity(s, update_approximate_line_to_order(orderInfo, line_to));
+	if(!Stream_EnsureRemainingCapacity(s, update_approximate_line_to_order(orderInfo, line_to)))
+		return FALSE;
 
 	orderInfo->fieldFlags = 0;
 
@@ -1381,7 +1389,8 @@ BOOL update_write_memblt_order(wStream* s, ORDER_INFO* orderInfo, MEMBLT_ORDER*
 {
 	UINT16 cacheId;
 
-	Stream_EnsureRemainingCapacity(s, update_approximate_memblt_order(orderInfo, memblt));
+	if(!Stream_EnsureRemainingCapacity(s, update_approximate_memblt_order(orderInfo, memblt)))
+		return FALSE;
 
 	cacheId = (memblt->cacheId & 0xFF) | ((memblt->colorIndex & 0xFF) << 8);
 
@@ -1517,9 +1526,11 @@ int update_approximate_glyph_index_order(ORDER_INFO* orderInfo, GLYPH_INDEX_ORDE
 
 BOOL update_write_glyph_index_order(wStream* s, ORDER_INFO* orderInfo, GLYPH_INDEX_ORDER* glyph_index)
 {
-	orderInfo->fieldFlags = 0;
 
-	Stream_EnsureRemainingCapacity(s, update_approximate_glyph_index_order(orderInfo, glyph_index));
+	if(!Stream_EnsureRemainingCapacity(s, update_approximate_glyph_index_order(orderInfo, glyph_index)))
+		return FALSE;
+
+	orderInfo->fieldFlags = 0;
 
 	orderInfo->fieldFlags |= ORDER_FIELD_01;
 	Stream_Write_UINT8(s, glyph_index->cacheId);
@@ -1911,9 +1922,11 @@ int update_approximate_cache_bitmap_order(CACHE_BITMAP_ORDER* cache_bitmap, BOOL
 
 BOOL update_write_cache_bitmap_order(wStream* s, CACHE_BITMAP_ORDER* cache_bitmap, BOOL compressed, UINT16* flags)
 {
-	*flags = NO_BITMAP_COMPRESSION_HDR;
 
-	Stream_EnsureRemainingCapacity(s, update_approximate_cache_bitmap_order(cache_bitmap, compressed, flags));
+	if(!Stream_EnsureRemainingCapacity(s, update_approximate_cache_bitmap_order(cache_bitmap, compressed, flags)))
+		return FALSE;
+
+	*flags = NO_BITMAP_COMPRESSION_HDR;
 
 	if ((*flags & NO_BITMAP_COMPRESSION_HDR) == 0)
 		cache_bitmap->bitmapLength += 8;
@@ -2028,7 +2041,8 @@ BOOL update_write_cache_bitmap_v2_order(wStream* s, CACHE_BITMAP_V2_ORDER* cache
 {
 	BYTE bitsPerPixelId;
 
-	Stream_EnsureRemainingCapacity(s, update_approximate_cache_bitmap_v2_order(cache_bitmap_v2, compressed, flags));
+	if(!Stream_EnsureRemainingCapacity(s, update_approximate_cache_bitmap_v2_order(cache_bitmap_v2, compressed, flags)))
+		return FALSE;
 
 	bitsPerPixelId = BPP_CBR2[cache_bitmap_v2->bitmapBpp];
 
@@ -2071,12 +2085,14 @@ BOOL update_write_cache_bitmap_v2_order(wStream* s, CACHE_BITMAP_V2_ORDER* cache
 			cache_bitmap_v2->bitmapLength = cache_bitmap_v2->cbCompMainBodySize;
 		}
 
-		Stream_EnsureRemainingCapacity(s, cache_bitmap_v2->bitmapLength);
+		if(!Stream_EnsureRemainingCapacity(s, cache_bitmap_v2->bitmapLength))
+			return FALSE;
 		Stream_Write(s, cache_bitmap_v2->bitmapDataStream, cache_bitmap_v2->bitmapLength);
 	}
 	else
 	{
-		Stream_EnsureRemainingCapacity(s, cache_bitmap_v2->bitmapLength);
+		if(!Stream_EnsureRemainingCapacity(s, cache_bitmap_v2->bitmapLength))
+			return FALSE;
 		Stream_Write(s, cache_bitmap_v2->bitmapDataStream, cache_bitmap_v2->bitmapLength);
 	}
 
@@ -2144,10 +2160,10 @@ BOOL update_write_cache_bitmap_v3_order(wStream* s, CACHE_BITMAP_V3_ORDER* cache
 	BYTE bitsPerPixelId;
 	BITMAP_DATA_EX* bitmapData;
 
+	if(!Stream_EnsureRemainingCapacity(s, update_approximate_cache_bitmap_v3_order(cache_bitmap_v3, flags)))
+		return FALSE;
+
 	bitmapData = &cache_bitmap_v3->bitmapData;
-
-	Stream_EnsureRemainingCapacity(s, update_approximate_cache_bitmap_v3_order(cache_bitmap_v3, flags));
-
 	bitsPerPixelId = BPP_CBR23[cache_bitmap_v3->bpp];
 
 	*flags = (cache_bitmap_v3->cacheId & 0x00000003) |
@@ -2214,7 +2230,8 @@ BOOL update_write_cache_color_table_order(wStream* s, CACHE_COLOR_TABLE_ORDER* c
 	if (cache_color_table->numberColors != 256)
 		return FALSE;
 
-	Stream_EnsureRemainingCapacity(s, update_approximate_cache_color_table_order(cache_color_table, flags));
+	if(!Stream_EnsureRemainingCapacity(s, update_approximate_cache_color_table_order(cache_color_table, flags)))
+		return FALSE;
 
 	Stream_Write_UINT8(s, cache_color_table->cacheIndex); /* cacheIndex (1 byte) */
 	Stream_Write_UINT16(s, cache_color_table->numberColors); /* numberColors (2 bytes) */
@@ -2285,7 +2302,8 @@ BOOL update_write_cache_glyph_order(wStream* s, CACHE_GLYPH_ORDER* cache_glyph,
 	INT16 lsi16;
 	GLYPH_DATA* glyph;
 
-	Stream_EnsureRemainingCapacity(s, update_approximate_cache_glyph_order(cache_glyph, flags));
+	if(!Stream_EnsureRemainingCapacity(s, update_approximate_cache_glyph_order(cache_glyph, flags)))
+		return FALSE;
 
 	Stream_Write_UINT8(s, cache_glyph->cacheId); /* cacheId (1 byte) */
 	Stream_Write_UINT8(s, cache_glyph->cGlyphs); /* cGlyphs (1 byte) */
@@ -2373,7 +2391,8 @@ BOOL update_write_cache_glyph_v2_order(wStream* s, CACHE_GLYPH_V2_ORDER* cache_g
 	int i;
 	GLYPH_DATA_V2* glyph;
 
-	Stream_EnsureRemainingCapacity(s, update_approximate_cache_glyph_v2_order(cache_glyph_v2, flags));
+	if(!Stream_EnsureRemainingCapacity(s, update_approximate_cache_glyph_v2_order(cache_glyph_v2, flags)))
+		return FALSE;
 
 	*flags = (cache_glyph_v2->cacheId & 0x000F) |
 			((cache_glyph_v2->flags & 0x000F) << 4) |
@@ -2531,7 +2550,8 @@ BOOL update_write_cache_brush_order(wStream* s, CACHE_BRUSH_ORDER* cache_brush,
 	BYTE iBitmapFormat;
 	BOOL compressed = FALSE;
 
-	Stream_EnsureRemainingCapacity(s, update_approximate_cache_brush_order(cache_brush, flags));
+	if(!Stream_EnsureRemainingCapacity(s, update_approximate_cache_brush_order(cache_brush, flags)))
+		return FALSE;
 
 	iBitmapFormat = BPP_BMF[cache_brush->bpp];
 
@@ -2660,9 +2680,10 @@ BOOL update_write_create_offscreen_bitmap_order(wStream* s, CREATE_OFFSCREEN_BIT
 	BOOL deleteListPresent;
 	OFFSCREEN_DELETE_LIST* deleteList;
 
-	deleteList = &(create_offscreen_bitmap->deleteList);
+	if(!Stream_EnsureRemainingCapacity(s, update_approximate_create_offscreen_bitmap_order(create_offscreen_bitmap)))
+		return FALSE;
 
-	Stream_EnsureRemainingCapacity(s, update_approximate_create_offscreen_bitmap_order(create_offscreen_bitmap));
+	deleteList = &(create_offscreen_bitmap->deleteList);
 
 	flags = create_offscreen_bitmap->id & 0x7FFF;
 
@@ -2708,7 +2729,8 @@ int update_approximate_switch_surface_order(SWITCH_SURFACE_ORDER* switch_surface
 
 BOOL update_write_switch_surface_order(wStream* s, SWITCH_SURFACE_ORDER* switch_surface)
 {
-	Stream_EnsureRemainingCapacity(s, update_approximate_switch_surface_order(switch_surface));
+	if(!Stream_EnsureRemainingCapacity(s, update_approximate_switch_surface_order(switch_surface)))
+		return FALSE;
 
 	Stream_Write_UINT16(s, switch_surface->bitmapId); /* bitmapId (2 bytes) */
 

libfreerdp/core/peer.c

@@ -145,6 +145,8 @@ static int freerdp_peer_virtual_channel_write(freerdp_peer* client, HANDLE hChan
 	while (length > 0)
 	{
 		s = rdp_send_stream_init(rdp);
+		if (!s)
+			return -1;
 
 		if (length > maxChunkSize)
 		{
@@ -161,10 +163,18 @@ static int freerdp_peer_virtual_channel_write(freerdp_peer* client, HANDLE hChan
 
 		Stream_Write_UINT32(s, totalLength);
 		Stream_Write_UINT32(s, flags);
-		Stream_EnsureRemainingCapacity(s, chunkSize);
+		if(!Stream_EnsureRemainingCapacity(s, chunkSize))
+		{
+			Stream_Release(s);
+			return -1;
+		}
 		Stream_Write(s, buffer, chunkSize);
 
-		rdp_send(rdp, s, peerChannel->channelId);
+		if(!rdp_send(rdp, s, peerChannel->channelId))
+		{
+			Stream_Release(s);
+			return -1;
+		}
 
 		buffer += chunkSize;
 		length -= chunkSize;

libfreerdp/core/rdp.c

@@ -271,6 +271,8 @@ wStream* rdp_message_channel_pdu_init(rdpRdp* rdp)
 {
 	wStream* s;
 	s = transport_send_stream_init(rdp->transport, 2048);
+	if (!s)
+		return NULL;
 	Stream_Seek(s, RDP_PACKET_HEADER_MAX_LENGTH);
 	rdp_security_stream_init(rdp, s, TRUE);
 	return s;
@@ -710,7 +712,8 @@ BOOL rdp_write_monitor_layout_pdu(wStream* s, UINT32 monitorCount, MONITOR_DEF*
 {
 	UINT32 index;
 	MONITOR_DEF* monitor;
-	Stream_EnsureRemainingCapacity(s, 4 + (monitorCount * 20));
+	if(!Stream_EnsureRemainingCapacity(s, 4 + (monitorCount * 20)))
+		return FALSE;
 	Stream_Write_UINT32(s, monitorCount); /* monitorCount (4 bytes) */
 
 	for (index = 0; index < monitorCount; index++)
@@ -755,7 +758,10 @@ int rdp_recv_data_pdu(rdpRdp* rdp, wStream* s)
 		if (bulk_decompress(rdp->bulk, Stream_Pointer(s), SrcSize, &pDstData, &DstSize, compressedType))
 		{
 			if (!(cs = StreamPool_Take(rdp->transport->ReceivePool, DstSize)))
+			{
+				WLog_ERR(TAG, "Coudn't take stream from pool");
 				return -1;
+			}
 
 			Stream_SetPosition(cs, 0);
 			Stream_Write(cs, pDstData, DstSize);

libfreerdp/core/transport.c

@@ -63,7 +63,11 @@ wStream* transport_send_stream_init(rdpTransport* transport, int size)
 	wStream* s;
 	if (!(s = StreamPool_Take(transport->ReceivePool, size)))
 		return NULL;
-	Stream_EnsureCapacity(s, size);
+	if(!Stream_EnsureCapacity(s, size))
+	{
+		Stream_Release(s);
+		return NULL;
+	}
 	Stream_SetPosition(s, 0);
 	return s;
 }
@@ -446,7 +450,8 @@ int transport_read_pdu(rdpTransport* transport, wStream* s)
 
 	position = Stream_GetPosition(s);
 	/* Make sure there is enough space for the longest header within the stream */
-	Stream_EnsureCapacity(s, 4);
+	if(!Stream_EnsureCapacity(s, 4))
+		return -1;
 
 	/* Make sure at least two bytes are read for further processing */
 	if (position < 2 && (status = transport_read_layer_bytes(transport, s, 2 - position)) != 1)
@@ -552,7 +557,8 @@ int transport_read_pdu(rdpTransport* transport, wStream* s)
 		}
 	}
 
-	Stream_EnsureCapacity(s, Stream_GetPosition(s) + pduLength);
+	if(!Stream_EnsureCapacity(s, Stream_GetPosition(s) + pduLength))
+		return -1;
 	status = transport_read_layer_bytes(transport, s, pduLength - Stream_GetPosition(s));
 
 	if (status != 1)
@@ -736,9 +742,7 @@ int transport_check_fds(rdpTransport* transport)
 
 		received = transport->ReceiveBuffer;
 		if (!(transport->ReceiveBuffer = StreamPool_Take(transport->ReceivePool, 0)))
-		{
 			return -1;
-		}
 		/**
 		 * status:
 		 * 	-1: error

libfreerdp/core/update.c

@@ -114,7 +114,8 @@ BOOL update_read_bitmap_data(rdpUpdate* update, wStream* s, BITMAP_DATA* bitmapD
 
 BOOL update_write_bitmap_data(rdpUpdate* update, wStream* s, BITMAP_DATA* bitmapData)
 {
-	Stream_EnsureRemainingCapacity(s, 64 + bitmapData->bitmapLength);
+	if(!Stream_EnsureRemainingCapacity(s, 64 + bitmapData->bitmapLength))
+		return TRUE;
 
 	bitmapData->flags = 0;
 	bitmapData->cbCompFirstRowSize = 0;
@@ -202,7 +203,8 @@ BOOL update_write_bitmap_update(rdpUpdate* update, wStream* s, BITMAP_UPDATE* bi
 {
 	int i;
 
-	Stream_EnsureRemainingCapacity(s, 32);
+	if(!Stream_EnsureRemainingCapacity(s, 32))
+		return TRUE;
 
 	Stream_Write_UINT16(s, UPDATE_TYPE_BITMAP); /* updateType */
 

winpr/include/winpr/stream.h

@@ -44,8 +44,8 @@ struct _wStream
 };
 typedef struct _wStream wStream;
 
-WINPR_API void Stream_EnsureCapacity(wStream* s, size_t size);
+WINPR_API BOOL Stream_EnsureCapacity(wStream* s, size_t size);
-WINPR_API void Stream_EnsureRemainingCapacity(wStream* s, size_t size);
+WINPR_API BOOL Stream_EnsureRemainingCapacity(wStream* s, size_t size);
 
 WINPR_API wStream* Stream_New(BYTE* buffer, size_t size);
 WINPR_API void Stream_Free(wStream* s, BOOL bFreeBuffer);

winpr/libwinpr/utils/collections/StreamPool.c

@@ -169,13 +169,13 @@ wStream* StreamPool_Take(wStreamPool* pool, size_t size)
 	if (foundIndex < 0)
 	{
 		s = Stream_New(NULL, size);
+		if (!s)
+			goto out_fail;
 	}
 	else
 	{
-		StreamPool_ShiftAvailable(pool, foundIndex, -1);
-
 		Stream_SetPosition(s, 0);
-		Stream_EnsureCapacity(s, size);
+		StreamPool_ShiftAvailable(pool, foundIndex, -1);
 	}
 
 	if (s)
@@ -185,6 +185,7 @@ wStream* StreamPool_Take(wStreamPool* pool, size_t size)
 		StreamPool_AddUsed(pool, s);
 	}
 
+out_fail:
 	if (pool->synchronized)
 		LeaveCriticalSection(&pool->lock);
 

winpr/libwinpr/utils/stream.c

@@ -25,7 +25,7 @@
 #include <winpr/crt.h>
 #include <winpr/stream.h>
 
-void Stream_EnsureCapacity(wStream* s, size_t size)
+BOOL Stream_EnsureCapacity(wStream* s, size_t size)
 {
 	if (s->capacity < size)
 	{
@@ -48,7 +48,7 @@ void Stream_EnsureCapacity(wStream* s, size_t size)
 
 		new_buf = (BYTE*) realloc(s->buffer, new_capacity);
 		if (!new_buf)
-			return;
+			return FALSE;
 		s->buffer = new_buf;
 		s->capacity = new_capacity;
 		s->length = new_capacity;
@@ -56,12 +56,14 @@ void Stream_EnsureCapacity(wStream* s, size_t size)
 
 		Stream_SetPosition(s, position);
 	}
+	return TRUE;
 }
 
-void Stream_EnsureRemainingCapacity(wStream* s, size_t size)
+BOOL Stream_EnsureRemainingCapacity(wStream* s, size_t size)
 {
 	if (Stream_GetPosition(s) + size > Stream_Capacity(s))
-		Stream_EnsureCapacity(s, Stream_Capacity(s) + size);
+		return Stream_EnsureCapacity(s, Stream_Capacity(s) + size);
+	return TRUE;
 }
 
 wStream* Stream_New(BYTE* buffer, size_t size)

winpr/libwinpr/utils/test/TestStream.c

@@ -148,9 +148,15 @@ static BOOL TestStream_Extent(UINT32 maxSize)
 	for (i = 1; i < maxSize; i++)
 	{
 		if (i % 2)
-			Stream_EnsureRemainingCapacity(s, i);
+		{
+			if(!Stream_EnsureRemainingCapacity(s, i))
+				goto fail;
+		}
 		else
-			Stream_EnsureCapacity(s, i);
+		{
+			if(!Stream_EnsureCapacity(s, i))
+				goto fail;
+		}
 
 		Stream_SetPosition(s, i);
 		Stream_SealLength(s);