zgfx.c

 
#include <freerdp/config.h>
 
#include <winpr/assert.h>
#include <winpr/cast.h>
#include <winpr/crt.h>
#include <winpr/print.h>
#include <winpr/bitstream.h>
 
#include <freerdp/log.h>
#include <freerdp/codec/zgfx.h>
 
#define TAG FREERDP_TAG("codec")
 
typedef struct
{
  UINT32 prefixLength;
  UINT32 prefixCode;
  UINT32 valueBits;
  UINT32 tokenType;
  UINT32 valueBase;
} ZGFX_TOKEN;
 
struct S_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
  { 1, 0, 8, 0, 0 },           // 0
  { 5, 17, 5, 1, 0 },          // 10001
  { 5, 18, 7, 1, 32 },         // 10010
  { 5, 19, 9, 1, 160 },        // 10011
  { 5, 20, 10, 1, 672 },       // 10100
  { 5, 21, 12, 1, 1696 },      // 10101
  { 5, 24, 0, 0, 0x00 },       // 11000
  { 5, 25, 0, 0, 0x01 },       // 11001
  { 6, 44, 14, 1, 5792 },      // 101100
  { 6, 45, 15, 1, 22176 },     // 101101
  { 6, 52, 0, 0, 0x02 },       // 110100
  { 6, 53, 0, 0, 0x03 },       // 110101
  { 6, 54, 0, 0, 0xFF },       // 110110
  { 7, 92, 18, 1, 54944 },     // 1011100
  { 7, 93, 20, 1, 317088 },    // 1011101
  { 7, 110, 0, 0, 0x04 },      // 1101110
  { 7, 111, 0, 0, 0x05 },      // 1101111
  { 7, 112, 0, 0, 0x06 },      // 1110000
  { 7, 113, 0, 0, 0x07 },      // 1110001
  { 7, 114, 0, 0, 0x08 },      // 1110010
  { 7, 115, 0, 0, 0x09 },      // 1110011
  { 7, 116, 0, 0, 0x0A },      // 1110100
  { 7, 117, 0, 0, 0x0B },      // 1110101
  { 7, 118, 0, 0, 0x3A },      // 1110110
  { 7, 119, 0, 0, 0x3B },      // 1110111
  { 7, 120, 0, 0, 0x3C },      // 1111000
  { 7, 121, 0, 0, 0x3D },      // 1111001
  { 7, 122, 0, 0, 0x3E },      // 1111010
  { 7, 123, 0, 0, 0x3F },      // 1111011
  { 7, 124, 0, 0, 0x40 },      // 1111100
  { 7, 125, 0, 0, 0x80 },      // 1111101
  { 8, 188, 20, 1, 1365664 },  // 10111100
  { 8, 189, 21, 1, 2414240 },  // 10111101
  { 8, 252, 0, 0, 0x0C },      // 11111100
  { 8, 253, 0, 0, 0x38 },      // 11111101
  { 8, 254, 0, 0, 0x39 },      // 11111110
  { 8, 255, 0, 0, 0x66 },      // 11111111
  { 9, 380, 22, 1, 4511392 },  // 101111100
  { 9, 381, 23, 1, 8705696 },  // 101111101
  { 9, 382, 24, 1, 17094304 }, // 101111110
  { 0 }
};
 
static INLINE BOOL zgfx_GetBits(ZGFX_CONTEXT* WINPR_RESTRICT 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 INLINE void zgfx_history_buffer_ring_write(ZGFX_CONTEXT* WINPR_RESTRICT zgfx,
                                                  const BYTE* WINPR_RESTRICT src, size_t count)
{
  if (count <= 0)
    return;
 
  if (count > zgfx->HistoryBufferSize)
  {
    const size_t residue = count - zgfx->HistoryBufferSize;
    count = zgfx->HistoryBufferSize;
    src += residue;
    zgfx->HistoryIndex = (zgfx->HistoryIndex + residue) % zgfx->HistoryBufferSize;
  }
 
  if (zgfx->HistoryIndex + count <= zgfx->HistoryBufferSize)
  {
    CopyMemory(&(zgfx->HistoryBuffer[zgfx->HistoryIndex]), src, count);
 
    zgfx->HistoryIndex += WINPR_ASSERTING_INT_CAST(uint32_t, count);
    if (zgfx->HistoryIndex == zgfx->HistoryBufferSize)
      zgfx->HistoryIndex = 0;
  }
  else
  {
    const UINT32 front = zgfx->HistoryBufferSize - zgfx->HistoryIndex;
    CopyMemory(&(zgfx->HistoryBuffer[zgfx->HistoryIndex]), src, front);
    CopyMemory(zgfx->HistoryBuffer, &src[front], count - front);
    zgfx->HistoryIndex = (UINT32)(count - front);
  }
}
 
static INLINE void zgfx_history_buffer_ring_read(ZGFX_CONTEXT* WINPR_RESTRICT zgfx, int offset,
                                                 BYTE* WINPR_RESTRICT dst, UINT32 count)
{
  UINT32 front = 0;
  UINT32 index = 0;
  INT32 bytes = 0;
  UINT32 valid = 0;
  INT32 bytesLeft = 0;
  BYTE* dptr = dst;
  BYTE* origDst = dst;
 
  if ((count <= 0) || (count > INT32_MAX))
    return;
 
  bytesLeft = (INT32)count;
  index = (zgfx->HistoryIndex + zgfx->HistoryBufferSize -
           WINPR_ASSERTING_INT_CAST(uint32_t, offset)) %
          zgfx->HistoryBufferSize;
  bytes = MIN(bytesLeft, offset);
 
  if ((index + WINPR_ASSERTING_INT_CAST(uint32_t, bytes)) <= zgfx->HistoryBufferSize)
  {
    CopyMemory(dptr, &(zgfx->HistoryBuffer[index]), WINPR_ASSERTING_INT_CAST(size_t, bytes));
  }
  else
  {
    front = zgfx->HistoryBufferSize - index;
    CopyMemory(dptr, &(zgfx->HistoryBuffer[index]), front);
    CopyMemory(&dptr[front], zgfx->HistoryBuffer,
               WINPR_ASSERTING_INT_CAST(uint32_t, bytes) - front);
  }
 
  if ((bytesLeft -= bytes) == 0)
    return;
 
  dptr += bytes;
  valid = WINPR_ASSERTING_INT_CAST(uint32_t, bytes);
 
  do
  {
    bytes = WINPR_ASSERTING_INT_CAST(int32_t, valid);
 
    if (bytes > bytesLeft)
      bytes = bytesLeft;
 
    CopyMemory(dptr, origDst, WINPR_ASSERTING_INT_CAST(size_t, bytes));
    dptr += bytes;
    valid <<= 1;
  } while ((bytesLeft -= bytes) > 0);
}
 
static INLINE BOOL zgfx_decompress_segment(ZGFX_CONTEXT* WINPR_RESTRICT zgfx,
                                           wStream* WINPR_RESTRICT stream, size_t segmentSize)
{
  BYTE c = 0;
  BYTE flags = 0;
  UINT32 extra = 0;
  int opIndex = 0;
  UINT32 haveBits = 0;
  UINT32 inPrefix = 0;
  UINT32 count = 0;
  UINT32 distance = 0;
  BYTE* pbSegment = NULL;
 
  WINPR_ASSERT(zgfx);
  WINPR_ASSERT(stream);
 
  if (segmentSize < 2)
    return FALSE;
 
  const size_t cbSegment = segmentSize - 1;
 
  if (!Stream_CheckAndLogRequiredLength(TAG, stream, segmentSize) || (segmentSize > UINT32_MAX))
    return FALSE;
 
  Stream_Read_UINT8(stream, flags); /* header (1 byte) */
  zgfx->OutputCount = 0;
  pbSegment = Stream_Pointer(stream);
  if (!Stream_SafeSeek(stream, cbSegment))
    return FALSE;
 
  if (!(flags & PACKET_COMPRESSED))
  {
    zgfx_history_buffer_ring_write(zgfx, pbSegment, cbSegment);
 
    if (cbSegment > sizeof(zgfx->OutputBuffer))
      return FALSE;
 
    CopyMemory(zgfx->OutputBuffer, pbSegment, cbSegment);
    zgfx->OutputCount = (UINT32)cbSegment;
    return TRUE;
  }
 
  zgfx->pbInputCurrent = pbSegment;
  zgfx->pbInputEnd = &pbSegment[cbSegment - 1];
  /* NumberOfBitsToDecode = ((NumberOfBytesToDecode - 1) * 8) - ValueOfLastByte */
  const size_t bits = 8u * (cbSegment - 1u);
  if (bits > UINT32_MAX)
    return FALSE;
  if (bits < *zgfx->pbInputEnd)
    return FALSE;
 
  zgfx->cBitsRemaining = (UINT32)(bits - *zgfx->pbInputEnd);
  zgfx->cBitsCurrent = 0;
  zgfx->BitsCurrent = 0;
 
  while (zgfx->cBitsRemaining)
  {
    haveBits = 0;
    inPrefix = 0;
 
    for (opIndex = 0; ZGFX_TOKEN_TABLE[opIndex].prefixLength != 0; opIndex++)
    {
      while (haveBits < ZGFX_TOKEN_TABLE[opIndex].prefixLength)
      {
        zgfx_GetBits(zgfx, 1);
        inPrefix = (inPrefix << 1) + zgfx->bits;
        haveBits++;
      }
 
      if (inPrefix == ZGFX_TOKEN_TABLE[opIndex].prefixCode)
      {
        if (ZGFX_TOKEN_TABLE[opIndex].tokenType == 0)
        {
          /* Literal */
          zgfx_GetBits(zgfx, ZGFX_TOKEN_TABLE[opIndex].valueBits);
          c = (BYTE)(ZGFX_TOKEN_TABLE[opIndex].valueBase + zgfx->bits);
          zgfx->HistoryBuffer[zgfx->HistoryIndex] = c;
 
          if (++zgfx->HistoryIndex == zgfx->HistoryBufferSize)
            zgfx->HistoryIndex = 0;
 
          if (zgfx->OutputCount >= sizeof(zgfx->OutputBuffer))
            return FALSE;
 
          zgfx->OutputBuffer[zgfx->OutputCount++] = c;
        }
        else
        {
          zgfx_GetBits(zgfx, ZGFX_TOKEN_TABLE[opIndex].valueBits);
          distance = ZGFX_TOKEN_TABLE[opIndex].valueBase + zgfx->bits;
 
          if (distance != 0)
          {
            /* Match */
            zgfx_GetBits(zgfx, 1);
 
            if (zgfx->bits == 0)
            {
              count = 3;
            }
            else
            {
              count = 4;
              extra = 2;
              zgfx_GetBits(zgfx, 1);
 
              while (zgfx->bits == 1)
              {
                count *= 2;
                extra++;
                zgfx_GetBits(zgfx, 1);
              }
 
              zgfx_GetBits(zgfx, extra);
              count += zgfx->bits;
            }
 
            if (count > sizeof(zgfx->OutputBuffer) - zgfx->OutputCount)
              return FALSE;
 
            zgfx_history_buffer_ring_read(zgfx, WINPR_ASSERTING_INT_CAST(int, distance),
                                          &(zgfx->OutputBuffer[zgfx->OutputCount]),
                                          count);
            zgfx_history_buffer_ring_write(
                zgfx, &(zgfx->OutputBuffer[zgfx->OutputCount]), count);
            zgfx->OutputCount += count;
          }
          else
          {
            /* Unencoded */
            zgfx_GetBits(zgfx, 15);
            count = zgfx->bits;
            zgfx->cBitsRemaining -= zgfx->cBitsCurrent;
            zgfx->cBitsCurrent = 0;
            zgfx->BitsCurrent = 0;
 
            if (count > sizeof(zgfx->OutputBuffer) - zgfx->OutputCount)
              return FALSE;
            else if (count > zgfx->cBitsRemaining / 8)
              return FALSE;
            else if (zgfx->pbInputCurrent + count > zgfx->pbInputEnd)
              return FALSE;
 
            CopyMemory(&(zgfx->OutputBuffer[zgfx->OutputCount]), zgfx->pbInputCurrent,
                       count);
            zgfx_history_buffer_ring_write(zgfx, zgfx->pbInputCurrent, count);
            zgfx->pbInputCurrent += count;
            zgfx->cBitsRemaining -= (8 * count);
            zgfx->OutputCount += count;
          }
        }
 
        break;
      }
    }
  }
 
  return TRUE;
}
 
static INLINE BOOL zgfx_append(ZGFX_CONTEXT* WINPR_RESTRICT zgfx,
                               BYTE** WINPR_RESTRICT ppConcatenated, size_t uncompressedSize,
                               size_t* WINPR_RESTRICT pUsed)
{
  WINPR_ASSERT(zgfx);
  WINPR_ASSERT(ppConcatenated);
  WINPR_ASSERT(pUsed);
 
  const size_t used = *pUsed;
  if (zgfx->OutputCount > UINT32_MAX - used)
    return FALSE;
 
  if (used + zgfx->OutputCount > uncompressedSize)
    return FALSE;
 
  BYTE* tmp = realloc(*ppConcatenated, used + zgfx->OutputCount + 64ull);
  if (!tmp)
    return FALSE;
  *ppConcatenated = tmp;
  CopyMemory(&tmp[used], zgfx->OutputBuffer, zgfx->OutputCount);
  *pUsed = used + zgfx->OutputCount;
  return TRUE;
}
 
int zgfx_decompress(ZGFX_CONTEXT* WINPR_RESTRICT zgfx, const BYTE* WINPR_RESTRICT pSrcData,
                    UINT32 SrcSize, BYTE** WINPR_RESTRICT ppDstData,
                    UINT32* WINPR_RESTRICT pDstSize, WINPR_ATTR_UNUSED UINT32 flags)
{
  int status = -1;
  BYTE descriptor = 0;
  wStream sbuffer = { 0 };
  size_t used = 0;
  BYTE* pConcatenated = NULL;
  wStream* stream = Stream_StaticConstInit(&sbuffer, pSrcData, SrcSize);
 
  WINPR_ASSERT(zgfx);
  WINPR_ASSERT(stream);
  WINPR_ASSERT(ppDstData);
  WINPR_ASSERT(pDstSize);
 
  *ppDstData = NULL;
  *pDstSize = 0;
 
  if (!Stream_CheckAndLogRequiredLength(TAG, stream, 1))
    goto fail;
 
  Stream_Read_UINT8(stream, descriptor); /* descriptor (1 byte) */
 
  if (descriptor == ZGFX_SEGMENTED_SINGLE)
  {
    if (!zgfx_decompress_segment(zgfx, stream, Stream_GetRemainingLength(stream)))
      goto fail;
 
    if (zgfx->OutputCount > 0)
    {
      if (!zgfx_append(zgfx, &pConcatenated, zgfx->OutputCount, &used))
        goto fail;
      if (used != zgfx->OutputCount)
        goto fail;
      *ppDstData = pConcatenated;
      *pDstSize = zgfx->OutputCount;
    }
  }
  else if (descriptor == ZGFX_SEGMENTED_MULTIPART)
  {
    UINT32 segmentSize = 0;
    UINT16 segmentNumber = 0;
    UINT16 segmentCount = 0;
    UINT32 uncompressedSize = 0;
 
    if (!Stream_CheckAndLogRequiredLength(TAG, stream, 6))
      goto fail;
 
    Stream_Read_UINT16(stream, segmentCount);     /* segmentCount (2 bytes) */
    Stream_Read_UINT32(stream, uncompressedSize); /* uncompressedSize (4 bytes) */
 
    for (segmentNumber = 0; segmentNumber < segmentCount; segmentNumber++)
    {
      if (!Stream_CheckAndLogRequiredLength(TAG, stream, sizeof(UINT32)))
        goto fail;
 
      Stream_Read_UINT32(stream, segmentSize); /* segmentSize (4 bytes) */
 
      if (!zgfx_decompress_segment(zgfx, stream, segmentSize))
        goto fail;
 
      if (!zgfx_append(zgfx, &pConcatenated, uncompressedSize, &used))
        goto fail;
    }
 
    if (used != uncompressedSize)
      goto fail;
 
    *ppDstData = pConcatenated;
    *pDstSize = uncompressedSize;
  }
  else
  {
    goto fail;
  }
 
  status = 1;
fail:
  if (status < 0)
    free(pConcatenated);
  return status;
}
 
static BOOL zgfx_compress_segment(WINPR_ATTR_UNUSED ZGFX_CONTEXT* WINPR_RESTRICT zgfx,
                                  wStream* WINPR_RESTRICT s, const BYTE* WINPR_RESTRICT pSrcData,
                                  UINT32 SrcSize, UINT32* WINPR_RESTRICT pFlags)
{
  /* FIXME: Currently compression not implemented. Just copy the raw source */
  if (!Stream_EnsureRemainingCapacity(s, SrcSize + 1))
  {
    WLog_ERR(TAG, "Stream_EnsureRemainingCapacity failed!");
    return FALSE;
  }
 
  (*pFlags) |= ZGFX_PACKET_COMPR_TYPE_RDP8; /* RDP 8.0 compression format */
  Stream_Write_UINT8(s, WINPR_ASSERTING_INT_CAST(uint8_t, *pFlags)); /* header (1 byte) */
  Stream_Write(s, pSrcData, SrcSize);
  return TRUE;
}
 
int zgfx_compress_to_stream(ZGFX_CONTEXT* WINPR_RESTRICT zgfx, wStream* WINPR_RESTRICT sDst,
                            const BYTE* WINPR_RESTRICT pUncompressed, UINT32 uncompressedSize,
                            UINT32* WINPR_RESTRICT pFlags)
{
  int fragment = 0;
  UINT16 maxLength = 0;
  UINT32 totalLength = 0;
  size_t posSegmentCount = 0;
  const BYTE* pSrcData = NULL;
  int status = 0;
  maxLength = ZGFX_SEGMENTED_MAXSIZE;
  totalLength = uncompressedSize;
  pSrcData = pUncompressed;
 
  for (; (totalLength > 0) || (fragment == 0); fragment++)
  {
    size_t posDstSize = 0;
    size_t posDataStart = 0;
 
    const UINT32 SrcSize = (totalLength > maxLength) ? maxLength : totalLength;
    posDstSize = 0;
    totalLength -= SrcSize;
 
    /* Ensure we have enough space for headers */
    if (!Stream_EnsureRemainingCapacity(sDst, 12))
    {
      WLog_ERR(TAG, "Stream_EnsureRemainingCapacity failed!");
      return -1;
    }
 
    if (fragment == 0)
    {
      /* First fragment */
      /* descriptor (1 byte) */
      Stream_Write_UINT8(sDst, (totalLength == 0) ? ZGFX_SEGMENTED_SINGLE
                                                  : ZGFX_SEGMENTED_MULTIPART);
 
      if (totalLength > 0)
      {
        posSegmentCount = Stream_GetPosition(sDst); /* segmentCount (2 bytes) */
        Stream_Seek(sDst, 2);
        Stream_Write_UINT32(sDst, uncompressedSize); /* uncompressedSize (4 bytes) */
      }
    }
 
    if (fragment > 0 || totalLength > 0)
    {
      /* Multipart */
      posDstSize = Stream_GetPosition(sDst); /* size (4 bytes) */
      Stream_Seek(sDst, 4);
    }
 
    posDataStart = Stream_GetPosition(sDst);
 
    if (!zgfx_compress_segment(zgfx, sDst, pSrcData, SrcSize, pFlags))
      return -1;
 
    if (posDstSize)
    {
      /* Fill segment data size */
      const size_t DstSize = Stream_GetPosition(sDst) - posDataStart;
      if (DstSize > UINT32_MAX)
        return -1;
      Stream_SetPosition(sDst, posDstSize);
      Stream_Write_UINT32(sDst, (UINT32)DstSize);
      Stream_SetPosition(sDst, posDataStart + DstSize);
    }
 
    pSrcData += SrcSize;
  }
 
  Stream_SealLength(sDst);
 
  /* fill back segmentCount */
  if (posSegmentCount)
  {
    Stream_SetPosition(sDst, posSegmentCount);
    Stream_Write_UINT16(sDst, WINPR_ASSERTING_INT_CAST(uint16_t, fragment));
    Stream_SetPosition(sDst, Stream_Length(sDst));
  }
 
  return status;
}
 
int zgfx_compress(ZGFX_CONTEXT* WINPR_RESTRICT zgfx, const BYTE* WINPR_RESTRICT pSrcData,
                  UINT32 SrcSize, BYTE** WINPR_RESTRICT ppDstData, UINT32* WINPR_RESTRICT pDstSize,
                  UINT32* WINPR_RESTRICT pFlags)
{
  int status = 0;
  wStream* s = Stream_New(NULL, SrcSize);
  status = zgfx_compress_to_stream(zgfx, s, pSrcData, SrcSize, pFlags);
  const size_t pos = Stream_GetPosition(s);
  if (pos > UINT32_MAX)
    status = -1;
  else
  {
    (*ppDstData) = Stream_Buffer(s);
    (*pDstSize) = (UINT32)pos;
  }
  Stream_Free(s, FALSE);
  return status;
}
 
void zgfx_context_reset(ZGFX_CONTEXT* WINPR_RESTRICT zgfx, WINPR_ATTR_UNUSED BOOL flush)
{
  zgfx->HistoryIndex = 0;
}
 
ZGFX_CONTEXT* zgfx_context_new(BOOL Compressor)
{
  ZGFX_CONTEXT* zgfx = NULL;
  zgfx = (ZGFX_CONTEXT*)calloc(1, sizeof(ZGFX_CONTEXT));
 
  if (zgfx)
  {
    zgfx->Compressor = Compressor;
    zgfx->HistoryBufferSize = sizeof(zgfx->HistoryBuffer);
    zgfx_context_reset(zgfx, FALSE);
  }
 
  return zgfx;
}
 
void zgfx_context_free(ZGFX_CONTEXT* zgfx)
{
  free(zgfx);
}