--- FreeRDP-2.0.0-rc3/libfreerdp/codec/zgfx.c	2019-06-05 01:21:31.264000000 -0400
+++ FreeRDP-2.0.0-rc3/libfreerdp/codec/zgfx_2.c	2019-06-05 01:28:52.362000000 -0400
@@ -44,14 +44,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
@@ -98,18 +118,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);
-	return TRUE;
+}
 
 static void zgfx_history_buffer_ring_write(ZGFX_CONTEXT* zgfx, const BYTE* src, size_t count)
 {
@@ -194,7 +222,7 @@ static BOOL zgfx_decompress_segment(ZGFX
 {
 	BYTE c;
 	BYTE flags;
-	int extra;
+	UINT32 extra = 0;
 	int opIndex;
 	int haveBits;
 	int inPrefix;
@@ -338,8 +366,8 @@ int zgfx_decompress(ZGFX_CONTEXT* zgfx,
 {
 	int status = -1;
 	BYTE descriptor;
-
 	wStream* stream = Stream_New((BYTE*)pSrcData, SrcSize);
+
 	if (!stream)
 		return -1;
 
@@ -354,6 +382,7 @@ int zgfx_decompress(ZGFX_CONTEXT* zgfx,
 			goto fail;
 
 		*ppDstData = NULL;
+
 		if (zgfx->OutputCount > 0)
 			*ppDstData = (BYTE*) malloc(zgfx->OutputCount);
 
@@ -370,6 +399,7 @@ int zgfx_decompress(ZGFX_CONTEXT* zgfx,
 		UINT16 segmentCount;
 		UINT32 uncompressedSize;
 		BYTE* pConcatenated;
+		size_t used = 0;
 
 		if (Stream_GetRemainingLength(stream) < 6)
 			goto fail;
@@ -398,8 +428,15 @@ int zgfx_decompress(ZGFX_CONTEXT* zgfx,
 			if (!zgfx_decompress_segment(zgfx, stream, segmentSize))
 				goto fail;
 
+			if (zgfx->OutputCount > UINT32_MAX - used)
+				goto fail;
+
+			if (used + zgfx->OutputCount > uncompressedSize)
+				goto fail;
+
 			CopyMemory(pConcatenated, zgfx->OutputBuffer, zgfx->OutputCount);
 			pConcatenated += zgfx->OutputCount;
+			used += zgfx->OutputCount;
 		}
 	}
 	else
diff --git a/include/freerdp/codec/zgfx.h b/include/freerdp/codec/zgfx.h
index b5c65319cb..b0faf6e0de 100644
--- a/include/freerdp/codec/zgfx.h
+++ b/include/freerdp/codec/zgfx.h
@@ -32,34 +32,18 @@
 
 #define ZGFX_SEGMENTED_MAXSIZE			65535
 
-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;
-};
 typedef struct _ZGFX_CONTEXT ZGFX_CONTEXT;
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
-FREERDP_API int zgfx_decompress(ZGFX_CONTEXT* zgfx, const BYTE* pSrcData, UINT32 SrcSize, BYTE** ppDstData, UINT32* pDstSize, UINT32 flags);
-FREERDP_API int zgfx_compress(ZGFX_CONTEXT* zgfx, const BYTE* pSrcData, UINT32 SrcSize, BYTE** ppDstData, UINT32* pDstSize, UINT32* pFlags);
-FREERDP_API int zgfx_compress_to_stream(ZGFX_CONTEXT* zgfx, wStream* sDst, const BYTE* pUncompressed, UINT32 uncompressedSize, UINT32* pFlags);
+FREERDP_API int zgfx_decompress(ZGFX_CONTEXT* zgfx, const BYTE* pSrcData, UINT32 SrcSize,
+                                BYTE** ppDstData, UINT32* pDstSize, UINT32 flags);
+FREERDP_API int zgfx_compress(ZGFX_CONTEXT* zgfx, const BYTE* pSrcData, UINT32 SrcSize,
+                              BYTE** ppDstData, UINT32* pDstSize, UINT32* pFlags);
+FREERDP_API int zgfx_compress_to_stream(ZGFX_CONTEXT* zgfx, wStream* sDst,
+                                        const BYTE* pUncompressed, UINT32 uncompressedSize, UINT32* pFlags);
 
 FREERDP_API void zgfx_context_reset(ZGFX_CONTEXT* zgfx, BOOL flush);
 
@@ -71,4 +55,4 @@ FREERDP_API void zgfx_context_free(ZGFX_CONTEXT* zgfx);
 #endif
 
 #endif /* FREERDP_CODEC_ZGFX_H */
- 
+