epan/tvbuff_rdp.c

   1 /* tvbuff_rdp.c
   2  * Decompression routines used in RDP
   3  * Copyright 2021, David Fort
   4  *
   5  * Wireshark - Network traffic analyzer
   6  * By Gerald Combs <gerald@wireshark.org>
   7  * Copyright 1998 Gerald Combs
   8  *
   9  * SPDX-License-Identifier: GPL-2.0-or-later
  10  */
  11 #include "config.h"
  12
  13 #include <glib.h>
  14 #include <stdbool.h>
  15
  16 #include <epan/tvbuff.h>
  17 #include <epan/proto.h>
  18 #include <epan/tvbuff_rdp.h>
  19 #include <wsutil/array.h>
  20
  21 enum {
  22         ZGFX_SEGMENTED_SINGLE = 0xe0,
  23         ZGFX_SEGMENTED_MULTIPART = 0xe1,
  24
  25         ZGX_PACKET_COMPRESSED = 0x20,
  26 };
  27
  28
  29 typedef struct {
  30         tvbuff_t *input;
  31         unsigned offset;
  32         unsigned remainingBits;
  33         uint32_t currentValue;
  34         unsigned currentBits;
  35 } bitstream_t;
  36
  37 static void
  38 bitstream_init(bitstream_t *b, tvbuff_t *input, unsigned blen) {
  39         b->input = input;
  40         b->offset = 0;
  41         b->remainingBits = blen;
  42         b->currentValue = 0;
  43         b->currentBits = 0;
  44 }
  45
  46 static uint32_t
  47 bitstream_getbits(bitstream_t *b, uint8_t nbits, bool *ok) {
  48         uint32_t ret = 0;
  49
  50         if (nbits > b->remainingBits) {
  51                 *ok = false;
  52                 return 0;
  53         }
  54
  55         while (b->currentBits < nbits) {
  56                 if (!tvb_reported_length_remaining(b->input, b->offset)) {
  57                         *ok = false;
  58                         return 0;
  59                 }
  60
  61                 b->currentValue <<= 8;
  62                 b->currentValue += tvb_get_uint8(b->input, b->offset++);
  63
  64                 b->currentBits += 8;
  65         }
  66
  67         *ok = true;
  68         ret = b->currentValue >> (b->currentBits-nbits);
  69         b->currentBits -= nbits;
  70         b->remainingBits -= nbits;
  71         b->currentValue &= (1 << b->currentBits) - 1;
  72
  73         return ret;
  74 }
  75
  76 static bool
  77 bitstream_copyraw(bitstream_t *b, uint8_t *dest, int nbytes)
  78 {
  79         if (tvb_captured_length_remaining(b->input, b->offset) < nbytes)
  80                 return false;
  81
  82         tvb_memcpy(b->input, dest, b->offset, nbytes);
  83
  84         return true;
  85 }
  86
  87 static bool
  88 bitstream_copyraw_advance(bitstream_t *b, uint8_t *dest, unsigned nbytes)
  89 {
  90         if (!bitstream_copyraw(b, dest, nbytes))
  91                 return false;
  92
  93         b->offset += nbytes;
  94         b->remainingBits -= (nbytes * 8);
  95         return true;
  96 }
  97
  98
  99 static void
 100 bitstream_realign(bitstream_t *b) {
 101         b->remainingBits -= b->currentBits;
 102         b->currentBits = 0;
 103         b->currentValue = 0;
 104 }
 105
 106 typedef struct {
 107         uint32_t prefixLength;
 108         uint32_t prefixCode;
 109         uint32_t valueBits;
 110         uint32_t valueBase;
 111 } zgfx_token_t;
 112
 113 static const zgfx_token_t ZGFX_LITERAL_TABLE[] = {
 114         // prefixLength prefixCode valueBits valueBase
 115         { 5,  24, 0, 0x00 },   // 11000
 116         { 5,  25, 0, 0x01 },   // 11001
 117         { 6,  52, 0, 0x02 },   // 110100
 118         { 6,  53, 0, 0x03 },   // 110101
 119         { 6,  54, 0, 0xFF },   // 110110
 120         { 7, 110, 0, 0x04 },   // 1101110
 121         { 7, 111, 0, 0x05 },   // 1101111
 122         { 7, 112, 0, 0x06 },   // 1110000
 123         { 7, 113, 0, 0x07 },   // 1110001
 124         { 7, 114, 0, 0x08 },   // 1110010
 125         { 7, 115, 0, 0x09 },   // 1110011
 126         { 7, 116, 0, 0x0A },   // 1110100
 127         { 7, 117, 0, 0x0B },   // 1110101
 128         { 7, 118, 0, 0x3A },   // 1110110
 129         { 7, 119, 0, 0x3B },   // 1110111
 130         { 7, 120, 0, 0x3C },   // 1111000
 131         { 7, 121, 0, 0x3D },   // 1111001
 132         { 7, 122, 0, 0x3E },   // 1111010
 133         { 7, 123, 0, 0x3F },   // 1111011
 134         { 7, 124, 0, 0x40 },   // 1111100
 135         { 7, 125, 0, 0x80 },   // 1111101
 136         { 8, 252, 0, 0x0C },   // 11111100
 137         { 8, 253, 0, 0x38 },   // 11111101
 138         { 8, 254, 0, 0x39 },   // 11111110
 139         { 8, 255, 0, 0x66 },   // 11111111
 140 };
 141
 142 static const zgfx_token_t ZGFX_MATCH_TABLE[] = {
 143         // prefixLength prefixCode valueBits tokenType valueBase
 144         { 5, 17, 5, 0 },          // 10001
 145         { 5, 18, 7, 32 },         // 10010
 146         { 5, 19, 9, 160 },        // 10011
 147         { 5, 20, 10, 672 },       // 10100
 148         { 5, 21, 12, 1696 },      // 10101
 149         { 6, 44, 14, 5792 },      // 101100
 150         { 6, 45, 15, 22176 },     // 101101
 151         { 7, 92, 18, 54944 },     // 1011100
 152         { 7, 93, 20, 317088 },    // 1011101
 153         { 8, 188, 20, 1365664 },  // 10111100
 154         { 8, 189, 21, 2414240 },  // 10111101
 155         { 9, 380, 22, 4511392 },  // 101111100
 156         { 9, 381, 23, 8705696 },  // 101111101
 157         { 9, 382, 24, 17094304 }, // 101111110
 158 };
 159
 160
 161 struct _zgfx_context_t{
 162         uint8_t historyBuffer[2500000];
 163         uint32_t historyIndex;
 164         uint32_t historyBufferSize;
 165
 166         uint32_t outputCount;
 167         uint8_t outputSegment[65536];
 168 };
 169
 170 zgfx_context_t *zgfx_context_new(wmem_allocator_t *allocator) {
 171         zgfx_context_t *ret = wmem_alloc0(allocator, sizeof(*ret));
 172         ret->historyBufferSize = sizeof(ret->historyBuffer);
 173         return ret;
 174 }
 175
 176 static void
 177 zgfx_write_history_literal(zgfx_context_t *zgfx, uint8_t c)
 178 {
 179         zgfx->historyBuffer[zgfx->historyIndex] = c;
 180         zgfx->historyIndex = (zgfx->historyIndex + 1) % zgfx->historyBufferSize;
 181 }
 182
 183 static void
 184 zgfx_write_history_buffer_tvb(zgfx_context_t *zgfx, tvbuff_t *src, uint32_t count)
 185 {
 186         int src_offset = 0;
 187         uint32_t front;
 188
 189         if (count > zgfx->historyBufferSize) {
 190                 const uint32_t residue = count - zgfx->historyBufferSize;
 191                 count = zgfx->historyBufferSize;
 192                 src_offset += residue;
 193                 zgfx->historyIndex = (zgfx->historyIndex + residue) % zgfx->historyBufferSize;
 194         }
 195
 196         if (zgfx->historyIndex + count <= zgfx->historyBufferSize)
 197         {
 198                 tvb_memcpy(src, &(zgfx->historyBuffer[zgfx->historyIndex]), src_offset, count);
 199         }
 200         else
 201         {
 202                 front = zgfx->historyBufferSize - zgfx->historyIndex;
 203                 tvb_memcpy(src, &(zgfx->historyBuffer[zgfx->historyIndex]), src_offset, front);
 204                 tvb_memcpy(src, &(zgfx->historyBuffer), src_offset + count, count - front);
 205         }
 206
 207         zgfx->historyIndex = (zgfx->historyIndex + count) % zgfx->historyBufferSize;
 208 }
 209
 210
 211 static void
 212 zgfx_write_history_buffer(zgfx_context_t *zgfx, const uint8_t *src, uint32_t count)
 213 {
 214         uint32_t front;
 215
 216         if (count > zgfx->historyBufferSize) {
 217                 const uint32_t residue = count - zgfx->historyBufferSize;
 218                 count = zgfx->historyBufferSize;
 219                 zgfx->historyIndex = (zgfx->historyIndex + residue) % zgfx->historyBufferSize;
 220         }
 221
 222         if (zgfx->historyIndex + count <= zgfx->historyBufferSize)
 223         {
 224                 memcpy(&(zgfx->historyBuffer[zgfx->historyIndex]), src, count);
 225         }
 226         else
 227         {
 228                 front = zgfx->historyBufferSize - zgfx->historyIndex;
 229                 memcpy(&(zgfx->historyBuffer[zgfx->historyIndex]), src, front);
 230                 memcpy(&(zgfx->historyBuffer), src + front, count - front);
 231         }
 232
 233         zgfx->historyIndex = (zgfx->historyIndex + count) % zgfx->historyBufferSize;
 234 }
 235
 236
 237 static bool
 238 zgfx_write_literal(zgfx_context_t *zgfx, uint8_t c)
 239 {
 240         if (zgfx->outputCount == 65535)
 241                 return false;
 242
 243         zgfx->outputSegment[zgfx->outputCount++] = c;
 244
 245         zgfx_write_history_literal(zgfx, c);
 246         return true;
 247 }
 248
 249 static bool
 250 zgfx_write_raw(zgfx_context_t *zgfx, bitstream_t *b, uint32_t count)
 251 {
 252         uint32_t rest, tocopy;
 253
 254         /* first copy in the output buffer */
 255         if (zgfx->outputCount > 65535 - count)
 256                 return false;
 257
 258         if (!bitstream_copyraw(b, &(zgfx->outputSegment[zgfx->outputCount]), count))
 259                 return false;
 260
 261         zgfx->outputCount += count;
 262
 263         /* then update the history buffer */
 264         rest = (zgfx->historyBufferSize - zgfx->historyIndex);
 265         tocopy = count;
 266         if (rest < count)
 267                 tocopy = rest;
 268
 269         if (!bitstream_copyraw_advance(b, &(zgfx->historyBuffer[zgfx->historyIndex]), tocopy))
 270                 return false;
 271
 272         zgfx->historyIndex = (zgfx->historyIndex + tocopy) % zgfx->historyBufferSize;
 273         count -= tocopy;
 274         if (count) {
 275                 if (!bitstream_copyraw_advance(b, &(zgfx->historyBuffer[zgfx->historyIndex]), tocopy))
 276                         return false;
 277
 278                 zgfx->historyIndex = (zgfx->historyIndex + tocopy) % zgfx->historyBufferSize;
 279         }
 280
 281         return true;
 282 }
 283
 284 static bool
 285 zgfx_write_from_history(zgfx_context_t *zgfx, uint32_t distance, uint32_t count)
 286 {
 287         unsigned idx;
 288         uint32_t remainingCount, copyTemplateSize, toCopy;
 289         uint8_t *outputPtr;
 290
 291         if (zgfx->outputCount > 65535 - count)
 292                 return false;
 293
 294         remainingCount = count;
 295         idx = (zgfx->historyIndex + zgfx->historyBufferSize - distance) % zgfx->historyBufferSize;
 296         copyTemplateSize = (distance > count) ? count : distance;
 297
 298         /* first do copy a single copy in output */
 299         outputPtr = &(zgfx->outputSegment[zgfx->outputCount]);
 300         toCopy = copyTemplateSize;
 301         if (idx + toCopy < zgfx->historyBufferSize) {
 302                 memcpy(outputPtr, &(zgfx->historyBuffer[idx]), toCopy);
 303         } else {
 304                 uint32_t partial = zgfx->historyBufferSize - idx;
 305                 memcpy(outputPtr, &(zgfx->historyBuffer[idx]), partial);
 306                 memcpy(outputPtr + partial, zgfx->historyBuffer, toCopy - partial);
 307         }
 308         outputPtr += toCopy;
 309         remainingCount -= toCopy;
 310
 311         /* then duplicate output as much as needed by count, at each loop turn we double
 312          * the size of the template we can copy */
 313         while (remainingCount) {
 314                 toCopy = (remainingCount < copyTemplateSize) ? remainingCount : copyTemplateSize;
 315                 memcpy(outputPtr, &(zgfx->outputSegment[zgfx->outputCount]), toCopy);
 316
 317                 outputPtr += toCopy;
 318                 remainingCount -= toCopy;
 319                 copyTemplateSize *= 2;
 320         }
 321
 322         /* let's update the history from output and update counters */
 323         zgfx_write_history_buffer(zgfx, &(zgfx->outputSegment[zgfx->outputCount]), count);
 324         zgfx->outputCount += count;
 325         return true;
 326 }
 327
 328
 329 static bool
 330 rdp8_decompress_segment(zgfx_context_t *zgfx, tvbuff_t *tvb)
 331 {
 332         bitstream_t bitstream;
 333         int offset = 0;
 334         int len = tvb_reported_length(tvb);
 335         uint8_t flags = tvb_get_uint8(tvb, offset);
 336         uint8_t v;
 337         offset++;
 338         len--;
 339
 340         if (!(flags & ZGX_PACKET_COMPRESSED)) {
 341                 tvbuff_t *raw = tvb_new_subset_remaining(tvb, 1);
 342                 zgfx_write_history_buffer_tvb(zgfx, raw, len);
 343
 344                 tvb_memcpy(tvb, zgfx->outputSegment, 1, len);
 345                 zgfx->outputCount += len;
 346                 return true;
 347         }
 348
 349         v = tvb_get_uint8(tvb, offset + len - 1);
 350         if (v > 7)
 351                 return false;
 352         len--;
 353
 354         bitstream_init(&bitstream, tvb_new_subset_length(tvb, offset, len), (len * 8) - v);
 355         while (bitstream.remainingBits) {
 356                 bool ok, ismatch, found;
 357                 uint32_t bits_val = bitstream_getbits(&bitstream, 1, &ok);
 358                 uint32_t inPrefix;
 359                 const zgfx_token_t *tokens;
 360                 int ntokens, i;
 361                 uint32_t prefixBits;
 362
 363                 if (!ok)
 364                         return false;
 365
 366                 // 0 - literal
 367                 if (bits_val == 0) {
 368
 369                         bits_val = bitstream_getbits(&bitstream, 8, &ok);
 370                         if (!zgfx_write_literal(zgfx, bits_val))
 371                                 return false;
 372                         continue;
 373                 }
 374
 375                 // 1x - match or literal branch
 376                 bits_val = bitstream_getbits(&bitstream, 1, &ok);
 377                 if (bits_val == 0) {
 378                         // 10 - match
 379                         ismatch = true;
 380                         tokens = ZGFX_MATCH_TABLE;
 381                         ntokens = array_length(ZGFX_MATCH_TABLE);
 382                         inPrefix = 2;
 383                 } else {
 384                         // 11 - literal
 385                         ismatch = false;
 386                         tokens = ZGFX_LITERAL_TABLE;
 387                         ntokens = array_length(ZGFX_LITERAL_TABLE);
 388                         inPrefix = 3;
 389                 }
 390
 391                 prefixBits = 2;
 392                 found = false;
 393                 for (i = 0; i < ntokens; i++) {
 394                         if (prefixBits != tokens[i].prefixLength) {
 395                                 uint32_t missingBits = (tokens[i].prefixLength - prefixBits);
 396                                 inPrefix <<= missingBits;
 397                                 inPrefix |= bitstream_getbits(&bitstream, missingBits, &ok);
 398                                 if (!ok)
 399                                         return false;
 400                                 prefixBits = tokens[i].prefixLength;
 401                         }
 402
 403                         if (inPrefix == tokens[i].prefixCode) {
 404                                 found = true;
 405                                 break;
 406                         }
 407                 }
 408
 409                 if (!found) // TODO: is it an error ?
 410                         continue;
 411
 412                 if (ismatch) {
 413                         /* It's a match */
 414                         uint32_t count, distance, extra = 0;
 415
 416                         distance = tokens[i].valueBase + bitstream_getbits(&bitstream, tokens[i].valueBits, &ok);
 417                         if (!ok)
 418                                 return false;
 419
 420                         if (distance != 0) {
 421                                 bits_val = bitstream_getbits(&bitstream, 1, &ok);
 422                                 if (!ok)
 423                                         return false;
 424
 425                                 if (bits_val == 0) {
 426                                         count = 3;
 427                                 } else {
 428                                         count = 4;
 429                                         extra = 2;
 430
 431                                         bits_val = bitstream_getbits(&bitstream, 1, &ok);
 432                                         if (!ok)
 433                                                 return false;
 434
 435                                         while (bits_val == 1) {
 436                                                 count *= 2;
 437                                                 extra ++;
 438                                                 bits_val = bitstream_getbits(&bitstream, 1, &ok);
 439                                                 if (!ok)
 440                                                         return false;
 441                                         }
 442
 443                                         count += bitstream_getbits(&bitstream, extra, &ok);
 444                                         if (!ok)
 445                                                 return false;
 446                                 }
 447
 448                                 if (count > sizeof(zgfx->outputSegment) - zgfx->outputCount)
 449                                         return false;
 450
 451                                 if (!zgfx_write_from_history(zgfx, distance, count))
 452                                         return false;
 453                         } else {
 454                                 /* Unencoded */
 455                                 count = bitstream_getbits(&bitstream, 15, &ok);
 456                                 if (!ok)
 457                                         return false;
 458
 459                                 bitstream_realign(&bitstream);
 460                                 if (!zgfx_write_raw(zgfx, &bitstream, count))
 461                                         return false;
 462                         }
 463                 } else {
 464                         /* literal */
 465                         bits_val = tokens[i].valueBase;
 466                         if (!zgfx_write_literal(zgfx, bits_val))
 467                                 return false;
 468                 }
 469         }
 470
 471         return true;
 472 }
 473
 474
 475
 476 tvbuff_t *
 477 rdp8_decompress(zgfx_context_t *zgfx, wmem_allocator_t *allocator, tvbuff_t *tvb, unsigned offset)
 478 {
 479         void *output;
 480         uint8_t descriptor;
 481
 482         descriptor = tvb_get_uint8(tvb, offset);
 483         offset++;
 484
 485         switch (descriptor) {
 486         case ZGFX_SEGMENTED_SINGLE:
 487                 zgfx->outputCount = 0;
 488                 if (!rdp8_decompress_segment(zgfx, tvb_new_subset_remaining(tvb, offset)))
 489                         return NULL;
 490
 491                 output = wmem_alloc(allocator, zgfx->outputCount);
 492                 memcpy(output, zgfx->outputSegment, zgfx->outputCount);
 493                 return tvb_new_real_data(output, zgfx->outputCount, zgfx->outputCount);
 494
 495         case ZGFX_SEGMENTED_MULTIPART: {
 496                 uint16_t segment_count, i;
 497                 uint32_t output_consumed, uncompressed_size;
 498                 uint8_t *output_ptr;
 499
 500                 segment_count = tvb_get_uint16(tvb, offset, ENC_LITTLE_ENDIAN);
 501                 offset += 2;
 502                 uncompressed_size = tvb_get_uint32(tvb, offset, ENC_LITTLE_ENDIAN);
 503                 offset += 4;
 504
 505                 output = output_ptr = wmem_alloc(allocator, uncompressed_size);
 506                 output_consumed = 0;
 507                 for (i = 0; i < segment_count; i++) {
 508                         uint32_t segment_size = tvb_get_uint32(tvb, offset, ENC_LITTLE_ENDIAN);
 509                         offset += 4;
 510
 511                         zgfx->outputCount = 0;
 512                         if (!rdp8_decompress_segment(zgfx, tvb_new_subset_length(tvb, offset, segment_size)))
 513                                 return NULL;
 514
 515                         output_consumed += zgfx->outputCount;
 516                         if (output_consumed > uncompressed_size) {
 517                                 // TODO: error message ?
 518                                 return NULL;
 519                         }
 520                         memcpy(output_ptr, zgfx->outputSegment, zgfx->outputCount);
 521
 522                         offset += segment_size;
 523                         output_ptr += zgfx->outputCount;
 524                 }
 525                 return tvb_new_real_data(output, uncompressed_size, uncompressed_size);
 526         }
 527         default:
 528                 return tvb;
 529         }
 530 }