fix tjanczuk#225: memory leak in websocket response processing

Azure · Oct 23, 2012 · 2c069e3 · 2c069e3
1 parent c1d021d
commit 2c069e3
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 6 deletions.
diff --git a/src/iisnode/cnodehttpstoredcontext.cpp b/src/iisnode/cnodehttpstoredcontext.cpp
@@ -5,11 +5,15 @@ CNodeHttpStoredContext::CNodeHttpStoredContext(CNodeApplication* nodeApplication
 	chunkLength(0), chunkTransmitted(0), isChunked(FALSE), pipe(INVALID_HANDLE_VALUE), result(S_OK), isLastChunk(FALSE),
 	requestNotificationStatus(RQ_NOTIFICATION_PENDING), connectionRetryCount(0), pendingAsyncOperationCount(1),
 	targetUrl(NULL), targetUrlLength(0), childContext(NULL), isConnectionFromPool(FALSE), expectResponseBody(TRUE),
-	closeConnection(FALSE), isUpgrade(FALSE), upgradeContext(NULL), opaqueFlagSet(FALSE), requestPumpStarted(FALSE)
+	closeConnection(FALSE), isUpgrade(FALSE), upgradeContext(NULL), opaqueFlagSet(FALSE), requestPumpStarted(FALSE),
+	responseChunkBufferSize(0)
 {
 	IHttpTraceContext* tctx;
 	LPCGUID pguid;
 
+	this->responseChunk.DataChunkType = HttpDataChunkFromMemory;
+	this->responseChunk.FromMemory.pBuffer = NULL;
+
 	RtlZeroMemory(&this->asyncContext, sizeof(ASYNC_CONTEXT));
 	if (NULL != (tctx = context->GetTraceContext()) && NULL != (pguid = tctx->GetTraceActivityId()))
 	{
@@ -35,6 +39,36 @@ CNodeHttpStoredContext::~CNodeHttpStoredContext()
 		CloseHandle(this->pipe);
 		this->pipe = INVALID_HANDLE_VALUE;
 	}
+
+	if (this->responseChunk.FromMemory.pBuffer) {
+		free(this->responseChunk.FromMemory.pBuffer);
+		this->responseChunk.FromMemory.pBuffer = NULL;
+		responseChunkBufferSize = 0;
+	}
+}
+
+HRESULT CNodeHttpStoredContext::EnsureResponseChunk(DWORD size, HTTP_DATA_CHUNK** chunk)
+{
+	HRESULT hr;
+
+	if (size > this->responseChunkBufferSize)
+	{
+		if (this->responseChunk.FromMemory.pBuffer) 
+		{
+			free(this->responseChunk.FromMemory.pBuffer);
+			this->responseChunk.FromMemory.pBuffer = NULL;
+			this->responseChunkBufferSize = 0;
+		}
+
+		ErrorIf(NULL == (this->responseChunk.FromMemory.pBuffer = malloc(size)), ERROR_NOT_ENOUGH_MEMORY);
+		this->responseChunkBufferSize = size;
+	}
+
+	*chunk = &this->responseChunk;
+
+	return S_OK;
+Error:
+	return hr;
 }
 
 IHttpContext* CNodeHttpStoredContext::GetHttpContext() 

diff --git a/src/iisnode/cnodehttpstoredcontext.h b/src/iisnode/cnodehttpstoredcontext.h
@@ -37,6 +37,8 @@ class CNodeHttpStoredContext : public IHttpStoredContext
 	BOOL opaqueFlagSet;
 	BOOL requestPumpStarted;
 	FILETIME startTime;
+	HTTP_DATA_CHUNK responseChunk;
+	DWORD responseChunkBufferSize;
 
 public:
 
@@ -104,6 +106,7 @@ class CNodeHttpStoredContext : public IHttpStoredContext
 	void SetRequestPumpStarted();
 	BOOL GetRequestPumpStarted();
 	FILETIME* GetStartTime();
+	HRESULT EnsureResponseChunk(DWORD size, HTTP_DATA_CHUNK** chunk);
 
 	static CNodeHttpStoredContext* Get(LPOVERLAPPED overlapped);
 

diff --git a/src/iisnode/cprotocolbridge.cpp b/src/iisnode/cprotocolbridge.cpp
@@ -1537,12 +1537,8 @@ void WINAPI CProtocolBridge::ProcessResponseBody(DWORD error, DWORD bytesTransfe
 			DWORD remainingChunkSize = ctx->GetChunkLength() - ctx->GetChunkTransmitted();
 			DWORD bytesToSend = dataInBuffer < remainingChunkSize ? dataInBuffer : remainingChunkSize;
 
-			// CR: consider using malloc here (memory can be released after Flush)
-
-			ErrorIf(NULL == (chunk = (HTTP_DATA_CHUNK*) ctx->GetHttpContext()->AllocateRequestMemory(sizeof HTTP_DATA_CHUNK)), ERROR_NOT_ENOUGH_MEMORY);
-			chunk->DataChunkType = HttpDataChunkFromMemory;
+			CheckError(ctx->EnsureResponseChunk(bytesToSend, &chunk));
 			chunk->FromMemory.BufferLength = bytesToSend;
-			ErrorIf(NULL == (chunk->FromMemory.pBuffer = ctx->GetHttpContext()->AllocateRequestMemory(chunk->FromMemory.BufferLength)), ERROR_NOT_ENOUGH_MEMORY);
 			memcpy(chunk->FromMemory.pBuffer, (char*)ctx->GetBuffer() + ctx->GetParsingOffset(), chunk->FromMemory.BufferLength);
 
 			if (bytesToSend == dataInBuffer)