Coverage for /Syzygy/pe/transforms/add_imports_transform.cc

CoverageLines executed / instrumented / missingexe / inst / missLanguageGroup
78.5%1942470.C++source

Line-by-line coverage:

   1    :  // Copyright 2012 Google Inc.
   2    :  //
   3    :  // Licensed under the Apache License, Version 2.0 (the "License");
   4    :  // you may not use this file except in compliance with the License.
   5    :  // You may obtain a copy of the License at
   6    :  //
   7    :  //     http://www.apache.org/licenses/LICENSE-2.0
   8    :  //
   9    :  // Unless required by applicable law or agreed to in writing, software
  10    :  // distributed under the License is distributed on an "AS IS" BASIS,
  11    :  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12    :  // See the License for the specific language governing permissions and
  13    :  // limitations under the License.
  14    :  //
  15    :  // The AddImportsTransform can be summed up as follows:
  16    :  //
  17    :  // (1) Make sure that the imports and IAT data directories exist.
  18    :  // (2) For each module to be imported, either find it in the import data
  19    :  //     directory, or add a new entry. The entry is always added to the end
  20    :  //     of the list so that module indices are strictly increasing, allowing
  21    :  //     the transform to be stacked. Adding a new entry also causes the creation
  22    :  //     of two new blocks (for the INT and the module filename), as well as
  23    :  //     extends the existing IAT block.
  24    :  // (3) For each symbol to be imported, either find it in the module's INT/IAT,
  25    :  //     or add a new entry. Adding a new entry causes the existing INT and IAT
  26    :  //     blocks to be extended. The new entry is always added to the end of the
  27    :  //     module's table so that symbol indices are strictly increasing, again
  28    :  //     allowing the transform to be stacked. Rather than allocating a new
  29    :  //     block for the name of the symbol we reuse the module filename block and
  30    :  //     insert the name of the symbol immediately prior to the module filename.
  31    :  //     This ensures that all of the strings for a module are laid out together,
  32    :  //     mimicking the observed behavior of the MS linker.
  33    :  //
  34    :  // We give a quick rundown of the PE structures involved, their layout in
  35    :  // typical PE images and how we parse them into blocks. This helps visualize
  36    :  // the work performed by the transform.
  37    :  //
  38    :  // headers:
  39    :  //
  40    :  //   ...
  41    :  //   nt_headers
  42    :  //     DataDirectory
  43    :  //       ...
  44    :  //       IMAGE_DIRECTORY_ENTRY_IMPORT -> IMAGE_IMPORT_DESCRIPTOR array
  45    :  //       ...
  46    :  //       IMAGE_DIRECTORY_ENTRY_IAT -> Import Address Table
  47    :  //       ...
  48    :  //
  49    :  // .rdata:
  50    :  //
  51    :  //   Import Address Table
  52    :  //   NOTE: All entries in this table must remain consecutive as it is also
  53    :  //       exposed directly via a data directory. At runtime these get patched to
  54    :  //       point to the actual functions rather than the thunks. This is stored
  55    :  //       at the very beginning of .rdata and parsed as a single Block.
  56    :  //     IAT[0,0] -> thunk[0, 0]  \
  57    :  //     ...                      |
  58    :  //     IAT[0,j] -> thunk[0, j]  |
  59    :  //     NULL terminator          |
  60    :  //     ...                      |- Block
  61    :  //     IAT[i,0] -> thunk[i, 0]  |
  62    :  //     ...                      |
  63    :  //     IAT[i,k] -> thunk[i, k]  |
  64    :  //     NULL terminator          /
  65    :  //
  66    :  //   ... whole bunch of other .rdata here ...
  67    :  //   NOTE: The following are stored at the end of .rdata, in the order
  68    :  //       shown (they are not quite last, being immediately prior to export
  69    :  //       information).
  70    :  //
  71    :  //   IMAGE_IMPORT_DESCRIPTOR array  \
  72    :  //     IMAGE_IMPORT_DESCRIPTOR[0]   |
  73    :  //       -> module_name[0]          |
  74    :  //       -> INT[0,0]                |
  75    :  //       -> IAT[0,0]                |
  76    :  //     ...                          |- Block
  77    :  //     IMAGE_IMPORT_DESCRIPTOR[i]   |
  78    :  //       -> module_name[i]          |
  79    :  //       -> INT[i,0]                |
  80    :  //       -> IAT[i,0]                |
  81    :  //     NULL terminator              /
  82    :  //
  83    :  //   Import Name Table (also known as Hint Name Array)
  84    :  //   NOTE: The entries for each module need be consecutive. While the entries
  85    :  //       across all modules are consecutive, they need not be.
  86    :  //     INT[0,0] -> thunk[0, 0]  \
  87    :  //     ...                      |_ Block
  88    :  //     INT[0,j] -> thunk[0, j]  |
  89    :  //     NULL terminator          /
  90    :  //     ...
  91    :  //     INT[i,0] -> thunk[i, 0]  \
  92    :  //     ...                      |_ Block
  93    :  //     INT[i,k] -> thunk[i, k]  |
  94    :  //     NULL terminator          /
  95    :  //
  96    :  //   Array of names
  97    :  //   NOTE: These are consecutive in typical PE images (with the layout shown
  98    :  //       below), but they need not be.
  99    :  //     thunk[0, 0]     } Block
 100    :  //     ...
 101    :  //     thunk[0, j]     } Block
 102    :  //     module_name[0]  } Block
 103    :  //     ...
 104    :  //     thunk[i, 0]     } Block
 105    :  //     ...
 106    :  //     thunk[i, k]     } Block
 107    :  //     module_name[i]  } Block
 108    :  
 109    :  #include "syzygy/pe/transforms/add_imports_transform.h"
 110    :  
 111    :  #include "base/string_piece.h"
 112    :  #include "base/string_util.h"
 113    :  #include "base/stringprintf.h"
 114    :  #include "syzygy/block_graph/typed_block.h"
 115    :  #include "syzygy/common/align.h"
 116    :  #include "syzygy/pe/pe_utils.h"
 117    :  
 118    :  namespace pe {
 119    :  namespace transforms {
 120    :  
 121    :  using block_graph::BlockGraph;
 122    :  using block_graph::ConstTypedBlock;
 123    :  using block_graph::TypedBlock;
 124    :  using core::RelativeAddress;
 125    :  
 126    :  // A simple struct that can be used to let us access strings using TypedBlock.
 127    :  struct StringStruct {
 128    :    const char string[1];
 129    :  };
 130    :  
 131    :  typedef BlockGraph::Offset Offset;
 132    :  typedef TypedBlock<IMAGE_DOS_HEADER> DosHeader;
 133    :  typedef TypedBlock<IMAGE_IMPORT_BY_NAME> ImageImportByName;
 134    :  typedef TypedBlock<IMAGE_IMPORT_DESCRIPTOR> ImageImportDescriptor;
 135    :  typedef TypedBlock<IMAGE_NT_HEADERS> NtHeaders;
 136    :  typedef TypedBlock<IMAGE_THUNK_DATA32> ImageThunkData32;
 137    :  typedef TypedBlock<StringStruct> String;
 138    :  
 139    :  namespace {
 140    :  
 141    :  const size_t kPtrSize = sizeof(core::RelativeAddress);
 142    :  
 143    :  // Looks up the given data directory and checks that it points to valid data.
 144    :  // If it doesn't, it will allocate a block with the given name and size.
 145    :  bool EnsureDataDirectoryExists(size_t directory_index,
 146    :                                 const char* block_name,
 147    :                                 size_t block_size,
 148    :                                 BlockGraph* block_graph,
 149  E :                                 BlockGraph::Block* nt_headers_block) {
 150  E :    DCHECK(block_name != NULL);
 151    :    DCHECK_LT(directory_index,
 152  E :              static_cast<size_t>(IMAGE_NUMBEROF_DIRECTORY_ENTRIES));
 153  E :    DCHECK_GT(block_size, 0u);
 154  E :    DCHECK(block_graph != NULL);
 155  E :    DCHECK(nt_headers_block != NULL);
 156    :  
 157  E :    NtHeaders nt_headers;
 158  E :    if (!nt_headers.Init(0, nt_headers_block)) {
 159  i :      LOG(ERROR) << "Unable to cast NT headers.";
 160  i :      return false;
 161    :    }
 162    :  
 163    :    IMAGE_DATA_DIRECTORY* data_directory =
 164  E :        nt_headers->OptionalHeader.DataDirectory + directory_index;
 165    :  
 166    :    // No entry? Then make a zero initialized block that is stored in .rdata,
 167    :    // where all of these structures live.
 168  E :    if (!nt_headers.HasReference(data_directory->VirtualAddress)) {
 169    :      BlockGraph::Section* section = block_graph->FindOrAddSection(
 170  E :          kReadOnlyDataSectionName, kReadOnlyDataCharacteristics);
 171  E :      DCHECK(section != NULL);
 172    :  
 173    :      BlockGraph::Block* block = block_graph->AddBlock(
 174  E :          BlockGraph::DATA_BLOCK, block_size, block_name);
 175  E :      DCHECK(block != NULL);
 176  E :      block->set_section(section->id());
 177  E :      block->set_attribute(BlockGraph::PE_PARSED);
 178    :  
 179    :      // We need to actually allocate the data so that future TypedBlock
 180    :      // dereferences will work.
 181  E :      if (block->AllocateData(block_size) == NULL) {
 182  i :        LOG(ERROR) << "Failed to allocate block data.";
 183  i :        return false;
 184    :      }
 185    :  
 186    :      // Hook it up to the NT header.
 187    :      nt_headers.SetReference(BlockGraph::RELATIVE_REF,
 188    :                              data_directory->VirtualAddress,
 189    :                              block,
 190  E :                              0, 0);
 191  E :      data_directory->Size = block_size;
 192    :    }
 193    :  
 194  E :    return true;
 195  E :  }
 196    :  
 197    :  // Finds or creates an Image Import Descriptor block for the given library.
 198    :  // Returns true on success, false otherwise.
 199    :  bool FindOrAddImageImportDescriptor(const char* module_name,
 200    :                                      BlockGraph* block_graph,
 201    :                                      BlockGraph::Block* iida_block,
 202    :                                      BlockGraph::Block* iat_block,
 203    :                                      ImageImportDescriptor* iid,
 204  E :                                      bool* added) {
 205  E :    DCHECK(module_name != NULL);
 206  E :    DCHECK(block_graph != NULL);
 207  E :    DCHECK(iida_block != NULL);
 208  E :    DCHECK(iat_block != NULL);
 209  E :    DCHECK(iid != NULL);
 210  E :    DCHECK(added != NULL);
 211    :  
 212  E :    *added = false;
 213    :  
 214  E :    ImageImportDescriptor iida;
 215  E :    if (!iida.Init(0, iida_block)) {
 216  i :      LOG(ERROR) << "Unable to cast Image Import Descriptor.";
 217  i :      return false;
 218    :    }
 219    :  
 220    :    // The array is NULL terminated with a potentially incomplete descriptor so
 221    :    // we can't use ElementCount - 1.
 222    :    size_t descriptor_count =
 223    :        (common::AlignUp(iida_block->size(), sizeof(IMAGE_IMPORT_DESCRIPTOR)) /
 224  E :         sizeof(IMAGE_IMPORT_DESCRIPTOR)) - 1;
 225    :  
 226  E :    for (size_t iida_index = 0; iida_index < descriptor_count; ++iida_index) {
 227  E :      String dll_name;
 228  E :      if (!iida.Dereference(iida[iida_index].Name, &dll_name)) {
 229  i :        LOG(ERROR) << "Unable to dereference DLL name.";
 230  i :        return false;
 231    :      }
 232    :  
 233  E :      size_t max_len = dll_name.block()->size() - dll_name.offset();
 234  E :      if (base::strncasecmp(dll_name->string, module_name, max_len) == 0) {
 235    :        // This should never fail, but we sanity check it nonetheless.
 236  E :        bool result = iid->Init(iida.OffsetOf(iida[iida_index]), iida.block());
 237  E :        DCHECK(result);
 238  E :        return true;
 239    :      }
 240  E :    }
 241    :  
 242    :    // Create room for the new descriptor, which we'll tack on to the end of the
 243    :    // array, but before the NULL terminator. We use 'InsertData' so that all
 244    :    // labels are patched up.
 245  E :    Offset new_iid_offset = descriptor_count * sizeof(IMAGE_IMPORT_DESCRIPTOR);
 246    :    iida_block->InsertData(
 247  E :        new_iid_offset, sizeof(IMAGE_IMPORT_DESCRIPTOR), true);
 248    :    iida_block->SetLabel(
 249    :        new_iid_offset,
 250    :        base::StringPrintf("Image Import Descriptor: %s", module_name),
 251  E :        BlockGraph::DATA_LABEL);
 252    :  
 253    :    // We expect the new entry to be dereferencable using iida[descriptor_count].
 254  E :    DCHECK_GT(iida.ElementCount(), descriptor_count);
 255    :  
 256    :    // Create the various child structures that will be pointed to by the
 257    :    // import descriptor. The INT block and the IAT block are NULL terminated
 258    :    // lists of pointers, and the terminating NULL is allocated. We don't yet
 259    :    // allocate a block to hold the import names, deferring that for later.
 260  E :    BlockGraph::SectionId iida_section_id = iida_block->section();
 261  E :    size_t name_len = strlen(module_name);
 262    :    BlockGraph::Block* int_block = block_graph->AddBlock(
 263    :        BlockGraph::DATA_BLOCK, kPtrSize,
 264  E :        base::StringPrintf("Import Name Table: %s", module_name));
 265    :    BlockGraph::Block* dll_name_block = block_graph->AddBlock(
 266    :        BlockGraph::DATA_BLOCK, name_len + 1,
 267  E :        base::StringPrintf("Import Name: %s", module_name));
 268  E :    if (int_block == NULL || dll_name_block == NULL) {
 269  i :      LOG(ERROR) << "Unable to create blocks for Image Import Descriptor.";
 270  i :      return false;
 271    :    }
 272    :  
 273    :    // NOTE: If PEParser was modified to parse a single INT block, we could be
 274    :    //     extending/reusing it rather than creating a new INT per module.
 275  E :    int_block->set_section(iida_section_id);
 276  E :    int_block->set_attribute(BlockGraph::PE_PARSED);
 277    :    int_block->SetLabel(
 278    :        0,
 279    :        StringPrintf("%s INT: NULL entry", module_name),
 280  E :        BlockGraph::DATA_LABEL);
 281  E :    if (int_block->AllocateData(kPtrSize) == NULL) {
 282  i :      LOG(ERROR) << "Failed to allocate block data.";
 283  i :      return false;
 284    :    }
 285    :  
 286    :    // We use the DLL name block and extend it. This keeps things well ordered
 287    :    // when writing back the image using a canonical ordering.
 288  E :    dll_name_block->set_section(iida_section_id);
 289  E :    dll_name_block->set_attribute(BlockGraph::PE_PARSED);
 290  E :    if (dll_name_block->CopyData(name_len + 1, module_name) == NULL) {
 291  i :      LOG(ERROR) << "Failed to copy block data.";
 292  i :      return false;
 293    :    }
 294    :  
 295    :    // Add another NULL entry to the IAT block, but only if it does not already
 296    :    // consist of a single NULL entry (meaning it was just created). We are purely
 297    :    // extending this block, so no need to use the data insertion functions.
 298  E :    Offset iat_offset = 0;
 299  E :    if (iat_block->size() != kPtrSize) {
 300  E :      iat_offset = iat_block->size();
 301  E :      size_t iat_size = iat_offset + kPtrSize;
 302  E :      iat_block->set_size(iat_size);
 303  E :      iat_block->ResizeData(iat_size);
 304  E :      DCHECK_EQ(iat_size, iat_block->size());
 305  E :      DCHECK_EQ(iat_size, iat_block->data_size());
 306    :    }
 307    :  
 308    :    // Add a label for debugging purposes.
 309    :    iat_block->SetLabel(iat_offset,
 310    :                        StringPrintf("%s: NULL thunk", module_name),
 311  E :                        BlockGraph::DATA_LABEL);
 312    :  
 313    :    // Hook up these blocks.
 314    :    iida.SetReference(BlockGraph::RELATIVE_REF,
 315  E :                      iida[descriptor_count].OriginalFirstThunk, int_block, 0, 0);
 316    :    iida.SetReference(BlockGraph::RELATIVE_REF,
 317    :                      iida[descriptor_count].FirstThunk, iat_block, iat_offset,
 318  E :                      iat_offset);
 319    :    iida.SetReference(BlockGraph::RELATIVE_REF,
 320  E :                      iida[descriptor_count].Name, dll_name_block, 0, 0);
 321    :  
 322    :    // Finally, return the descriptor.
 323  E :    if (!iid->Init(new_iid_offset, iida_block)) {
 324  i :      LOG(ERROR) << "Unable to cast Image Import Descriptor.";
 325  i :      return false;
 326    :    }
 327    :  
 328  E :    *added = true;
 329    :  
 330  E :    return true;
 331  E :  }
 332    :  
 333    :  // Finds or adds an imported symbol to the given module (represented by its
 334    :  // import descriptor). Returns true on success, false otherwise. On success
 335    :  // returns a reference to the module's IAT entry. New entries are always added
 336    :  // to the end of the table so as not to invalidate any other unlinked references
 337    :  // (not part of the BlockGraph, so unable to be patched up) into the table.
 338    :  bool FindOrAddImportedSymbol(const char* symbol_name,
 339    :                               const ImageImportDescriptor& iid,
 340    :                               BlockGraph* block_graph,
 341    :                               BlockGraph::Block* iat_block,
 342    :                               size_t* iat_index,
 343  E :                               bool* added) {
 344  E :    DCHECK(symbol_name != NULL);
 345  E :    DCHECK(block_graph != NULL);
 346  E :    DCHECK(iat_block != NULL);
 347  E :    DCHECK(iat_index != NULL);
 348  E :    DCHECK(added != NULL);
 349    :  
 350  E :    *added = false;
 351    :  
 352  E :    TypedBlock<IMAGE_IMPORT_BY_NAME*> hna, iat;
 353    :    if (!iid.Dereference(iid->OriginalFirstThunk, &hna) ||
 354  E :        !iid.Dereference(iid->FirstThunk, &iat)) {
 355  i :      LOG(ERROR) << "Unable to dereference OriginalFirstThunk/FirstThunk.";
 356  i :      return false;
 357    :    }
 358    :  
 359    :    // Loop through the existing imports and see if we can't find a match. If so,
 360    :    // we don't need to import the symbol as it is already imported. The array is
 361    :    // NULL terminated so we loop through all elements except for the last one.
 362  E :    size_t i = 0;
 363  E :    for (; i < hna.ElementCount() && i < iat.ElementCount(); ++i) {
 364  E :      ConstTypedBlock<IMAGE_THUNK_DATA32> thunk;
 365  E :      if (!thunk.Init(hna.OffsetOf(hna[i]), hna.block())) {
 366  i :        LOG(ERROR) << "Unable to dereference IMAGE_THUNK_DATA32.";
 367  i :        return false;
 368    :      }
 369    :  
 370    :      // Is this an ordinal import? Skip it, as we have no way of
 371    :      // knowing the actual name of the symbol.
 372  E :      if (IMAGE_SNAP_BY_ORDINAL(thunk->u1.Ordinal))
 373  E :        continue;
 374    :  
 375    :      // Have no reference? Then terminate the iteration.
 376  E :      if (!thunk.HasReference(thunk->u1.AddressOfData)) {
 377    :        // We sanity check that the actual data is null.
 378  E :        DCHECK_EQ(0u, thunk->u1.AddressOfData);
 379  E :        break;
 380    :      }
 381    :  
 382    :      // Otherwise this should point to an IMAGE_IMPORT_BY_NAME structure.
 383  E :      ImageImportByName iibn;
 384  E :      if (!hna.Dereference(hna[i], &iibn)) {
 385  i :        LOG(ERROR) << "Unable to dereference IMAGE_IMPORT_BY_NAME.";
 386  i :        return false;
 387    :      }
 388    :  
 389    :      // Check to see if this symbol matches that of the current image import
 390    :      // by name.
 391    :      size_t max_len = iibn.block()->data_size() - iibn.offset() -
 392  E :          offsetof(IMAGE_IMPORT_BY_NAME, Name);
 393  E :      const char* import_name = reinterpret_cast<const char*>(iibn->Name);
 394  E :      if (::strncmp(import_name, symbol_name, max_len) == 0) {
 395  E :        *iat_index = i;
 396  E :        return true;
 397    :      }
 398  E :    }
 399    :  
 400    :    // Figure out how large the data needs to be to hold the name of this exported
 401    :    // symbol.  The IMAGE_IMPORT_BY_NAME struct has a WORD ordinal and a variable
 402    :    // sized field for the null-terminated function name. Each entry should be
 403    :    // WORD aligned, and will be referenced from the import address table and the
 404    :    // import name table.
 405  E :    size_t symbol_name_len = strlen(symbol_name);
 406    :    size_t iibn_size = sizeof(WORD) + common::AlignUp(symbol_name_len + 1,
 407  E :                                                      sizeof(WORD));
 408    :  
 409    :    // Get the DLL name. We will be inserting the IIBN entry to the block
 410    :    // containing it immediately prior to the DLL name.
 411  E :    String dll_name;
 412  E :    if (!iid.Dereference(iid->Name, &dll_name)) {
 413  i :      LOG(ERROR) << "Unable to dereference DLL name.";
 414  i :      return false;
 415    :    }
 416  E :    Offset iibn_offset = dll_name.offset();
 417  E :    dll_name.block()->InsertData(iibn_offset, iibn_size, true);
 418    :  
 419    :    // Populate the import struct.
 420  E :    TypedBlock<IMAGE_IMPORT_BY_NAME> iibn;
 421  E :    if (!iibn.InitWithSize(iibn_offset, iibn_size, dll_name.block())) {
 422  i :      LOG(ERROR) << "Unable to dereference new IMAGE_IMPORT_BY_NAME.";
 423  i :      return false;
 424    :    }
 425  E :    iibn->Hint = 0;
 426    :    base::strlcpy(reinterpret_cast<char*>(iibn->Name), symbol_name,
 427  E :                  symbol_name_len + 1);
 428    :  
 429    :    // Make room in the INT and the IAT for the new symbol. We place it
 430    :    // after the last entry for this module.
 431  E :    Offset int_offset = hna.OffsetOf(hna[i]);
 432  E :    Offset iat_offset = iat.OffsetOf(iat[i]);
 433    :    // We're pointed at the terminating zero. The position we're pointing at can
 434    :    // be the destination for references (in the normal case where someone is
 435    :    // using the import). However, in the special case where the IAT and the INT
 436    :    // are empty, our slot may also be pointed at by the import descriptor.
 437    :    // If we were to insert data at this position, we'd push the import
 438    :    // descriptor's pointer forward, past our new entry. To avoid this, we insert
 439    :    // the new data after the terminating zero we're pointing at, then usurp the
 440    :    // previously terminating zero for our entry.
 441  E :    hna.block()->InsertData(int_offset + kPtrSize, kPtrSize, true);
 442  E :    iat.block()->InsertData(iat_offset + kPtrSize, kPtrSize, true);
 443    :  
 444    :    // Because of the usurping mentioned above, we manually move any existing
 445    :    // labels.
 446  E :    BlockGraph::Label label;
 447  E :    if (hna.block()->GetLabel(int_offset, &label)) {
 448  E :      hna.block()->RemoveLabel(int_offset);
 449  E :      hna.block()->SetLabel(int_offset + kPtrSize, label);
 450    :    }
 451  E :    if (iat.block()->GetLabel(iat_offset, &label)) {
 452  E :      iat.block()->RemoveLabel(iat_offset);
 453  E :      iat.block()->SetLabel(iat_offset + kPtrSize, label);
 454    :    }
 455    :  
 456    :    // Add the new labels. We have to get the module_name at this point
 457    :    // because it may have been moved with our insertions above.
 458  E :    String module_name;
 459  E :    if (!iid.Dereference(iid->Name, &module_name)) {
 460  i :      LOG(ERROR) << "Unable to dereference import name.";
 461  i :      return false;
 462    :    }
 463    :    hna.block()->SetLabel(
 464    :        int_offset,
 465    :        StringPrintf("%s INT: %s", module_name->string, symbol_name),
 466  E :        BlockGraph::DATA_LABEL);
 467    :    iat.block()->SetLabel(
 468    :        iat_offset,
 469    :        StringPrintf("%s IAT: %s", module_name->string, symbol_name),
 470  E :        BlockGraph::DATA_LABEL);
 471    :  
 472    :    // Hook up the newly created IMAGE_IMPORT_BY_NAME to both tables.
 473    :    BlockGraph::Reference iibn_ref(BlockGraph::RELATIVE_REF,
 474    :                                   kPtrSize,
 475    :                                   iibn.block(),
 476    :                                   iibn.offset(),
 477  E :                                   iibn.offset());
 478  E :    hna.block()->SetReference(int_offset, iibn_ref);
 479  E :    iat.block()->SetReference(iat_offset, iibn_ref);
 480    :  
 481    :    // Return the reference to the IAT entry for the newly imported symbol.
 482  E :    *iat_index = i;
 483  E :    *added = true;
 484    :  
 485  E :    return true;
 486  E :  }
 487    :  
 488    :  }  // namespace
 489    :  
 490    :  
 491    :  const char AddImportsTransform::kTransformName[] = "AddImportsTransform";
 492    :  
 493    :  AddImportsTransform::AddImportsTransform()
 494  E :      : modules_added_(0), symbols_added_(0) {
 495  E :  }
 496    :  
 497    :  bool AddImportsTransform::TransformBlockGraph(
 498  E :      BlockGraph* block_graph, BlockGraph::Block* dos_header_block) {
 499  E :    DCHECK(block_graph != NULL);
 500  E :    DCHECK(dos_header_block != NULL);
 501    :  
 502  E :    modules_added_ = 0;
 503  E :    symbols_added_ = 0;
 504    :  
 505  E :    DosHeader dos_header;
 506  E :    NtHeaders nt_headers;
 507    :    if (!dos_header.Init(0, dos_header_block) ||
 508  E :        !dos_header.Dereference(dos_header->e_lfanew, &nt_headers)) {
 509  i :      LOG(ERROR) << "Unable to cast image headers.";
 510  i :      return false;
 511    :    }
 512    :  
 513    :    // Get the block containing the image import directory. In general it's not
 514    :    // safe to keep around a pointer into a TypedBlock, but we won't be resizing
 515    :    // or reallocating the data within nt_headers so this is safe.
 516    :    if (!EnsureDataDirectoryExists(IMAGE_DIRECTORY_ENTRY_IMPORT,
 517    :                                   "Image Import Descriptor Array",
 518    :                                   sizeof(IMAGE_IMPORT_DESCRIPTOR),
 519    :                                   block_graph,
 520  E :                                   nt_headers.block())) {
 521  i :      LOG(ERROR) << "Failed to create Image Import Descriptor Array.";
 522  i :      return false;
 523    :    }
 524    :    IMAGE_DATA_DIRECTORY* import_directory =
 525  E :        nt_headers->OptionalHeader.DataDirectory + IMAGE_DIRECTORY_ENTRY_IMPORT;
 526  E :    DCHECK(nt_headers.HasReference(import_directory->VirtualAddress));
 527    :  
 528  E :    ImageImportDescriptor image_import_descriptor;
 529    :    if (!nt_headers.Dereference(import_directory->VirtualAddress,
 530  E :                                &image_import_descriptor)) {
 531    :      // This could happen if the image import descriptor array is empty, and
 532    :      // terminated by a *partial* null entry. However, we've not yet seen that.
 533  i :      LOG(ERROR) << "Failed to dereference Image Import Descriptor Array.";
 534  i :      return false;
 535    :    }
 536    :  
 537    :    // We expect the image import descriptor to have been parsed as its own block,
 538    :    // so the reference needs to be to offset 0.
 539  E :    if (image_import_descriptor.offset() != 0) {
 540  i :      LOG(ERROR) << "Unexpected offset on Image Import Descriptor.";
 541  i :      return false;
 542    :    }
 543    :  
 544  E :    image_import_descriptor_block_ = image_import_descriptor.block();
 545    :  
 546    :    // Similarly, get the block containing the IAT.
 547    :    if (!EnsureDataDirectoryExists(IMAGE_DIRECTORY_ENTRY_IAT,
 548    :                                   "Import Address Table",
 549    :                                   kPtrSize,
 550    :                                   block_graph,
 551  E :                                   nt_headers.block())) {
 552  i :      LOG(ERROR) << "Failed to create Import Address Table.";
 553  i :      return false;
 554    :    }
 555    :    IMAGE_DATA_DIRECTORY* iat_directory =
 556  E :        nt_headers->OptionalHeader.DataDirectory + IMAGE_DIRECTORY_ENTRY_IAT;
 557  E :    DCHECK(nt_headers.HasReference(iat_directory->VirtualAddress));
 558  E :    TypedBlock<RelativeAddress> iat;
 559  E :    if (!nt_headers.Dereference(iat_directory->VirtualAddress, &iat)) {
 560  i :      LOG(ERROR) << "Failed to dereference Import Address Table.";
 561  i :      return false;
 562    :    }
 563    :  
 564    :    // We expect the import address table to have been parsed as its own block,
 565    :    // so the reference needs to be to offset 0.
 566  E :    if (iat.offset() != 0) {
 567  i :      LOG(ERROR) << "Unexpected offset on Image Address Table";
 568  i :      return false;
 569    :    }
 570  E :    import_address_table_block_ = iat.block();
 571    :  
 572    :    // Handle each library individually.
 573  E :    for (size_t i = 0; i < imported_modules_.size(); ++i) {
 574  E :      ImportedModule* module = imported_modules_[i];
 575  E :      if (module->size() == 0)
 576  i :        continue;
 577    :  
 578    :      // First find or create an entry for this module in the Image Import
 579    :      // Descriptor Array.
 580  E :      ImageImportDescriptor iid;
 581  E :      bool module_added = false;
 582    :      if (!FindOrAddImageImportDescriptor(module->name().c_str(),
 583    :                                          block_graph,
 584    :                                          image_import_descriptor_block_,
 585    :                                          import_address_table_block_,
 586    :                                          &iid,
 587  E :                                          &module_added)) {
 588  i :        LOG(ERROR) << "Failed to find or import module.";
 589  i :        return false;
 590    :      }
 591  E :      modules_added_ += module_added;
 592    :  
 593    :      // This should always succeed as iid was initialized earlier.
 594  E :      bool inited = module->import_descriptor_.Init(iid.offset(), iid.block());
 595  E :      DCHECK(inited);
 596    :  
 597  E :      for (size_t j = 0; j < module->size(); ++j) {
 598  E :        ImportedModule::Symbol& symbol = module->symbols_[j];
 599    :  
 600    :        // Now, for each symbol get the offset of the IAT entry. This will create
 601    :        // the entry (and all accompanying structures) if necessary.
 602  E :        size_t symbol_index = ImportedModule::kInvalidIndex;
 603  E :        bool symbol_added = false;
 604    :        if (!FindOrAddImportedSymbol(symbol.name.c_str(),
 605    :                                     iid,
 606    :                                     block_graph,
 607    :                                     import_address_table_block_,
 608    :                                     &symbol_index,
 609  E :                                     &symbol_added)) {
 610  i :          LOG(ERROR) << "Failed to find or import symbol.";
 611  i :          return false;
 612    :        }
 613  E :        symbols_added_ += symbol_added;
 614  E :        symbol.index = symbol_index;
 615  E :      }
 616  E :    }
 617    :  
 618    :    // Update the data directory sizes.
 619  E :    import_directory->Size = image_import_descriptor_block_->size();
 620  E :    iat_directory->Size = import_address_table_block_->size();
 621    :  
 622  E :    return true;
 623  E :  }
 624    :  
 625    :  const size_t AddImportsTransform::ImportedModule::kInvalidIndex = -1;
 626    :  
 627    :  size_t AddImportsTransform::ImportedModule::AddSymbol(
 628  E :      const base::StringPiece& symbol_name) {
 629  E :    Symbol symbol = {symbol_name.as_string(), kInvalidIndex};
 630  E :    symbols_.push_back(symbol);
 631  E :    return symbols_.size() - 1;
 632  E :  }
 633    :  
 634    :  bool AddImportsTransform::ImportedModule::GetSymbolReference(
 635  E :      size_t index, BlockGraph::Reference* abs_reference) const {
 636  E :    DCHECK_LT(index, symbols_.size());
 637  E :    DCHECK(abs_reference != NULL);
 638    :  
 639  E :    size_t symbol_index = symbols_[index].index;
 640  E :    if (import_descriptor_.block() == NULL || symbol_index == kInvalidIndex) {
 641  i :      LOG(ERROR) << "Called GetAbsReference on an unitialized ImportedSymbol.";
 642  i :      return false;
 643    :    }
 644    :  
 645  E :    ImageThunkData32 thunks;
 646    :    if (!import_descriptor_.Dereference(import_descriptor_->FirstThunk,
 647  E :                                        &thunks)) {
 648  i :      LOG(ERROR) << "Unable to dereference IMAGE_THUNK_DATA32.";
 649  i :      return false;
 650    :    }
 651    :  
 652  E :    if (symbol_index >= thunks.ElementCount()) {
 653  i :      LOG(ERROR) << "Invalid symbol_index for IAT.";
 654  i :      return false;
 655    :    }
 656    :  
 657  E :    Offset offset = thunks.OffsetOf(thunks[symbol_index].u1.AddressOfData);
 658    :    *abs_reference = BlockGraph::Reference(
 659  E :        BlockGraph::ABSOLUTE_REF, kPtrSize, thunks.block(), offset, offset);
 660    :  
 661  E :    return true;
 662  E :  }
 663    :  
 664    :  }  // namespace transforms
 665    :  }  // namespace pe

Coverage information generated Thu Sep 06 11:30:46 2012.