wimboot/wimboot-2.7.3/src/wimpatch.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  * WIM dynamic patching
  24  *
  25  */
  26
  27 #include <stddef.h>
  28 #include <stdio.h>
  29 #include <string.h>
  30 #include <strings.h>
  31 #include <assert.h>
  32 #include "wimboot.h"
  33 #include "cmdline.h"
  34 #include "vdisk.h"
  35 #include "sha1.h"
  36 #include "wim.h"
  37 #include "wimpatch.h"
  38
  39 /** Directory into which files are injected */
  40 #define WIM_INJECT_DIR "\\Windows\\System32"
  41
  42 struct wim_patch;
  43
  44 /** A region of a patched WIM file */
  45 struct wim_patch_region {
  46         /** Name */
  47         const char *name;
  48         /** Opaque token */
  49         void *opaque;
  50         /** Starting offset of region */
  51         size_t offset;
  52         /** Length of region */
  53         size_t len;
  54         /** Patch region
  55          *
  56          * @v patch             WIM patch
  57          * @v region            Patch region
  58          * @v data              Data buffer
  59          * @v offset            Relative offset
  60          * @v len               Length
  61          * @ret rc              Return status code
  62          */
  63         int ( * patch ) ( struct wim_patch *patch,
  64                           struct wim_patch_region *region,
  65                           void *data, size_t offset, size_t len );
  66 };
  67
  68 /** Regions of a patched WIM directory containing injected files */
  69 struct wim_patch_dir_regions {
  70         /** Subdirectory offset within parent entry */
  71         struct wim_patch_region subdir;
  72         /** Copy of original directory entries */
  73         struct wim_patch_region copy;
  74         /** Injected file directory entries */
  75         struct wim_patch_region file[VDISK_MAX_FILES];
  76 } __attribute__ (( packed ));
  77
  78 /** Regions of a patched WIM file */
  79 union wim_patch_regions {
  80         /** Structured list of regions */
  81         struct {
  82                 /** WIM header */
  83                 struct wim_patch_region header;
  84                 /** Injected file contents */
  85                 struct wim_patch_region file[VDISK_MAX_FILES];
  86                 /** Injected lookup table */
  87                 struct {
  88                         /** Uncompressed copy of original lookup table */
  89                         struct wim_patch_region copy;
  90                         /** Injected boot image metadata lookup table entry */
  91                         struct wim_patch_region boot;
  92                         /** Injected file lookup table entries */
  93                         struct wim_patch_region file[VDISK_MAX_FILES];
  94                 } __attribute__ (( packed )) lookup;
  95                 /** Injected boot image metadata */
  96                 struct {
  97                         /** Uncompressed copy of original metadata */
  98                         struct wim_patch_region copy;
  99                         /** Patched directory containing injected files */
 100                         struct wim_patch_dir_regions dir;
 101                 } __attribute__ (( packed )) boot;
 102         } __attribute__ (( packed ));
 103         /** Unstructured list of regions */
 104         struct wim_patch_region region[0];
 105 };
 106
 107 /** An injected directory entry */
 108 struct wim_patch_dir_entry {
 109         /** Directory entry */
 110         struct wim_directory_entry dir;
 111         /** Name */
 112         wchar_t name[ VDISK_NAME_LEN + 1 /* wNUL */ ];
 113 } __attribute__ (( packed ));
 114
 115 /** A directory containing injected files */
 116 struct wim_patch_dir {
 117         /** Name */
 118         const char *name;
 119         /** Offset to parent directory entry */
 120         size_t parent;
 121         /** Offset to original directory entries */
 122         size_t offset;
 123         /** Length of original directory entries (excluding terminator) */
 124         size_t len;
 125         /** Offset to modified directory entries */
 126         size_t subdir;
 127 };
 128
 129 /** A patched WIM file */
 130 struct wim_patch {
 131         /** Virtual file */
 132         struct vdisk_file *file;
 133         /** Patched WIM header */
 134         struct wim_header header;
 135         /** Original lookup table */
 136         struct wim_resource_header lookup;
 137         /** Original boot image metadata */
 138         struct wim_resource_header boot;
 139         /** Original boot index */
 140         uint32_t boot_index;
 141         /** Directory containing injected files */
 142         struct wim_patch_dir dir;
 143         /** Patched regions */
 144         union wim_patch_regions regions;
 145 };
 146
 147 /**
 148  * Align WIM offset to nearest qword
 149  *
 150  * @v len               Length
 151  * @ret len             Aligned length
 152  */
 153 static size_t wim_align ( size_t len ) {
 154         return ( ( len + 0x07 ) & ~0x07 );
 155 }
 156
 157 /**
 158  * Calculate WIM hash
 159  *
 160  * @v vfile             Virtual file
 161  * @v hash              Hash to fill in
 162  */
 163 static void wim_hash ( struct vdisk_file *vfile, struct wim_hash *hash ) {
 164         uint8_t ctx[SHA1_CTX_SIZE];
 165         uint8_t buf[512];
 166         size_t offset;
 167         size_t len;
 168
 169         /* Calculate SHA-1 digest */
 170         sha1_init ( ctx );
 171         for ( offset = 0 ; offset < vfile->len ; offset += len ) {
 172
 173                 /* Read block */
 174                 len = ( vfile->len - offset );
 175                 if ( len > sizeof ( buf ) )
 176                         len = sizeof ( buf );
 177                 vfile->read ( vfile, buf, offset, len );
 178
 179                 /* Update digest */
 180                 sha1_update ( ctx, buf, len );
 181         }
 182         sha1_final ( ctx, hash->sha1 );
 183 }
 184
 185 /**
 186  * Determine whether or not to inject file
 187  *
 188  * @v vfile             Virtual file
 189  * @ret inject          Inject this file
 190  */
 191 static int wim_inject_file ( struct vdisk_file *vfile ) {
 192         size_t name_len;
 193         const char *ext;
 194
 195         /* Ignore non-existent files */
 196         if ( ! vfile->read )
 197                 return 0;
 198
 199         /* Ignore wimboot itself */
 200         if ( strcasecmp ( vfile->name, "wimboot" ) == 0 )
 201                 return 0;
 202
 203         /* Ignore bootmgr files */
 204         if ( strcasecmp ( vfile->name, "bootmgr" ) == 0 )
 205                 return 0;
 206         if ( strcasecmp ( vfile->name, "bootmgr.exe" ) == 0 )
 207                 return 0;
 208
 209         /* Ignore BCD files */
 210         if ( strcasecmp ( vfile->name, "BCD" ) == 0 )
 211                 return 0;
 212
 213         /* Locate file extension */
 214         name_len = strlen ( vfile->name );
 215         ext = ( ( name_len > 4 ) ? ( vfile->name + name_len - 4 ) : "" );
 216
 217         /* Ignore .wim files */
 218         if ( strcasecmp ( ext, ".wim" ) == 0 )
 219                 return 0;
 220
 221         /* Ignore .sdi files */
 222         if ( strcasecmp ( ext, ".sdi" ) == 0 )
 223                 return 0;
 224
 225         /* Ignore .efi files */
 226         if ( strcasecmp ( ext, ".efi" ) == 0 )
 227                 return 0;
 228
 229         /* Ignore .ttf files */
 230         if ( strcasecmp ( ext, ".ttf" ) == 0 )
 231                 return 0;
 232
 233         return 1;
 234 }
 235
 236 /**
 237  * Patch WIM header
 238  *
 239  * @v patch             WIM patch
 240  * @v region            Patch region
 241  * @v data              Data buffer
 242  * @v offset            Relative offset
 243  * @v len               Length
 244  * @ret rc              Return status code
 245  */
 246 static int wim_patch_header ( struct wim_patch *patch,
 247                               struct wim_patch_region *region,
 248                               void *data, size_t offset, size_t len ) {
 249         struct wim_header *header = &patch->header;
 250
 251         /* Sanity checks */
 252         assert ( offset < sizeof ( *header ) );
 253         assert ( len <= ( sizeof ( *header ) - offset ) );
 254
 255         /* Copy patched header */
 256         if ( patch->lookup.offset != patch->header.lookup.offset ) {
 257                 DBG2 ( "...patched WIM %s lookup table %#llx->%#llx\n",
 258                        region->name, patch->lookup.offset,
 259                        patch->header.lookup.offset );
 260         }
 261         if ( patch->boot.offset != patch->header.boot.offset ) {
 262                 DBG2 ( "...patched WIM %s boot metadata %#llx->%#llx\n",
 263                        region->name, patch->boot.offset,
 264                        patch->header.boot.offset );
 265         }
 266         if ( patch->boot_index != patch->header.boot_index ) {
 267                 DBG2 ( "...patched WIM %s boot index %d->%d\n", region->name,
 268                        patch->boot_index, patch->header.boot_index );
 269         }
 270         memcpy ( data, ( ( ( void * ) &patch->header ) + offset ), len );
 271
 272         return 0;
 273 }
 274
 275 /**
 276  * Patch injected file content
 277  *
 278  * @v patch             WIM patch
 279  * @v region            Patch region
 280  * @v data              Data buffer
 281  * @v offset            Relative offset
 282  * @v len               Length
 283  * @ret rc              Return status code
 284  */
 285 static int wim_patch_file ( struct wim_patch *patch __unused,
 286                             struct wim_patch_region *region,
 287                             void *data, size_t offset, size_t len ) {
 288         struct vdisk_file *vfile = region->opaque;
 289
 290         /* Read from file */
 291         vfile->read ( vfile, data, offset, len );
 292
 293         return 0;
 294 }
 295
 296 /**
 297  * Patch uncompressed copy of original lookup table
 298  *
 299  * @v patch             WIM patch
 300  * @v region            Patch region
 301  * @v data              Data buffer
 302  * @v offset            Relative offset
 303  * @v len               Length
 304  * @ret rc              Return status code
 305  */
 306 static int wim_patch_lookup_copy ( struct wim_patch *patch,
 307                                    struct wim_patch_region *region __unused,
 308                                    void *data, size_t offset, size_t len ) {
 309         int rc;
 310
 311         /* Read original lookup table */
 312         if ( ( rc = wim_read ( patch->file, &patch->header, &patch->lookup,
 313                                data, offset, len ) ) != 0 )
 314                 return rc;
 315
 316         return 0;
 317 }
 318
 319 /**
 320  * Patch injected boot image metadata lookup table entry
 321  *
 322  * @v patch             WIM patch
 323  * @v region            Patch region
 324  * @v data              Data buffer
 325  * @v offset            Relative offset
 326  * @v len               Length
 327  * @ret rc              Return status code
 328  */
 329 static int wim_patch_lookup_boot ( struct wim_patch *patch,
 330                                    struct wim_patch_region *region __unused,
 331                                    void *data, size_t offset, size_t len ) {
 332         struct wim_lookup_entry entry;
 333
 334         /* Sanity checks */
 335         assert ( offset < sizeof ( entry ) );
 336         assert ( len <= ( sizeof ( entry ) - offset ) );
 337
 338         /* Construct lookup table entry */
 339         memset ( &entry, 0, sizeof ( entry ) );
 340         memcpy ( &entry.resource, &patch->header.boot,
 341                  sizeof ( entry.resource ) );
 342
 343         /* Copy lookup table entry */
 344         memcpy ( data, ( ( ( void * ) &entry ) + offset ), len );
 345
 346         return 0;
 347 }
 348
 349 /**
 350  * Patch injected file lookup table entry
 351  *
 352  * @v patch             WIM patch
 353  * @v region            Patch region
 354  * @v data              Data buffer
 355  * @v offset            Relative offset
 356  * @v len               Length
 357  * @ret rc              Return status code
 358  */
 359 static int wim_patch_lookup_file ( struct wim_patch *patch __unused,
 360                                    struct wim_patch_region *region,
 361                                    void *data, size_t offset, size_t len ) {
 362         struct wim_patch_region *rfile = region->opaque;
 363         struct vdisk_file *vfile = rfile->opaque;
 364         struct wim_lookup_entry entry;
 365
 366         /* Sanity checks */
 367         assert ( offset < sizeof ( entry ) );
 368         assert ( len <= ( sizeof ( entry ) - offset ) );
 369
 370         /* Construct lookup table entry */
 371         memset ( &entry, 0, sizeof ( entry ) );
 372         entry.resource.offset = rfile->offset;
 373         entry.resource.len = vfile->len;
 374         entry.resource.zlen__flags = entry.resource.len;
 375         entry.refcnt = 1;
 376         wim_hash ( vfile, &entry.hash );
 377
 378         /* Copy lookup table entry */
 379         memcpy ( data, ( ( ( void * ) &entry ) + offset ), len );
 380         DBG2 ( "...patched WIM %s %s\n", region->name, vfile->name );
 381
 382         return 0;
 383 }
 384
 385 /**
 386  * Patch uncompressed copy of original boot metadata
 387  *
 388  * @v patch             WIM patch
 389  * @v region            Patch region
 390  * @v data              Data buffer
 391  * @v offset            Relative offset
 392  * @v len               Length
 393  * @ret rc              Return status code
 394  */
 395 static int wim_patch_boot_copy ( struct wim_patch *patch,
 396                                  struct wim_patch_region *region __unused,
 397                                  void *data, size_t offset, size_t len ) {
 398         int rc;
 399
 400         /* Read original boot metadata */
 401         if ( ( rc = wim_read ( patch->file, &patch->header, &patch->boot,
 402                                data, offset, len ) ) != 0 )
 403                 return rc;
 404
 405         return 0;
 406 }
 407
 408 /**
 409  * Patch subdirectory offset within parent directory entry
 410  *
 411  * @v patch             WIM patch
 412  * @v region            Patch region
 413  * @v data              Data buffer
 414  * @v offset            Relative offset
 415  * @v len               Length
 416  * @ret rc              Return status code
 417  */
 418 static int wim_patch_dir_subdir ( struct wim_patch *patch,
 419                                   struct wim_patch_region *region,
 420                                   void *data, size_t offset, size_t len ) {
 421         struct wim_patch_dir *dir = region->opaque;
 422         uint64_t subdir = dir->subdir;
 423
 424         /* Sanity checks */
 425         assert ( offset < sizeof ( subdir ) );
 426         assert ( len <= ( sizeof ( subdir ) - offset ) );
 427
 428         /* Copy subdirectory offset */
 429         memcpy ( data, ( ( ( void * ) &subdir ) + offset ), len );
 430         DBG2 ( "...patched WIM %s %s %#llx\n", region->name, dir->name,
 431                ( patch->header.boot.offset + subdir ) );
 432
 433         return 0;
 434 }
 435
 436 /**
 437  * Patch copy of original directory entries
 438  *
 439  * @v patch             WIM patch
 440  * @v region            Patch region
 441  * @v data              Data buffer
 442  * @v offset            Relative offset
 443  * @v len               Length
 444  * @ret rc              Return status code
 445  */
 446 static int wim_patch_dir_copy ( struct wim_patch *patch,
 447                                 struct wim_patch_region *region,
 448                                 void *data, size_t offset, size_t len ) {
 449         struct wim_patch_dir *dir = region->opaque;
 450         int rc;
 451
 452         /* Read portion of original boot metadata */
 453         if ( ( rc = wim_read ( patch->file, &patch->header, &patch->boot,
 454                                data, ( dir->offset + offset ), len ) ) != 0 )
 455                 return rc;
 456
 457         return 0;
 458 }
 459
 460 /**
 461  * Patch injected directory entries
 462  *
 463  * @v patch             WIM patch
 464  * @v region            Patch region
 465  * @v data              Data buffer
 466  * @v offset            Relative offset
 467  * @v len               Length
 468  * @ret rc              Return status code
 469  */
 470 static int wim_patch_dir_file ( struct wim_patch *patch __unused,
 471                                 struct wim_patch_region *region,
 472                                 void *data, size_t offset, size_t len ) {
 473         struct wim_patch_region *rfile = region->opaque;
 474         struct vdisk_file *vfile = rfile->opaque;
 475         struct wim_patch_dir_entry entry;
 476         size_t name_len = strlen ( vfile->name );
 477         unsigned int i;
 478
 479         /* Sanity checks */
 480         assert ( offset < sizeof ( entry ) );
 481         assert ( len <= ( sizeof ( entry ) - offset ) );
 482
 483         /* Construct directory entry */
 484         memset ( &entry, 0, sizeof ( entry ) );
 485         entry.dir.len = wim_align ( sizeof ( entry ) );
 486         entry.dir.attributes = WIM_ATTR_NORMAL;
 487         entry.dir.security = WIM_NO_SECURITY;
 488         entry.dir.created = WIM_MAGIC_TIME;
 489         entry.dir.accessed = WIM_MAGIC_TIME;
 490         entry.dir.written = WIM_MAGIC_TIME;
 491         wim_hash ( vfile, &entry.dir.hash );
 492         entry.dir.name_len = ( name_len * sizeof ( entry.name[0] ) );
 493         for ( i = 0 ; i < name_len ; i++ )
 494                 entry.name[i] = vfile->name[i];
 495
 496         /* Copy directory entry */
 497         memcpy ( data, ( ( ( void * ) &entry ) + offset ), len );
 498         DBG2 ( "...patched WIM %s %s\n", region->name, vfile->name );
 499
 500         return 0;
 501 }
 502
 503 /**
 504  * Patch WIM region
 505  *
 506  * @v patch             WIM patch
 507  * @v region            Patch region
 508  * @v data              Data buffer
 509  * @v offset            Relative offset
 510  * @v len               Length
 511  * @ret rc              Return status code
 512  */
 513 static int wim_patch_region ( struct wim_patch *patch,
 514                               struct wim_patch_region *region,
 515                               void *data, size_t offset, size_t len ) {
 516         size_t skip;
 517         int rc;
 518
 519         /* Skip unused regions */
 520         if ( ! region->patch )
 521                 return 0;
 522
 523         /* Skip any data before this region */
 524         skip = ( ( region->offset > offset ) ?
 525                  ( region->offset - offset ) : 0 );
 526         if ( skip >= len )
 527                 return 0;
 528         data += skip;
 529         offset += skip;
 530         len -= skip;
 531
 532         /* Convert to relative offset within this region */
 533         offset -= region->offset;
 534
 535         /* Skip any data after this region */
 536         if ( offset >= region->len )
 537                 return 0;
 538         if ( len > ( region->len - offset ) )
 539                 len = ( region->len - offset );
 540
 541         /* Patch this region */
 542         if ( ( rc = region->patch ( patch, region, data, offset, len ) ) != 0 )
 543                 return rc;
 544         DBG2 ( "...patched WIM %s at [%#zx,%#zx)\n", region->name,
 545                ( region->offset + offset ), ( region->offset + offset + len ) );
 546
 547         return 0;
 548 }
 549
 550 /**
 551  * Construct patched WIM region
 552  *
 553  * @v region            Patched region to fill in
 554  * @v name              Name
 555  * @v opaque            Opaque data
 556  * @v offset            Offset
 557  * @v len               Length
 558  * @v patch             Patch method
 559  * @ret offset          Next offset
 560  */
 561 static inline __attribute__ (( always_inline )) size_t
 562 wim_construct_region ( struct wim_patch_region *region, const char *name,
 563                        void *opaque, size_t offset, size_t len,
 564                        int ( * patch ) ( struct wim_patch *patch,
 565                                          struct wim_patch_region *region,
 566                                          void *data, size_t offset,
 567                                          size_t len ) ) {
 568
 569         DBG ( "...patching WIM %s at [%#zx,%#zx)\n",
 570               name, offset, ( offset + len ) );
 571         region->name = name;
 572         region->opaque = opaque;
 573         region->offset = offset;
 574         region->len = len;
 575         region->patch = patch;
 576         return ( offset + len );
 577 }
 578
 579 /**
 580  * Construct patch WIM directory regions
 581  *
 582  * @v patch             WIM patch
 583  * @v dir               Patched directory
 584  * @v offset            Offset
 585  * @v regions           Patched directory regions to fill in
 586  * @ret offset          Next offset
 587  */
 588 static size_t wim_construct_dir ( struct wim_patch *patch,
 589                                   struct wim_patch_dir *dir, size_t offset,
 590                                   struct wim_patch_dir_regions *regions ) {
 591         struct wim_patch_dir_entry *entry;
 592         struct wim_patch_region *rfile;
 593         size_t boot_offset = patch->header.boot.offset;
 594         unsigned int i;
 595
 596         DBG ( "...patching WIM directory at %#zx from [%#zx,%#zx)\n",
 597               ( boot_offset + dir->parent ), ( boot_offset + dir->offset ),
 598               ( boot_offset + dir->offset + dir->len ) );
 599
 600         /* Align directory entries */
 601         offset = wim_align ( offset );
 602         dir->subdir = ( offset - patch->header.boot.offset );
 603
 604         /* Construct injected file directory entries */
 605         for ( i = 0 ; i < VDISK_MAX_FILES ; i++ ) {
 606                 rfile = &patch->regions.file[i];
 607                 if ( ! rfile->patch )
 608                         continue;
 609                 offset = wim_construct_region ( &regions->file[i], "dir.file",
 610                                                 rfile, offset,
 611                                                 sizeof ( *entry ),
 612                                                 wim_patch_dir_file );
 613                 offset = wim_align ( offset );
 614         }
 615
 616         /* Construct copy of original directory entries */
 617         offset = wim_construct_region ( &regions->copy, dir->name, dir, offset,
 618                                         dir->len, wim_patch_dir_copy );
 619
 620         /* Allow space for directory terminator */
 621         offset += sizeof ( entry->dir.len );
 622
 623         /* Construct subdirectory offset within parent directory entry */
 624         wim_construct_region ( &regions->subdir, "dir.subdir", dir,
 625                                ( boot_offset + dir->parent +
 626                                  offsetof ( typeof ( entry->dir ), subdir ) ),
 627                                sizeof ( entry->dir.subdir ),
 628                                wim_patch_dir_subdir );
 629
 630         return offset;
 631 }
 632
 633 /**
 634  * Construct WIM patch
 635  *
 636  * @v file              Virtual file
 637  * @v boot_index        New boot index (or zero)
 638  * @v inject            Inject files into WIM
 639  * @v patch             Patch to fill in
 640  * @ret rc              Return status code
 641  */
 642 static int wim_construct_patch ( struct vdisk_file *file,
 643                                  unsigned int boot_index, int inject,
 644                                  struct wim_patch *patch ) {
 645         union wim_patch_regions *regions = &patch->regions;
 646         struct wim_patch_region *rfile;
 647         struct wim_resource_header *lookup;
 648         struct wim_resource_header *boot;
 649         struct wim_directory_entry direntry;
 650         struct wim_lookup_entry *entry;
 651         struct vdisk_file *vfile;
 652         size_t offset;
 653         unsigned int injected = 0;
 654         unsigned int i;
 655         int rc;
 656
 657         /* Initialise patch */
 658         memset ( patch, 0, sizeof ( *patch ) );
 659         patch->file = file;
 660         DBG ( "...patching WIM %s\n", file->name );
 661
 662         /* Reset file length */
 663         file->xlen = file->len;
 664         offset = file->len;
 665
 666         /* Read WIM header */
 667         if ( ( rc = wim_header ( file, &patch->header ) ) != 0 )
 668                 return rc;
 669         lookup = &patch->header.lookup;
 670         boot = &patch->header.boot;
 671
 672         /* Patch header within original image body */
 673         wim_construct_region ( &regions->header, "header", NULL, 0,
 674                                sizeof ( patch->header ), wim_patch_header );
 675
 676         /* Record original lookup table */
 677         memcpy ( &patch->lookup, lookup, sizeof ( patch->lookup ) );
 678
 679         /* Record original metadata for selected boot image (which may
 680          * not be the originally selected boot image).
 681          */
 682         if ( ( rc = wim_metadata ( file, &patch->header, boot_index,
 683                                    &patch->boot ) ) != 0 )
 684                 return rc;
 685
 686         /* Record original boot index */
 687         patch->boot_index = patch->header.boot_index;
 688
 689         /* Update boot index in patched header, if applicable */
 690         if ( boot_index )
 691                 patch->header.boot_index = boot_index;
 692
 693         /* Do nothing more if injection is disabled */
 694         if ( ! inject )
 695                 return 0;
 696
 697         /* Construct injected files */
 698         for ( i = 0 ; i < VDISK_MAX_FILES ; i++ ) {
 699                 vfile = &vdisk_files[i];
 700                 if ( ! wim_inject_file ( vfile ) )
 701                         continue;
 702                 offset = wim_construct_region ( &regions->file[i], vfile->name,
 703                                                 vfile, offset, vfile->len,
 704                                                 wim_patch_file );
 705                 injected++;
 706         }
 707
 708         /* Do nothing more if no files are injected */
 709         if ( injected == 0 )
 710                 return 0;
 711
 712         /* Calculate boot index for injected image */
 713         if ( ( rc = wim_count ( file, &patch->header, &boot_index ) ) != 0 )
 714                 return rc;
 715         patch->header.images = ( boot_index + 1 );
 716         patch->header.boot_index = patch->header.images;
 717
 718         /* Construct injected lookup table */
 719         lookup->offset = offset = wim_align ( offset );
 720         offset = wim_construct_region ( &regions->lookup.copy, "lookup.copy",
 721                                         NULL, offset, patch->lookup.len,
 722                                         wim_patch_lookup_copy );
 723         offset = wim_construct_region ( &regions->lookup.boot, "lookup.boot",
 724                                         NULL, offset, sizeof ( *entry ),
 725                                         wim_patch_lookup_boot );
 726         for ( i = 0 ; i < VDISK_MAX_FILES ; i++ ) {
 727                 rfile = &regions->file[i];
 728                 if ( ! rfile->patch )
 729                         continue;
 730                 offset = wim_construct_region ( &regions->lookup.file[i],
 731                                                 "lookup.file", rfile,
 732                                                 offset, sizeof ( *entry ),
 733                                                 wim_patch_lookup_file );
 734         }
 735         lookup->offset = regions->lookup.copy.offset;
 736         lookup->len = ( offset - lookup->offset );
 737         lookup->zlen__flags = lookup->len;
 738
 739         /* Locate directory containing injected files */
 740         patch->dir.name = WIM_INJECT_DIR;
 741         if ( ( rc = wim_path ( file, &patch->header, &patch->boot,
 742                                L(WIM_INJECT_DIR), &patch->dir.parent,
 743                                &direntry ) ) != 0 )
 744                 return rc;
 745         patch->dir.offset = direntry.subdir;
 746         if ( ( rc = wim_dir_len ( file, &patch->header, &patch->boot,
 747                                   patch->dir.offset, &patch->dir.len ) ) != 0 )
 748                 return rc;
 749
 750         /* Construct injected boot image metadata */
 751         boot->offset = offset = wim_align ( offset );
 752         offset = wim_construct_region ( &regions->boot.copy, "boot.copy",
 753                                         NULL, offset, patch->boot.len,
 754                                         wim_patch_boot_copy );
 755         offset = wim_construct_dir ( patch, &patch->dir, offset,
 756                                      &regions->boot.dir );
 757         boot->len = ( offset - boot->offset );
 758         boot->zlen__flags = ( boot->len | WIM_RESHDR_METADATA );
 759
 760         /* Record patched length */
 761         file->xlen = offset;
 762         DBG ( "...patching WIM length %#zx->%#zx\n", file->len, file->xlen );
 763
 764         return 0;
 765 }
 766
 767 /**
 768  * Patch WIM file
 769  *
 770  * @v file              Virtual file
 771  * @v data              Data buffer
 772  * @v offset            Offset
 773  * @v len               Length
 774  */
 775 void patch_wim ( struct vdisk_file *file, void *data, size_t offset,
 776                  size_t len ) {
 777         static struct wim_patch cached_patch;
 778         struct wim_patch *patch = &cached_patch;
 779         struct wim_patch_region *region;
 780         unsigned int boot_index;
 781         unsigned int i;
 782         int inject;
 783         int rc;
 784
 785         /* Do nothing unless patching is required */
 786         boot_index = cmdline_index;
 787         inject = ( ! cmdline_rawwim );
 788         if ( ( boot_index == 0 ) && ( ! inject ) )
 789                 return;
 790
 791         /* Update cached patch if required */
 792         if ( file != patch->file ) {
 793                 if ( ( rc = wim_construct_patch ( file, boot_index, inject,
 794                                                   patch ) ) != 0 ) {
 795                         die ( "Could not patch WIM %s\n", file->name );
 796                 }
 797         }
 798         patch = &cached_patch;
 799
 800         /* Patch regions */
 801         for ( i = 0 ; i < ( sizeof ( patch->regions ) /
 802                             sizeof ( patch->regions.region[0] ) ) ; i++ ) {
 803                 region = &patch->regions.region[i];
 804                 if ( ( rc = wim_patch_region ( patch, region, data, offset,
 805                                                len ) ) != 0 ) {
 806                         die ( "Could not patch WIM %s %s at [%#zx,%#zx)\n",
 807                               file->name, region->name, offset,
 808                               ( offset + len ) );
 809                 }
 810         }
 811 }