From 602f4a2e14b41703b5f431de3154cd46a5750a2d Mon Sep 17 00:00:00 2001
From: Armin Novak <armin.novak@thincast.com>
Date: Mon, 22 Oct 2018 16:20:34 +0200
Subject: [PATCH 2/6] Fixed CVE-2018-8785

Thanks to Eyal Itkin from Check Point Software Technologies.
---
 include/freerdp/codec/zgfx.h | 30 +++++--------------
 libfreerdp/codec/zgfx.c      | 70 ++++++++++++++++++++++++++++++++++----------
 2 files changed, 61 insertions(+), 39 deletions(-)

Index: b/include/freerdp/codec/zgfx.h
===================================================================
--- a/include/freerdp/codec/zgfx.h	2019-01-08 19:49:48.770871914 +0800
+++ b/include/freerdp/codec/zgfx.h	2019-01-08 19:53:46.228463491 +0800
@@ -28,25 +28,6 @@
 #define ZGFX_SEGMENTED_SINGLE			0xE0
 #define ZGFX_SEGMENTED_MULTIPART		0xE1
 
-struct _ZGFX_CONTEXT
-{
-	BOOL Compressor;
-
-	BYTE* pbInputCurrent;
-	BYTE* pbInputEnd;
-
-	UINT32 bits;
-	UINT32 cBitsRemaining;
-	UINT32 BitsCurrent;
-	UINT32 cBitsCurrent;
-
-	BYTE OutputBuffer[65536];
-	UINT32 OutputCount;
-
-	BYTE HistoryBuffer[2500000];
-	UINT32 HistoryIndex;
-	UINT32 HistoryBufferSize;
-};
 typedef struct _ZGFX_CONTEXT ZGFX_CONTEXT;
 
 #ifdef __cplusplus
@@ -66,4 +47,4 @@ FREERDP_API void zgfx_context_free(ZGFX_
 #endif
 
 #endif /* FREERDP_CODEC_ZGFX_H */
- 
+
Index: b/libfreerdp/codec/zgfx.c
===================================================================
--- a/libfreerdp/codec/zgfx.c	2019-01-08 19:49:48.770871914 +0800
+++ b/libfreerdp/codec/zgfx.c	2019-01-08 19:50:22.975101170 +0800
@@ -39,14 +39,34 @@
 
 struct _ZGFX_TOKEN
 {
-	int prefixLength;
-	int prefixCode;
-	int valueBits;
-	int tokenType;
+	UINT32 prefixLength;
+	UINT32 prefixCode;
+	UINT32 valueBits;
+	UINT32 tokenType;
 	UINT32 valueBase;
 };
 typedef struct _ZGFX_TOKEN ZGFX_TOKEN;
 
+struct _ZGFX_CONTEXT
+{
+	BOOL Compressor;
+
+	const BYTE* pbInputCurrent;
+	const BYTE* pbInputEnd;
+
+	UINT32 bits;
+	UINT32 cBitsRemaining;
+	UINT32 BitsCurrent;
+	UINT32 cBitsCurrent;
+
+	BYTE OutputBuffer[65536];
+	UINT32 OutputCount;
+
+	BYTE HistoryBuffer[2500000];
+	UINT32 HistoryIndex;
+	UINT32 HistoryBufferSize;
+};
+
 static const ZGFX_TOKEN ZGFX_TOKEN_TABLE[] =
 {
 	// len code vbits type  vbase
@@ -93,17 +113,26 @@ static const ZGFX_TOKEN ZGFX_TOKEN_TABLE
 	{ 0 }
 };
 
-#define zgfx_GetBits(_zgfx, _nbits) \
-	while (_zgfx->cBitsCurrent < _nbits) { \
-		_zgfx->BitsCurrent <<= 8; \
-		if (_zgfx->pbInputCurrent < _zgfx->pbInputEnd) \
-			_zgfx->BitsCurrent += *(_zgfx->pbInputCurrent)++; \
-		_zgfx->cBitsCurrent += 8; \
-	} \
-	_zgfx->cBitsRemaining -= _nbits; \
-	_zgfx->cBitsCurrent -= _nbits; \
-	_zgfx->bits = _zgfx->BitsCurrent >> _zgfx->cBitsCurrent; \
+static INLINE BOOL zgfx_GetBits(ZGFX_CONTEXT* _zgfx, UINT32 _nbits)
+{
+	if (!_zgfx)
+		return FALSE;
+
+	while (_zgfx->cBitsCurrent < _nbits)
+	{
+		_zgfx->BitsCurrent <<= 8;
+
+		if (_zgfx->pbInputCurrent < _zgfx->pbInputEnd)
+			_zgfx->BitsCurrent += *(_zgfx->pbInputCurrent)++;
+
+		_zgfx->cBitsCurrent += 8;
+	}
+
+	_zgfx->cBitsRemaining -= _nbits;
+	_zgfx->cBitsCurrent -= _nbits;
+	_zgfx->bits = _zgfx->BitsCurrent >> _zgfx->cBitsCurrent;
 	_zgfx->BitsCurrent &= ((1 << _zgfx->cBitsCurrent) - 1);
+}
 
 void zgfx_history_buffer_ring_write(ZGFX_CONTEXT* zgfx, BYTE* src, UINT32 count)
 {
@@ -192,7 +221,7 @@ int zgfx_decompress_segment(ZGFX_CONTEXT
 {
 	BYTE c;
 	BYTE flags;
-	int extra;
+	UINT32 extra = 0;
 	int opIndex;
 	int haveBits;
 	int inPrefix;
@@ -351,6 +380,7 @@ int zgfx_decompress(ZGFX_CONTEXT* zgfx,
 		UINT32 segmentOffset;
 		UINT32 uncompressedSize;
 		BYTE* pConcatenated;
+		size_t used = 0;
 
 		segmentOffset = 7;
 		segmentCount = *((UINT16*) &pSrcData[1]); /* segmentCount (2 bytes) */
@@ -371,8 +401,15 @@ int zgfx_decompress(ZGFX_CONTEXT* zgfx,
 			status = zgfx_decompress_segment(zgfx, &pSrcData[segmentOffset], segmentSize);
 			segmentOffset += segmentSize;
 
+			if (zgfx->OutputCount > UINT32_MAX - used)
+				return -1;
+
+			if (used + zgfx->OutputCount > uncompressedSize)
+				return -1;
+
 			CopyMemory(pConcatenated, zgfx->OutputBuffer, zgfx->OutputCount);
 			pConcatenated += zgfx->OutputCount;
+			used += zgfx->OutputCount;
 		}
 	}
 	else
