wimboot/wimboot-2.7.3/src/lzx.c

   1 /*
   2  * Copyright (C) 2014 Michael Brown <mbrown@fensystems.co.uk>.
   3  *
   4  * This program is free software; you can redistribute it and/or
   5  * modify it under the terms of the GNU General Public License as
   6  * published by the Free Software Foundation; either version 2 of the
   7  * License, or (at your option) any later version.
   8  *
   9  * This program is distributed in the hope that it will be useful, but
  10  * WITHOUT ANY WARRANTY; without even the implied warranty of
  11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  12  * General Public License for more details.
  13  *
  14  * You should have received a copy of the GNU General Public License
  15  * along with this program; if not, write to the Free Software
  16  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
  17  * 02110-1301, USA.
  18  */
  19
  20 /**
  21  * @file
  22  *
  23  * LZX decompression
  24  *
  25  * This algorithm is derived jointly from the document "[MS-PATCH]:
  26  * LZX DELTA Compression and Decompression", available from
  27  *
  28  *     http://msdn.microsoft.com/en-us/library/cc483133.aspx
  29  *
  30  * and from the file lzx-decompress.c in the wimlib source code.
  31  *
  32  */
  33
  34 #include <stdint.h>
  35 #include <stddef.h>
  36 #include <string.h>
  37 #include <stdio.h>
  38 #include "wimboot.h"
  39 #include "huffman.h"
  40 #include "lzx.h"
  41
  42 /** Base positions, indexed by position slot */
  43 static unsigned int lzx_position_base[LZX_POSITION_SLOTS];
  44
  45 /**
  46  * Attempt to accumulate bits from LZX bitstream
  47  *
  48  * @v lzx               Decompressor
  49  * @v bits              Number of bits to accumulate
  50  * @v norm_value        Accumulated value (normalised to 16 bits)
  51  *
  52  * Note that there may not be sufficient accumulated bits in the
  53  * bitstream; callers must check that sufficient bits are available
  54  * before using the value.
  55  */
  56 static int lzx_accumulate ( struct lzx *lzx, unsigned int bits ) {
  57         const uint16_t *src16;
  58
  59         /* Accumulate more bits if required */
  60         if ( ( lzx->bits < bits ) &&
  61              ( lzx->input.offset < lzx->input.len ) ) {
  62                 src16 = ( ( void * ) lzx->input.data + lzx->input.offset );
  63                 lzx->input.offset += sizeof ( *src16 );
  64                 lzx->accumulator |= ( *src16 << ( 16 - lzx->bits ) );
  65                 lzx->bits += 16;
  66         }
  67
  68         return ( lzx->accumulator >> 16 );
  69 }
  70
  71 /**
  72  * Consume accumulated bits from LZX bitstream
  73  *
  74  * @v lzx               Decompressor
  75  * @v bits              Number of bits to consume
  76  * @ret rc              Return status code
  77  */
  78 static int lzx_consume ( struct lzx *lzx, unsigned int bits ) {
  79
  80         /* Fail if insufficient bits are available */
  81         if ( lzx->bits < bits ) {
  82                 DBG ( "LZX input overrun in %#zx/%#zx out %#zx)\n",
  83                       lzx->input.offset, lzx->input.len, lzx->output.offset );
  84                 return -1;
  85         }
  86
  87         /* Consume bits */
  88         lzx->accumulator <<= bits;
  89         lzx->bits -= bits;
  90
  91         return 0;
  92 }
  93
  94 /**
  95  * Get bits from LZX bitstream
  96  *
  97  * @v lzx               Decompressor
  98  * @v bits              Number of bits to fetch
  99  * @ret value           Value, or negative error
 100  */
 101 static int lzx_getbits ( struct lzx *lzx, unsigned int bits ) {
 102         int norm_value;
 103         int rc;
 104
 105         /* Accumulate more bits if required */
 106         norm_value = lzx_accumulate ( lzx, bits );
 107
 108         /* Consume bits */
 109         if ( ( rc = lzx_consume ( lzx, bits ) ) != 0 )
 110                 return rc;
 111
 112         return ( norm_value >> ( 16 - bits ) );
 113 }
 114
 115 /**
 116  * Align LZX bitstream for byte access
 117  *
 118  * @v lzx               Decompressor
 119  * @v bits              Minimum number of padding bits
 120  * @ret rc              Return status code
 121  */
 122 static int lzx_align ( struct lzx *lzx, unsigned int bits ) {
 123         int pad;
 124
 125         /* Get padding bits */
 126         pad = lzx_getbits ( lzx, bits );
 127         if ( pad < 0 )
 128                 return pad;
 129
 130         /* Consume all accumulated bits */
 131         lzx_consume ( lzx, lzx->bits );
 132
 133         return 0;
 134 }
 135
 136 /**
 137  * Get bytes from LZX bitstream
 138  *
 139  * @v lzx               Decompressor
 140  * @v data              Data buffer, or NULL
 141  * @v len               Length of data buffer
 142  * @ret rc              Return status code
 143  */
 144 static int lzx_getbytes ( struct lzx *lzx, void *data, size_t len ) {
 145
 146         /* Sanity check */
 147         if ( ( lzx->input.offset + len ) > lzx->input.len ) {
 148                 DBG ( "LZX input overrun in %#zx/%#zx out %#zx)\n",
 149                       lzx->input.offset, lzx->input.len, lzx->output.offset );
 150                 return -1;
 151         }
 152
 153         /* Copy data */
 154         if ( data )
 155                 memcpy ( data, ( lzx->input.data + lzx->input.offset ), len );
 156         lzx->input.offset += len;
 157
 158         return 0;
 159 }
 160
 161 /**
 162  * Decode LZX Huffman-coded symbol
 163  *
 164  * @v lzx               Decompressor
 165  * @v alphabet          Huffman alphabet
 166  * @ret raw             Raw symbol, or negative error
 167  */
 168 static int lzx_decode ( struct lzx *lzx, struct huffman_alphabet *alphabet ) {
 169         struct huffman_symbols *sym;
 170         int huf;
 171         int rc;
 172
 173         /* Accumulate sufficient bits */
 174         huf = lzx_accumulate ( lzx, HUFFMAN_BITS );
 175         if ( huf < 0 )
 176                 return huf;
 177
 178         /* Decode symbol */
 179         sym = huffman_sym ( alphabet, huf );
 180
 181         /* Consume bits */
 182         if ( ( rc = lzx_consume ( lzx, huffman_len ( sym ) ) ) != 0 )
 183                 return rc;
 184
 185         return huffman_raw ( sym, huf );
 186 }
 187
 188 /**
 189  * Generate Huffman alphabet from raw length table
 190  *
 191  * @v lzx               Decompressor
 192  * @v count             Number of symbols
 193  * @v bits              Length of each length (in bits)
 194  * @v lengths           Lengths table to fill in
 195  * @v alphabet          Huffman alphabet to fill in
 196  * @ret rc              Return status code
 197  */
 198 static int lzx_raw_alphabet ( struct lzx *lzx, unsigned int count,
 199                               unsigned int bits, uint8_t *lengths,
 200                               struct huffman_alphabet *alphabet ) {
 201         unsigned int i;
 202         int len;
 203         int rc;
 204
 205         /* Read lengths */
 206         for ( i = 0 ; i < count ; i++ ) {
 207                 len = lzx_getbits ( lzx, bits );
 208                 if ( len < 0 )
 209                         return len;
 210                 lengths[i] = len;
 211         }
 212
 213         /* Generate Huffman alphabet */
 214         if ( ( rc = huffman_alphabet ( alphabet, lengths, count ) ) != 0 )
 215                 return rc;
 216
 217         return 0;
 218 }
 219
 220 /**
 221  * Generate pretree
 222  *
 223  * @v lzx               Decompressor
 224  * @v count             Number of symbols
 225  * @v lengths           Lengths table to fill in
 226  * @ret rc              Return status code
 227  */
 228 static int lzx_pretree ( struct lzx *lzx, unsigned int count,
 229                          uint8_t *lengths ) {
 230         unsigned int i;
 231         unsigned int length;
 232         int dup = 0;
 233         int code;
 234         int rc;
 235
 236         /* Generate pretree alphabet */
 237         if ( ( rc =  lzx_raw_alphabet ( lzx, LZX_PRETREE_CODES,
 238                                         LZX_PRETREE_BITS, lzx->pretree_lengths,
 239                                         &lzx->pretree ) ) != 0 )
 240                 return rc;
 241
 242         /* Read lengths */
 243         for ( i = 0 ; i < count ; i++ ) {
 244
 245                 if ( dup ) {
 246
 247                         /* Duplicate previous length */
 248                         lengths[i] = lengths[ i - 1 ];
 249                         dup--;
 250
 251                 } else {
 252
 253                         /* Get next code */
 254                         code = lzx_decode ( lzx, &lzx->pretree );
 255                         if ( code < 0 )
 256                                 return code;
 257
 258                         /* Interpret code */
 259                         if ( code <= 16 ) {
 260                                 length = ( ( lengths[i] - code + 17 ) % 17 );
 261                         } else if ( code == 17 ) {
 262                                 length = 0;
 263                                 dup = lzx_getbits ( lzx, 4 );
 264                                 if ( dup < 0 )
 265                                         return dup;
 266                                 dup += 3;
 267                         } else if ( code == 18 ) {
 268                                 length = 0;
 269                                 dup = lzx_getbits ( lzx, 5 );
 270                                 if ( dup < 0 )
 271                                         return dup;
 272                                 dup += 19;
 273                         } else if ( code == 19 ) {
 274                                 length = 0;
 275                                 dup = lzx_getbits ( lzx, 1 );
 276                                 if ( dup < 0 )
 277                                         return dup;
 278                                 dup += 3;
 279                                 code = lzx_decode ( lzx, &lzx->pretree );
 280                                 if ( code < 0 )
 281                                         return code;
 282                                 length = ( ( lengths[i] - code + 17 ) % 17 );
 283                         } else {
 284                                 DBG ( "Unrecognised pretree code %d\n", code );
 285                                 return -1;
 286                         }
 287                         lengths[i] = length;
 288                 }
 289         }
 290
 291         /* Sanity check */
 292         if ( dup ) {
 293                 DBG ( "Pretree duplicate overrun\n" );
 294                 return -1;
 295         }
 296
 297         return 0;
 298 }
 299
 300 /**
 301  * Generate aligned offset Huffman alphabet
 302  *
 303  * @v lzx               Decompressor
 304  * @ret rc              Return status code
 305  */
 306 static int lzx_alignoffset_alphabet ( struct lzx *lzx ) {
 307         int rc;
 308
 309         /* Generate aligned offset alphabet */
 310         if ( ( rc = lzx_raw_alphabet ( lzx, LZX_ALIGNOFFSET_CODES,
 311                                        LZX_ALIGNOFFSET_BITS,
 312                                        lzx->alignoffset_lengths,
 313                                        &lzx->alignoffset ) ) != 0 )
 314                 return rc;
 315
 316         return 0;
 317 }
 318
 319 /**
 320  * Generate main Huffman alphabet
 321  *
 322  * @v lzx               Decompressor
 323  * @ret rc              Return status code
 324  */
 325 static int lzx_main_alphabet ( struct lzx *lzx ) {
 326         int rc;
 327
 328         /* Generate literal symbols pretree */
 329         if ( ( rc = lzx_pretree ( lzx, LZX_MAIN_LIT_CODES,
 330                                   lzx->main_lengths.literals ) ) != 0 ) {
 331                 DBG ( "Could not construct main literal pretree\n" );
 332                 return rc;
 333         }
 334
 335         /* Generate remaining symbols pretree */
 336         if ( ( rc = lzx_pretree ( lzx, ( LZX_MAIN_CODES - LZX_MAIN_LIT_CODES ),
 337                                   lzx->main_lengths.remainder ) ) != 0 ) {
 338                 DBG ( "Could not construct main remainder pretree\n" );
 339                 return rc;
 340         }
 341
 342         /* Generate Huffman alphabet */
 343         if ( ( rc = huffman_alphabet ( &lzx->main, lzx->main_lengths.literals,
 344                                        LZX_MAIN_CODES ) ) != 0 ) {
 345                 DBG ( "Could not generate main alphabet\n" );
 346                 return rc;
 347         }
 348
 349         return 0;
 350 }
 351
 352 /**
 353  * Generate length Huffman alphabet
 354  *
 355  * @v lzx               Decompressor
 356  * @ret rc              Return status code
 357  */
 358 static int lzx_length_alphabet ( struct lzx *lzx ) {
 359         int rc;
 360
 361         /* Generate pretree */
 362         if ( ( rc = lzx_pretree ( lzx, LZX_LENGTH_CODES,
 363                                   lzx->length_lengths ) ) != 0 ) {
 364                 DBG ( "Could not generate length pretree\n" );
 365                 return rc;
 366         }
 367
 368         /* Generate Huffman alphabet */
 369         if ( ( rc = huffman_alphabet ( &lzx->length, lzx->length_lengths,
 370                                        LZX_LENGTH_CODES ) ) != 0 ) {
 371                 DBG ( "Could not generate length alphabet\n" );
 372                 return rc;
 373         }
 374
 375         return 0;
 376 }
 377
 378 /**
 379  * Process LZX block header
 380  *
 381  * @v lzx               Decompressor
 382  * @ret rc              Return status code
 383  */
 384 static int lzx_block_header ( struct lzx *lzx ) {
 385         size_t block_len;
 386         int block_type;
 387         int default_len;
 388         int len_high;
 389         int len_low;
 390         int rc;
 391
 392         /* Get block type */
 393         block_type = lzx_getbits ( lzx, LZX_BLOCK_TYPE_BITS );
 394         if ( block_type < 0 )
 395                 return block_type;
 396         lzx->block_type = block_type;
 397
 398         /* Check block length */
 399         default_len = lzx_getbits ( lzx, 1 );
 400         if ( default_len < 0 )
 401                 return default_len;
 402         if ( default_len ) {
 403                 block_len = LZX_DEFAULT_BLOCK_LEN;
 404         } else {
 405                 len_high = lzx_getbits ( lzx, 8 );
 406                 if ( len_high < 0 )
 407                         return len_high;
 408                 len_low = lzx_getbits ( lzx, 8 );
 409                 if ( len_low < 0 )
 410                         return len_low;
 411                 block_len = ( ( len_high << 8 ) | len_low );
 412         }
 413         lzx->output.threshold = ( lzx->output.offset + block_len );
 414
 415         /* Handle block type */
 416         switch ( block_type ) {
 417         case LZX_BLOCK_ALIGNOFFSET :
 418                 /* Generated aligned offset alphabet */
 419                 if ( ( rc = lzx_alignoffset_alphabet ( lzx ) ) != 0 )
 420                         return rc;
 421                 /* Fall through */
 422         case LZX_BLOCK_VERBATIM :
 423                 /* Generate main alphabet */
 424                 if ( ( rc = lzx_main_alphabet ( lzx ) ) != 0 )
 425                         return rc;
 426                 /* Generate lengths alphabet */
 427                 if ( ( rc = lzx_length_alphabet ( lzx ) ) != 0 )
 428                         return rc;
 429                 break;
 430         case LZX_BLOCK_UNCOMPRESSED :
 431                 /* Align input stream */
 432                 if ( ( rc = lzx_align ( lzx, 1 ) ) != 0 )
 433                         return rc;
 434                 /* Read new repeated offsets */
 435                 if ( ( rc = lzx_getbytes ( lzx, &lzx->repeated_offset,
 436                                            sizeof ( lzx->repeated_offset )))!=0)
 437                         return rc;
 438                 break;
 439         default:
 440                 DBG ( "Unrecognised block type %d\n", block_type );
 441                 return -1;
 442         }
 443
 444         return 0;
 445 }
 446
 447 /**
 448  * Process uncompressed data
 449  *
 450  * @v lzx               Decompressor
 451  * @ret rc              Return status code
 452  */
 453 static int lzx_uncompressed ( struct lzx *lzx ) {
 454         void *data;
 455         size_t len;
 456         int rc;
 457
 458         /* Copy bytes */
 459         data = ( lzx->output.data ?
 460                  ( lzx->output.data + lzx->output.offset ) : NULL );
 461         len = ( lzx->output.threshold - lzx->output.offset );
 462         if ( ( rc = lzx_getbytes ( lzx, data, len ) ) != 0 )
 463                 return rc;
 464
 465         /* Align input stream */
 466         if ( len % 2 )
 467                 lzx->input.offset++;
 468
 469         return 0;
 470 }
 471
 472 /**
 473  * Process an LZX token
 474  *
 475  * @v lzx               Decompressor
 476  * @ret rc              Return status code
 477  *
 478  * Variable names are chosen to match the LZX specification
 479  * pseudo-code.
 480  */
 481 static int lzx_token ( struct lzx *lzx ) {
 482         unsigned int length_header;
 483         unsigned int position_slot;
 484         unsigned int offset_bits;
 485         unsigned int i;
 486         size_t match_offset;
 487         size_t match_length;
 488         int verbatim_bits;
 489         int aligned_bits;
 490         int main;
 491         int length;
 492         uint8_t *copy;
 493
 494         /* Get main symelse*/
 495         main = lzx_decode ( lzx, &lzx->main );
 496         if ( main < 0 )
 497                 return main;
 498
 499         /* Check for literals */
 500         if ( main < LZX_MAIN_LIT_CODES ) {
 501                 if ( lzx->output.data )
 502                         lzx->output.data[lzx->output.offset] = main;
 503                 lzx->output.offset++;
 504                 return 0;
 505         }
 506         main -= LZX_MAIN_LIT_CODES;
 507
 508         /* Calculate the match length */
 509         length_header = ( main & 7 );
 510         if ( length_header == 7 ) {
 511                 length = lzx_decode ( lzx, &lzx->length );
 512                 if ( length < 0 )
 513                         return length;
 514         } else {
 515                 length = 0;
 516         }
 517         match_length = ( length_header + 2 + length );
 518
 519         /* Calculate the position slot */
 520         position_slot = ( main >> 3 );
 521         if ( position_slot < LZX_REPEATED_OFFSETS ) {
 522
 523                 /* Repeated offset */
 524                 match_offset = lzx->repeated_offset[position_slot];
 525                 lzx->repeated_offset[position_slot] = lzx->repeated_offset[0];
 526                 lzx->repeated_offset[0] = match_offset;
 527
 528         } else {
 529
 530                 /* Non-repeated offset */
 531                 offset_bits = lzx_footer_bits ( position_slot );
 532                 if ( ( lzx->block_type == LZX_BLOCK_ALIGNOFFSET ) &&
 533                      ( offset_bits >= 3 ) ) {
 534                         verbatim_bits = lzx_getbits ( lzx, ( offset_bits - 3 ));
 535                         if ( verbatim_bits < 0 )
 536                                 return verbatim_bits;
 537                         verbatim_bits <<= 3;
 538                         aligned_bits = lzx_decode ( lzx, &lzx->alignoffset );
 539                         if ( aligned_bits < 0 )
 540                                 return aligned_bits;
 541                 } else {
 542                         verbatim_bits = lzx_getbits ( lzx, offset_bits );
 543                         if ( verbatim_bits < 0 )
 544                                 return verbatim_bits;
 545                         aligned_bits = 0;
 546                 }
 547                 match_offset = ( lzx_position_base[position_slot] +
 548                                  verbatim_bits + aligned_bits - 2 );
 549
 550                 /* Update repeated offset list */
 551                 for ( i = ( LZX_REPEATED_OFFSETS - 1 ) ; i > 0 ; i-- )
 552                         lzx->repeated_offset[i] = lzx->repeated_offset[ i - 1 ];
 553                 lzx->repeated_offset[0] = match_offset;
 554         }
 555
 556         /* Copy data */
 557         if ( match_offset > lzx->output.offset ) {
 558                 DBG ( "LZX match underrun out %#zx offset %#zx len %#zx\n",
 559                       lzx->output.offset, match_offset, match_length );
 560                 return -1;
 561         }
 562         if ( lzx->output.data ) {
 563                 copy = &lzx->output.data[lzx->output.offset];
 564                 for ( i = 0 ; i < match_length ; i++ )
 565                         copy[i] = copy[ i - match_offset ];
 566         }
 567         lzx->output.offset += match_length;
 568
 569         return 0;
 570 }
 571
 572 /**
 573  * Translate E8 jump addresses
 574  *
 575  * @v lzx               Decompressor
 576  */
 577 static void lzx_translate_jumps ( struct lzx *lzx ) {
 578         size_t offset;
 579         int32_t *target;
 580
 581         /* Sanity check */
 582         if ( lzx->output.offset < 10 )
 583                 return;
 584
 585         /* Scan for jump instructions */
 586         for ( offset = 0 ; offset < ( lzx->output.offset - 10 ) ; offset++ ) {
 587
 588                 /* Check for jump instruction */
 589                 if ( lzx->output.data[offset] != 0xe8 )
 590                         continue;
 591
 592                 /* Translate jump target */
 593                 target = ( ( int32_t * ) &lzx->output.data[ offset + 1 ] );
 594                 if ( *target >= 0 ) {
 595                         if ( *target < LZX_WIM_MAGIC_FILESIZE )
 596                                 *target -= offset;
 597                 } else {
 598                         if ( *target >= -( ( int32_t ) offset ) )
 599                                 *target += LZX_WIM_MAGIC_FILESIZE;
 600                 }
 601                 offset += sizeof ( *target );
 602         }
 603 }
 604
 605 /**
 606  * Decompress LZX-compressed data
 607  *
 608  * @v data              Compressed data
 609  * @v len               Length of compressed data
 610  * @v buf               Decompression buffer, or NULL
 611  * @ret out_len         Length of decompressed data, or negative error
 612  */
 613 ssize_t lzx_decompress ( const void *data, size_t len, void *buf ) {
 614         struct lzx lzx;
 615         unsigned int i;
 616         int rc;
 617
 618         /* Sanity check */
 619         if ( len % 2 ) {
 620                 DBG ( "LZX cannot handle odd-length input data\n" );
 621                 return -1;
 622         }
 623
 624         /* Initialise global state, if required */
 625         if ( ! lzx_position_base[ LZX_POSITION_SLOTS - 1 ] ) {
 626                 for ( i = 1 ; i < LZX_POSITION_SLOTS ; i++ ) {
 627                         lzx_position_base[i] =
 628                                 ( lzx_position_base[i-1] +
 629                                   ( 1 << lzx_footer_bits ( i - 1 ) ) );
 630                 }
 631         }
 632
 633         /* Initialise decompressor */
 634         memset ( &lzx, 0, sizeof ( lzx ) );
 635         lzx.input.data = data;
 636         lzx.input.len = len;
 637         lzx.output.data = buf;
 638         for ( i = 0 ; i < LZX_REPEATED_OFFSETS ; i++ )
 639                 lzx.repeated_offset[i] = 1;
 640
 641         /* Process blocks */
 642         while ( lzx.input.offset < lzx.input.len ) {
 643
 644                 /* Process block header */
 645                 if ( ( rc = lzx_block_header ( &lzx ) ) != 0 )
 646                         return rc;
 647
 648                 /* Process block contents */
 649                 if ( lzx.block_type == LZX_BLOCK_UNCOMPRESSED ) {
 650
 651                         /* Copy uncompressed data */
 652                         if ( ( rc = lzx_uncompressed ( &lzx ) ) != 0 )
 653                                 return rc;
 654
 655                 } else {
 656
 657                         /* Process token stream */
 658                         while ( lzx.output.offset < lzx.output.threshold ) {
 659                                 if ( ( rc = lzx_token ( &lzx ) ) != 0 )
 660                                         return rc;
 661                         }
 662                 }
 663         }
 664
 665         /* Postprocess to undo E8 jump compression */
 666         if ( lzx.output.data )
 667                 lzx_translate_jumps ( &lzx );
 668
 669         return lzx.output.offset;
 670 }