From ec030e82552052bbea8462d9052919466aa896b2 Mon Sep 17 00:00:00 2001
From: akallabeth <akallabeth@posteo.net>
Date: Thu, 27 Feb 2025 17:37:15 +0100
Subject: [PATCH] [winpr,stream] implement reference counting for streams

---
 winpr/include/winpr/stream.h                  | 13 ++++++++++
 winpr/libwinpr/utils/collections/StreamPool.c | 24 +++++++++----------
 winpr/libwinpr/utils/stream.c                 |  6 ++---
 3 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/winpr/include/winpr/stream.h b/winpr/include/winpr/stream.h
index 4ae7c5b78..8ef97b2ee 100644
--- a/winpr/include/winpr/stream.h
+++ b/winpr/include/winpr/stream.h
@@ -1368,7 +1368,20 @@ extern "C"
 
 	WINPR_API void StreamPool_Return(wStreamPool* pool, wStream* s);
 
+	/** @brief increment reference count of stream
+	 *
+	 *  @param s The stream to reference
+	 *  @bug versions < 3.13.0 did only handle streams returned by StreamPool_Take
+	 */
 	WINPR_API void Stream_AddRef(wStream* s);
+
+	/** @brief Release a reference to a stream.
+	 *  If the reference count reaches \b 0 it is returned to the StreamPool it was taken from or \b
+	 * Stream_Free is called.
+	 *
+	 *  @param s The stream to release
+	 *  @bug versions < 3.13.0 did only handle streams returned by StreamPool_Take
+	 */
 	WINPR_API void Stream_Release(wStream* s);
 
 	WINPR_ATTR_MALLOC(Stream_Release, 1)
diff --git a/winpr/libwinpr/utils/collections/StreamPool.c b/winpr/libwinpr/utils/collections/StreamPool.c
index b550c9bad..8870f4d2f 100644
--- a/winpr/libwinpr/utils/collections/StreamPool.c
+++ b/winpr/libwinpr/utils/collections/StreamPool.c
@@ -285,10 +285,7 @@ static void StreamPool_Remove(wStreamPool* pool, wStream* s)
 static void StreamPool_ReleaseOrReturn(wStreamPool* pool, wStream* s)
 {
 	StreamPool_Lock(pool);
-	if (s->count > 0)
-		s->count--;
-	if (s->count == 0)
-		StreamPool_Remove(pool, s);
+	StreamPool_Remove(pool, s);
 	StreamPool_Unlock(pool);
 }
 
@@ -310,12 +307,7 @@ void StreamPool_Return(wStreamPool* pool, wStream* s)
 void Stream_AddRef(wStream* s)
 {
 	WINPR_ASSERT(s);
-	if (s->pool)
-	{
-		StreamPool_Lock(s->pool);
-		s->count++;
-		StreamPool_Unlock(s->pool);
-	}
+	s->count++;
 }
 
 /**
@@ -325,8 +317,16 @@ void Stream_AddRef(wStream* s)
 void Stream_Release(wStream* s)
 {
 	WINPR_ASSERT(s);
-	if (s->pool)
-		StreamPool_ReleaseOrReturn(s->pool, s);
+
+	if (s->count > 0)
+		s->count--;
+	if (s->count == 0)
+	{
+		if (s->pool)
+			StreamPool_ReleaseOrReturn(s->pool, s);
+		else
+			Stream_Free(s, TRUE);
+	}
 }
 
 /**
diff --git a/winpr/libwinpr/utils/stream.c b/winpr/libwinpr/utils/stream.c
index 145b54f5e..d1ed7b718 100644
--- a/winpr/libwinpr/utils/stream.c
+++ b/winpr/libwinpr/utils/stream.c
@@ -98,7 +98,7 @@ wStream* Stream_New(BYTE* buffer, size_t size)
 	if (!buffer && !size)
 		return NULL;
 
-	s = malloc(sizeof(wStream));
+	s = calloc(1, sizeof(wStream));
 	if (!s)
 		return NULL;
 
@@ -118,7 +118,7 @@ wStream* Stream_New(BYTE* buffer, size_t size)
 	s->length = size;
 
 	s->pool = NULL;
-	s->count = 0;
+	s->count = 1;
 	s->isAllocatedStream = TRUE;
 	s->isOwner = TRUE;
 	return s;
@@ -147,7 +147,7 @@ wStream* Stream_StaticInit(wStream* s, BYTE* buffer, size_t size)
 	s->buffer = s->pointer = buffer;
 	s->capacity = s->length = size;
 	s->pool = NULL;
-	s->count = 0;
+	s->count = 1;
 	s->isAllocatedStream = FALSE;
 	s->isOwner = FALSE;
 	return s;