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

CoverageLines executed / instrumented / missingexe / inst / missLanguageGroup
80.6%2282830.C++source

Line-by-line coverage:

   1    :  // Copyright 2012 Google Inc. All Rights Reserved.
   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 exist and find_only is false, it will allocate a block with
 145    :  // the given name and size.
 146    :  bool FindOrAddDataDirectory(bool find_only,
 147    :                              size_t directory_index,
 148    :                              const char* block_name,
 149    :                              size_t block_size,
 150    :                              BlockGraph* block_graph,
 151    :                              BlockGraph::Block* nt_headers_block,
 152  E :                              bool* exists) {
 153  E :    DCHECK(block_name != NULL);
 154    :    DCHECK_LT(directory_index,
 155  E :              static_cast<size_t>(IMAGE_NUMBEROF_DIRECTORY_ENTRIES));
 156  E :    DCHECK_GT(block_size, 0u);
 157  E :    DCHECK(block_graph != NULL);
 158  E :    DCHECK(nt_headers_block != NULL);
 159  E :    DCHECK(exists != NULL);
 160    :  
 161  E :    *exists = false;
 162    :  
 163  E :    NtHeaders nt_headers;
 164  E :    if (!nt_headers.Init(0, nt_headers_block)) {
 165  i :      LOG(ERROR) << "Unable to cast NT headers.";
 166  i :      return false;
 167    :    }
 168    :  
 169    :    IMAGE_DATA_DIRECTORY* data_directory =
 170  E :        nt_headers->OptionalHeader.DataDirectory + directory_index;
 171    :  
 172    :    // No entry? Then make a zero initialized block that is stored in .rdata,
 173    :    // where all of these structures live.
 174  E :    if (!nt_headers.HasReference(data_directory->VirtualAddress)) {
 175    :      // We don't need to create the entry if we're exploring only.
 176  E :      if (find_only)
 177  i :        return true;
 178    :  
 179    :      BlockGraph::Section* section = block_graph->FindOrAddSection(
 180  E :          kReadOnlyDataSectionName, kReadOnlyDataCharacteristics);
 181  E :      DCHECK(section != NULL);
 182    :  
 183    :      BlockGraph::Block* block = block_graph->AddBlock(
 184  E :          BlockGraph::DATA_BLOCK, block_size, block_name);
 185  E :      DCHECK(block != NULL);
 186  E :      block->set_section(section->id());
 187  E :      block->set_attribute(BlockGraph::PE_PARSED);
 188    :  
 189    :      // We need to actually allocate the data so that future TypedBlock
 190    :      // dereferences will work.
 191  E :      if (block->AllocateData(block_size) == NULL) {
 192  i :        LOG(ERROR) << "Failed to allocate block data.";
 193  i :        return false;
 194    :      }
 195    :  
 196    :      // Hook it up to the NT header.
 197    :      nt_headers.SetReference(BlockGraph::RELATIVE_REF,
 198    :                              data_directory->VirtualAddress,
 199    :                              block,
 200  E :                              0, 0);
 201  E :      data_directory->Size = block_size;
 202    :    }
 203    :  
 204  E :    *exists = true;
 205  E :    return true;
 206  E :  }
 207    :  
 208    :  // Finds or creates an Image Import Descriptor block for the given library.
 209    :  // Returns true on success, false otherwise.
 210    :  bool FindOrAddImageImportDescriptor(bool find_only,
 211    :                                      const char* module_name,
 212    :                                      BlockGraph* block_graph,
 213    :                                      BlockGraph::Block* iida_block,
 214    :                                      BlockGraph::Block* iat_block,
 215    :                                      ImageImportDescriptor* iid,
 216    :                                      bool* added,
 217  E :                                      bool* exists) {
 218  E :    DCHECK(module_name != NULL);
 219  E :    DCHECK(block_graph != NULL);
 220  E :    DCHECK(iida_block != NULL);
 221  E :    DCHECK(iat_block != NULL);
 222  E :    DCHECK(iid != NULL);
 223  E :    DCHECK(added != NULL);
 224  E :    DCHECK(exists != NULL);
 225    :  
 226  E :    *added = false;
 227  E :    *exists = false;
 228    :  
 229  E :    ImageImportDescriptor iida;
 230  E :    if (!iida.Init(0, iida_block)) {
 231  i :      LOG(ERROR) << "Unable to cast Image Import Descriptor.";
 232  i :      return false;
 233    :    }
 234    :  
 235    :    // The array is NULL terminated with a potentially incomplete descriptor so
 236    :    // we can't use ElementCount - 1.
 237  E :    DCHECK_GT(iida_block->size(), 0U);
 238    :    size_t descriptor_count =
 239    :        (common::AlignUp(iida_block->size(), sizeof(IMAGE_IMPORT_DESCRIPTOR)) /
 240  E :            sizeof(IMAGE_IMPORT_DESCRIPTOR)) - 1;
 241    :  
 242  E :    for (size_t iida_index = 0; iida_index < descriptor_count; ++iida_index) {
 243  E :      String dll_name;
 244  E :      if (!iida.Dereference(iida[iida_index].Name, &dll_name)) {
 245  i :        LOG(ERROR) << "Unable to dereference DLL name.";
 246  i :        return false;
 247    :      }
 248    :  
 249  E :      size_t max_len = dll_name.ElementCount();
 250  E :      if (base::strncasecmp(dll_name->string, module_name, max_len) == 0) {
 251    :        // This should never fail, but we sanity check it nonetheless.
 252  E :        bool result = iid->Init(iida.OffsetOf(iida[iida_index]), iida.block());
 253  E :        DCHECK(result);
 254  E :        *exists = true;
 255  E :        return true;
 256    :      }
 257  E :    }
 258    :  
 259    :    // If we get here then the entry doesn't exist. If we've been asked to only
 260    :    // search for it then we can return early.
 261  E :    if (find_only)
 262  E :      return true;
 263    :  
 264    :    // Create room for the new descriptor, which we'll tack on to the end of the
 265    :    // array, but before the NULL terminator. We use 'InsertData' so that all
 266    :    // labels are patched up.
 267  E :    Offset new_iid_offset = descriptor_count * sizeof(IMAGE_IMPORT_DESCRIPTOR);
 268    :    iida_block->InsertData(
 269  E :        new_iid_offset, sizeof(IMAGE_IMPORT_DESCRIPTOR), true);
 270    :    iida_block->SetLabel(
 271    :        new_iid_offset,
 272    :        base::StringPrintf("Image Import Descriptor: %s", module_name),
 273  E :        BlockGraph::DATA_LABEL);
 274    :  
 275    :    // We expect the new entry to be dereferencable using iida[descriptor_count].
 276  E :    DCHECK_GT(iida.ElementCount(), descriptor_count);
 277    :  
 278    :    // Create the various child structures that will be pointed to by the
 279    :    // import descriptor. The INT block and the IAT block are NULL terminated
 280    :    // lists of pointers, and the terminating NULL is allocated. We don't yet
 281    :    // allocate a block to hold the import names, deferring that for later.
 282  E :    BlockGraph::SectionId iida_section_id = iida_block->section();
 283  E :    size_t name_len = strlen(module_name);
 284    :    BlockGraph::Block* int_block = block_graph->AddBlock(
 285    :        BlockGraph::DATA_BLOCK, kPtrSize,
 286  E :        base::StringPrintf("Import Name Table: %s", module_name));
 287    :    BlockGraph::Block* dll_name_block = block_graph->AddBlock(
 288    :        BlockGraph::DATA_BLOCK, name_len + 1,
 289  E :        base::StringPrintf("Import Name: %s", module_name));
 290  E :    if (int_block == NULL || dll_name_block == NULL) {
 291  i :      LOG(ERROR) << "Unable to create blocks for Image Import Descriptor.";
 292  i :      return false;
 293    :    }
 294    :  
 295    :    // NOTE: If PEParser was modified to parse a single INT block, we could be
 296    :    //     extending/reusing it rather than creating a new INT per module.
 297  E :    int_block->set_section(iida_section_id);
 298  E :    int_block->set_attribute(BlockGraph::PE_PARSED);
 299    :    int_block->SetLabel(
 300    :        0,
 301    :        base::StringPrintf("%s INT: NULL entry", module_name),
 302  E :        BlockGraph::DATA_LABEL);
 303  E :    if (int_block->AllocateData(kPtrSize) == NULL) {
 304  i :      LOG(ERROR) << "Failed to allocate block data.";
 305  i :      return false;
 306    :    }
 307    :  
 308    :    // We use the DLL name block and extend it. This keeps things well ordered
 309    :    // when writing back the image using a canonical ordering.
 310  E :    dll_name_block->set_section(iida_section_id);
 311  E :    dll_name_block->set_attribute(BlockGraph::PE_PARSED);
 312  E :    if (dll_name_block->CopyData(name_len + 1, module_name) == NULL) {
 313  i :      LOG(ERROR) << "Failed to copy block data.";
 314  i :      return false;
 315    :    }
 316    :  
 317    :    // Add another NULL entry to the IAT block, but only if it does not already
 318    :    // consist of a single NULL entry (meaning it was just created). We are purely
 319    :    // extending this block, so no need to use the data insertion functions.
 320  E :    Offset iat_offset = 0;
 321  E :    if (iat_block->size() != kPtrSize) {
 322  E :      iat_offset = iat_block->size();
 323  E :      size_t iat_size = iat_offset + kPtrSize;
 324  E :      iat_block->set_size(iat_size);
 325  E :      iat_block->ResizeData(iat_size);
 326  E :      DCHECK_EQ(iat_size, iat_block->size());
 327  E :      DCHECK_EQ(iat_size, iat_block->data_size());
 328    :    }
 329    :  
 330    :    // Add a label for debugging purposes.
 331    :    iat_block->SetLabel(iat_offset,
 332    :                        base::StringPrintf("%s: NULL thunk", module_name),
 333  E :                        BlockGraph::DATA_LABEL);
 334    :  
 335    :    // Hook up these blocks.
 336    :    iida.SetReference(BlockGraph::RELATIVE_REF,
 337  E :                      iida[descriptor_count].OriginalFirstThunk, int_block, 0, 0);
 338    :    iida.SetReference(BlockGraph::RELATIVE_REF,
 339    :                      iida[descriptor_count].FirstThunk, iat_block, iat_offset,
 340  E :                      iat_offset);
 341    :    iida.SetReference(BlockGraph::RELATIVE_REF,
 342  E :                      iida[descriptor_count].Name, dll_name_block, 0, 0);
 343    :  
 344    :    // Finally, return the descriptor.
 345  E :    if (!iid->Init(new_iid_offset, iida_block)) {
 346  i :      LOG(ERROR) << "Unable to cast Image Import Descriptor.";
 347  i :      return false;
 348    :    }
 349    :  
 350  E :    *added = true;
 351  E :    *exists = true;
 352    :  
 353  E :    return true;
 354  E :  }
 355    :  
 356    :  // Finds or adds an imported symbol to the given module (represented by its
 357    :  // import descriptor). Returns true on success, false otherwise. On success
 358    :  // returns a reference to the module's IAT entry. New entries are always added
 359    :  // to the end of the table so as not to invalidate any other unlinked references
 360    :  // (not part of the BlockGraph, so unable to be patched up) into the table.
 361    :  bool FindOrAddImportedSymbol(bool find_only,
 362    :                               const char* symbol_name,
 363    :                               const ImageImportDescriptor& iid,
 364    :                               BlockGraph* block_graph,
 365    :                               BlockGraph::Block* iat_block,
 366    :                               size_t* iat_index,
 367  E :                               bool* added) {
 368  E :    DCHECK(symbol_name != NULL);
 369  E :    DCHECK(block_graph != NULL);
 370  E :    DCHECK(iat_block != NULL);
 371  E :    DCHECK(iat_index != NULL);
 372  E :    DCHECK(added != NULL);
 373    :  
 374  E :    *iat_index = AddImportsTransform::ImportedModule::kInvalidIatIndex;
 375  E :    *added = false;
 376    :  
 377  E :    TypedBlock<IMAGE_IMPORT_BY_NAME*> hna, iat;
 378    :    if (!iid.Dereference(iid->OriginalFirstThunk, &hna) ||
 379  E :        !iid.Dereference(iid->FirstThunk, &iat)) {
 380  i :      LOG(ERROR) << "Unable to dereference OriginalFirstThunk/FirstThunk.";
 381  i :      return false;
 382    :    }
 383    :  
 384    :    // Loop through the existing imports and see if we can't find a match. If so,
 385    :    // we don't need to import the symbol as it is already imported. The array is
 386    :    // NULL terminated so we loop through all elements except for the last one.
 387  E :    size_t i = 0;
 388  E :    for (; i < hna.ElementCount() && i < iat.ElementCount(); ++i) {
 389  E :      ConstTypedBlock<IMAGE_THUNK_DATA32> thunk;
 390  E :      if (!thunk.Init(hna.OffsetOf(hna[i]), hna.block())) {
 391  i :        LOG(ERROR) << "Unable to dereference IMAGE_THUNK_DATA32.";
 392  i :        return false;
 393    :      }
 394    :  
 395    :      // Is this an ordinal import? Skip it, as we have no way of
 396    :      // knowing the actual name of the symbol.
 397  E :      if (IMAGE_SNAP_BY_ORDINAL(thunk->u1.Ordinal))
 398  E :        continue;
 399    :  
 400    :      // Have no reference? Then terminate the iteration.
 401  E :      if (!thunk.HasReference(thunk->u1.AddressOfData)) {
 402    :        // We sanity check that the actual data is null.
 403  E :        DCHECK_EQ(0u, thunk->u1.AddressOfData);
 404  E :        break;
 405    :      }
 406    :  
 407    :      // Otherwise this should point to an IMAGE_IMPORT_BY_NAME structure.
 408  E :      ImageImportByName iibn;
 409  E :      if (!hna.Dereference(hna[i], &iibn)) {
 410  i :        LOG(ERROR) << "Unable to dereference IMAGE_IMPORT_BY_NAME.";
 411  i :        return false;
 412    :      }
 413    :  
 414    :      // Check to see if this symbol matches that of the current image import
 415    :      // by name.
 416    :      size_t max_len = iibn.block()->data_size() - iibn.offset() -
 417  E :          offsetof(IMAGE_IMPORT_BY_NAME, Name);
 418  E :      const char* import_name = reinterpret_cast<const char*>(iibn->Name);
 419  E :      if (::strncmp(import_name, symbol_name, max_len) == 0) {
 420  E :        *iat_index = i;
 421  E :        return true;
 422    :      }
 423  E :    }
 424    :  
 425    :    // If we get here then the entry doesn't exist. If we've been asked to only
 426    :    // search for it then we can return early.
 427  E :    if (find_only)
 428  E :      return true;
 429    :  
 430    :    // Figure out how large the data needs to be to hold the name of this exported
 431    :    // symbol.  The IMAGE_IMPORT_BY_NAME struct has a WORD ordinal and a variable
 432    :    // sized field for the null-terminated function name. Each entry should be
 433    :    // WORD aligned, and will be referenced from the import address table and the
 434    :    // import name table.
 435  E :    size_t symbol_name_len = strlen(symbol_name);
 436    :    size_t iibn_size = sizeof(WORD) + common::AlignUp(symbol_name_len + 1,
 437  E :                                                      sizeof(WORD));
 438    :  
 439    :    // Get the DLL name. We will be inserting the IIBN entry to the block
 440    :    // containing it immediately prior to the DLL name.
 441  E :    String dll_name;
 442  E :    if (!iid.Dereference(iid->Name, &dll_name)) {
 443  i :      LOG(ERROR) << "Unable to dereference DLL name.";
 444  i :      return false;
 445    :    }
 446  E :    Offset iibn_offset = dll_name.offset();
 447  E :    dll_name.block()->InsertData(iibn_offset, iibn_size, true);
 448    :  
 449    :    // Populate the import struct.
 450  E :    TypedBlock<IMAGE_IMPORT_BY_NAME> iibn;
 451  E :    if (!iibn.InitWithSize(iibn_offset, iibn_size, dll_name.block())) {
 452  i :      LOG(ERROR) << "Unable to dereference new IMAGE_IMPORT_BY_NAME.";
 453  i :      return false;
 454    :    }
 455  E :    iibn->Hint = 0;
 456    :    base::strlcpy(reinterpret_cast<char*>(iibn->Name), symbol_name,
 457  E :                  symbol_name_len + 1);
 458    :  
 459    :    // Make room in the INT and the IAT for the new symbol. We place it
 460    :    // after the last entry for this module.
 461  E :    Offset int_offset = hna.OffsetOf(hna[i]);
 462  E :    Offset iat_offset = iat.OffsetOf(iat[i]);
 463    :    // We're pointed at the terminating zero. The position we're pointing at can
 464    :    // be the destination for references (in the normal case where someone is
 465    :    // using the import). However, in the special case where the IAT and the INT
 466    :    // are empty, our slot may also be pointed at by the import descriptor.
 467    :    // If we were to insert data at this position, we'd push the import
 468    :    // descriptor's pointer forward, past our new entry. To avoid this, we insert
 469    :    // the new data after the terminating zero we're pointing at, then usurp the
 470    :    // previously terminating zero for our entry.
 471  E :    hna.block()->InsertData(int_offset + kPtrSize, kPtrSize, true);
 472  E :    iat.block()->InsertData(iat_offset + kPtrSize, kPtrSize, true);
 473    :  
 474    :    // Because of the usurping mentioned above, we manually move any existing
 475    :    // labels.
 476  E :    BlockGraph::Label label;
 477  E :    if (hna.block()->GetLabel(int_offset, &label)) {
 478  E :      hna.block()->RemoveLabel(int_offset);
 479  E :      hna.block()->SetLabel(int_offset + kPtrSize, label);
 480    :    }
 481  E :    if (iat.block()->GetLabel(iat_offset, &label)) {
 482  E :      iat.block()->RemoveLabel(iat_offset);
 483  E :      iat.block()->SetLabel(iat_offset + kPtrSize, label);
 484    :    }
 485    :  
 486    :    // Add the new labels. We have to get the module_name at this point
 487    :    // because it may have been moved with our insertions above.
 488  E :    String module_name;
 489  E :    if (!iid.Dereference(iid->Name, &module_name)) {
 490  i :      LOG(ERROR) << "Unable to dereference import name.";
 491  i :      return false;
 492    :    }
 493    :    hna.block()->SetLabel(
 494    :        int_offset,
 495    :        base::StringPrintf("%s INT: %s", module_name->string, symbol_name),
 496  E :        BlockGraph::DATA_LABEL);
 497    :    iat.block()->SetLabel(
 498    :        iat_offset,
 499    :        base::StringPrintf("%s IAT: %s", module_name->string, symbol_name),
 500  E :        BlockGraph::DATA_LABEL);
 501    :  
 502    :    // Hook up the newly created IMAGE_IMPORT_BY_NAME to both tables.
 503    :    BlockGraph::Reference iibn_ref(BlockGraph::RELATIVE_REF,
 504    :                                   kPtrSize,
 505    :                                   iibn.block(),
 506    :                                   iibn.offset(),
 507  E :                                   iibn.offset());
 508  E :    hna.block()->SetReference(int_offset, iibn_ref);
 509  E :    iat.block()->SetReference(iat_offset, iibn_ref);
 510    :  
 511    :    // Return the reference to the IAT entry for the newly imported symbol.
 512  E :    *iat_index = i;
 513  E :    *added = true;
 514    :  
 515  E :    return true;
 516  E :  }
 517    :  
 518    :  }  // namespace
 519    :  
 520    :  const char AddImportsTransform::kTransformName[] = "AddImportsTransform";
 521    :  
 522    :  AddImportsTransform::AddImportsTransform()
 523  E :      : modules_added_(0), symbols_added_(0) {
 524  E :  }
 525    :  
 526    :  bool AddImportsTransform::TransformBlockGraph(
 527  E :      BlockGraph* block_graph, BlockGraph::Block* dos_header_block) {
 528  E :    DCHECK(block_graph != NULL);
 529  E :    DCHECK(dos_header_block != NULL);
 530    :  
 531    :    // Before doing anything, let's determine if we're on a strictly exploratory
 532    :    // mission. We don't want to add anything if all modules/symbols are
 533    :    // 'find only'.
 534  E :    bool find_only = true;
 535  E :    for (size_t i = 0; i < imported_modules_.size(); ++i) {
 536  E :      if (imported_modules_[i]->mode_ != ImportedModule::kFindOnly) {
 537  E :        find_only = false;
 538  E :        break;
 539    :      }
 540  E :    }
 541    :  
 542  E :    modules_added_ = 0;
 543  E :    symbols_added_ = 0;
 544    :  
 545  E :    DosHeader dos_header;
 546  E :    NtHeaders nt_headers;
 547    :    if (!dos_header.Init(0, dos_header_block) ||
 548  E :        !dos_header.Dereference(dos_header->e_lfanew, &nt_headers)) {
 549  i :      LOG(ERROR) << "Unable to cast image headers.";
 550  i :      return false;
 551    :    }
 552    :  
 553    :    // Get the block containing the image import directory. In general it's not
 554    :    // safe to keep around a pointer into a TypedBlock, but we won't be resizing
 555    :    // or reallocating the data within nt_headers so this is safe.
 556  E :    bool import_directory_exists = false;
 557    :    if (!FindOrAddDataDirectory(find_only,
 558    :                                IMAGE_DIRECTORY_ENTRY_IMPORT,
 559    :                                "Image Import Descriptor Array",
 560    :                                sizeof(IMAGE_IMPORT_DESCRIPTOR),
 561    :                                block_graph,
 562    :                                nt_headers.block(),
 563  E :                                &import_directory_exists)) {
 564  i :      LOG(ERROR) << "Failed to create Image Import Descriptor Array.";
 565  i :      return false;
 566    :    }
 567    :  
 568    :    // If we're on a find only mission and no import data directory exists then
 569    :    // there are *no* imports and we can quit early.
 570  E :    if (find_only && !import_directory_exists)
 571  i :      return true;
 572    :  
 573    :    // Look up the import directory.
 574  E :    DCHECK(import_directory_exists);
 575    :    IMAGE_DATA_DIRECTORY* import_directory =
 576  E :        nt_headers->OptionalHeader.DataDirectory + IMAGE_DIRECTORY_ENTRY_IMPORT;
 577  E :    DCHECK(nt_headers.HasReference(import_directory->VirtualAddress));
 578    :  
 579  E :    ImageImportDescriptor image_import_descriptor;
 580    :    if (!nt_headers.Dereference(import_directory->VirtualAddress,
 581  E :                                &image_import_descriptor)) {
 582    :      // This could happen if the image import descriptor array is empty, and
 583    :      // terminated by a *partial* null entry. However, we've not yet seen that.
 584  i :      LOG(ERROR) << "Failed to dereference Image Import Descriptor Array.";
 585  i :      return false;
 586    :    }
 587    :  
 588    :    // We expect the image import descriptor to have been parsed as its own block,
 589    :    // so the reference needs to be to offset 0.
 590  E :    if (image_import_descriptor.offset() != 0) {
 591  i :      LOG(ERROR) << "Unexpected offset on Image Import Descriptor.";
 592  i :      return false;
 593    :    }
 594    :  
 595  E :    image_import_descriptor_block_ = image_import_descriptor.block();
 596    :  
 597    :    // Similarly, get the block containing the IAT.
 598  E :    bool iat_directory_exists = false;
 599    :    if (!FindOrAddDataDirectory(find_only,
 600    :                                IMAGE_DIRECTORY_ENTRY_IAT,
 601    :                                "Import Address Table",
 602    :                                kPtrSize,
 603    :                                block_graph,
 604    :                                nt_headers.block(),
 605  E :                                &iat_directory_exists)) {
 606  i :      LOG(ERROR) << "Failed to create Import Address Table.";
 607  i :      return false;
 608    :    }
 609    :  
 610    :    // If we're on a find only mission and no import data directory exists then
 611    :    // there are *no* imports and we can quit early.
 612  E :    if (find_only && !iat_directory_exists)
 613  i :      return true;
 614    :  
 615    :    // Look up the IAT directory.
 616  E :    DCHECK(iat_directory_exists);
 617    :    IMAGE_DATA_DIRECTORY* iat_directory =
 618  E :        nt_headers->OptionalHeader.DataDirectory + IMAGE_DIRECTORY_ENTRY_IAT;
 619  E :    DCHECK(nt_headers.HasReference(iat_directory->VirtualAddress));
 620  E :    TypedBlock<RelativeAddress> iat;
 621  E :    if (!nt_headers.Dereference(iat_directory->VirtualAddress, &iat)) {
 622  i :      LOG(ERROR) << "Failed to dereference Import Address Table.";
 623  i :      return false;
 624    :    }
 625    :  
 626    :    // We expect the import address table to have been parsed as its own block,
 627    :    // so the reference needs to be to offset 0.
 628  E :    if (iat.offset() != 0) {
 629  i :      LOG(ERROR) << "Unexpected offset on Image Address Table";
 630  i :      return false;
 631    :    }
 632  E :    import_address_table_block_ = iat.block();
 633    :  
 634    :    // Handle each library individually.
 635  E :    for (size_t i = 0; i < imported_modules_.size(); ++i) {
 636  E :      ImportedModule* module = imported_modules_[i];
 637    :  
 638    :      // First find or create an entry for this module in the Image Import
 639    :      // Descriptor Array.
 640  E :      ImageImportDescriptor iid;
 641  E :      bool module_added = false;
 642  E :      bool module_exists = false;
 643    :      if (!FindOrAddImageImportDescriptor(
 644    :              module->mode_ == ImportedModule::kFindOnly,
 645    :              module->name().c_str(),
 646    :              block_graph,
 647    :              image_import_descriptor_block_,
 648    :              import_address_table_block_,
 649    :              &iid,
 650    :              &module_added,
 651  E :              &module_exists)) {
 652  i :        LOG(ERROR) << "Failed to find or import module.";
 653  i :        return false;
 654    :      }
 655    :  
 656    :      // If we're fact finding only and the module does not exist then we don't
 657    :      // need to look up its symbols.
 658  E :      if (module->mode_ == ImportedModule::kFindOnly && !module_exists) {
 659  E :        DCHECK(!module_added);
 660  E :        continue;
 661    :      }
 662    :  
 663  E :      DCHECK(module_exists);
 664  E :      module->added_ = module_added;
 665  E :      modules_added_ += module_added;
 666    :  
 667    :      // This should always succeed as iid was initialized earlier.
 668  E :      bool inited = module->import_descriptor_.Init(iid.offset(), iid.block());
 669  E :      DCHECK(inited);
 670    :  
 671  E :      for (size_t j = 0; j < module->size(); ++j) {
 672  E :        ImportedModule::Symbol& symbol = module->symbols_[j];
 673    :  
 674    :        // Now, for each symbol get the offset of the IAT entry. This will create
 675    :        // the entry (and all accompanying structures) if necessary.
 676  E :        size_t symbol_iat_index = ImportedModule::kInvalidIatIndex;
 677  E :        bool symbol_added = false;
 678    :        if (!FindOrAddImportedSymbol(symbol.mode == ImportedModule::kFindOnly,
 679    :                                     symbol.name.c_str(),
 680    :                                     iid,
 681    :                                     block_graph,
 682    :                                     import_address_table_block_,
 683    :                                     &symbol_iat_index,
 684  E :                                     &symbol_added)) {
 685  i :          LOG(ERROR) << "Failed to find or import symbol.";
 686  i :          return false;
 687    :        }
 688  E :        symbols_added_ += symbol_added;
 689  E :        symbol.iat_index = symbol_iat_index;
 690  E :        symbol.added = symbol_added;
 691  E :      }
 692  E :    }
 693    :  
 694    :    // Update the data directory sizes.
 695  E :    import_directory->Size = image_import_descriptor_block_->size();
 696  E :    iat_directory->Size = import_address_table_block_->size();
 697    :  
 698  E :    return true;
 699  E :  }
 700    :  
 701    :  const size_t AddImportsTransform::ImportedModule::kInvalidIatIndex = -1;
 702    :  
 703    :  size_t AddImportsTransform::ImportedModule::AddSymbol(
 704    :      const base::StringPiece& symbol_name,
 705  E :      ImportedModule::TransformMode mode) {
 706  E :    Symbol symbol = {symbol_name.as_string(), kInvalidIatIndex, mode};
 707  E :    symbols_.push_back(symbol);
 708    :  
 709  E :    if (mode != ImportedModule::kFindOnly)
 710  E :      mode_ = ImportedModule::kAlwaysImport;
 711    :  
 712  E :    return symbols_.size() - 1;
 713  E :  }
 714    :  
 715    :  bool AddImportsTransform::ImportedModule::GetSymbolReference(
 716  E :      size_t index, BlockGraph::Reference* abs_reference) const {
 717  E :    DCHECK_LT(index, symbols_.size());
 718  E :    DCHECK(abs_reference != NULL);
 719    :  
 720  E :    size_t symbol_iat_index = symbols_[index].iat_index;
 721    :    if (import_descriptor_.block() == NULL ||
 722  E :        symbol_iat_index == kInvalidIatIndex) {
 723  i :      LOG(ERROR) << "Called GetAbsReference on an uninitialized ImportedSymbol.";
 724  i :      return false;
 725    :    }
 726    :  
 727  E :    ImageThunkData32 thunks;
 728    :    if (!import_descriptor_.Dereference(import_descriptor_->FirstThunk,
 729  E :                                        &thunks)) {
 730  i :      LOG(ERROR) << "Unable to dereference IMAGE_THUNK_DATA32.";
 731  i :      return false;
 732    :    }
 733    :  
 734  E :    if (symbol_iat_index >= thunks.ElementCount()) {
 735  i :      LOG(ERROR) << "Invalid symbol_index for IAT.";
 736  i :      return false;
 737    :    }
 738    :  
 739  E :    Offset offset = thunks.OffsetOf(thunks[symbol_iat_index].u1.AddressOfData);
 740    :    *abs_reference = BlockGraph::Reference(
 741  E :        BlockGraph::ABSOLUTE_REF, kPtrSize, thunks.block(), offset, offset);
 742    :  
 743  E :    return true;
 744  E :  }
 745    :  
 746    :  }  // namespace transforms
 747    :  }  // namespace pe

Coverage information generated Thu Jul 04 09:34:53 2013.