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

CoverageLines executed / instrumented / missingexe / inst / missLanguageGroup
81.6%3053740.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 PEAddImportsTransform 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/pe_add_imports_transform.h"
 110    :  
 111    :  #include "base/strings/string_piece.h"
 112    :  #include "base/strings/string_util.h"
 113    :  #include "base/strings/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_DELAYLOAD_DESCRIPTOR> ImageDelayLoadDescriptor;
 133    :  typedef TypedBlock<IMAGE_DOS_HEADER> DosHeader;
 134    :  typedef TypedBlock<IMAGE_IMPORT_BY_NAME> ImageImportByName;
 135    :  typedef TypedBlock<IMAGE_IMPORT_DESCRIPTOR> ImageImportDescriptor;
 136    :  typedef TypedBlock<IMAGE_NT_HEADERS> NtHeaders;
 137    :  typedef TypedBlock<IMAGE_THUNK_DATA32> ImageThunkData32;
 138    :  typedef TypedBlock<StringStruct> String;
 139    :  
 140    :  namespace {
 141    :  
 142    :  const size_t kPtrSize = sizeof(core::RelativeAddress);
 143    :  const size_t kInvalidIndex = static_cast<size_t>(-1);
 144    :  
 145    :  // Looks up the given data directory and checks that it points to valid data.
 146    :  // If it doesn't exist and find_only is false, it will allocate a block with
 147    :  // the given name and size.
 148    :  bool FindOrAddDataDirectory(bool find_only,
 149    :                              size_t directory_index,
 150    :                              const base::StringPiece& block_name,
 151    :                              size_t block_size,
 152    :                              BlockGraph* block_graph,
 153    :                              BlockGraph::Block* nt_headers_block,
 154  E :                              BlockGraph::Block** directory_block) {
 155    :    DCHECK_LT(directory_index,
 156  E :              static_cast<size_t>(IMAGE_NUMBEROF_DIRECTORY_ENTRIES));
 157  E :    DCHECK_GT(block_size, 0u);
 158  E :    DCHECK_NE(reinterpret_cast<BlockGraph*>(NULL), block_graph);
 159  E :    DCHECK_NE(reinterpret_cast<BlockGraph::Block*>(NULL), nt_headers_block);
 160  E :    DCHECK_NE(reinterpret_cast<BlockGraph::Block**>(NULL), directory_block);
 161    :  
 162  E :    *directory_block = NULL;
 163    :  
 164  E :    NtHeaders nt_headers;
 165  E :    if (!nt_headers.Init(0, nt_headers_block)) {
 166  i :      LOG(ERROR) << "Unable to cast NT headers.";
 167  i :      return false;
 168    :    }
 169    :  
 170    :    IMAGE_DATA_DIRECTORY* data_directory =
 171  E :        nt_headers->OptionalHeader.DataDirectory + directory_index;
 172    :  
 173    :    BlockGraph::Offset offset = nt_headers.OffsetOf(
 174  E :        data_directory->VirtualAddress);
 175  E :    BlockGraph::Reference ref;
 176    :  
 177    :    // No entry? Then make a zero initialized block that is stored in .rdata,
 178    :    // where all of these structures live.
 179  E :    if (!nt_headers_block->GetReference(offset, &ref)) {
 180    :      // We don't need to create the entry if we're exploring only.
 181  E :      if (find_only)
 182  E :        return true;
 183    :  
 184    :      BlockGraph::Section* section = block_graph->FindOrAddSection(
 185  E :          kReadOnlyDataSectionName, kReadOnlyDataCharacteristics);
 186  E :      DCHECK(section != NULL);
 187    :  
 188    :      BlockGraph::Block* block = block_graph->AddBlock(
 189  E :          BlockGraph::DATA_BLOCK, block_size, block_name);
 190  E :      DCHECK(block != NULL);
 191  E :      block->set_section(section->id());
 192  E :      block->set_attribute(BlockGraph::PE_PARSED);
 193    :  
 194    :      // We need to actually allocate the data so that future TypedBlock
 195    :      // dereferences will work.
 196  E :      if (block->AllocateData(block_size) == NULL) {
 197  i :        LOG(ERROR) << "Failed to allocate block data.";
 198  i :        return false;
 199    :      }
 200    :  
 201    :      // Hook it up to the NT header.
 202    :      nt_headers.SetReference(BlockGraph::RELATIVE_REF,
 203    :                              data_directory->VirtualAddress,
 204    :                              block,
 205  E :                              0, 0);
 206  E :      data_directory->Size = block_size;
 207    :  
 208  E :      *directory_block = block;
 209  E :    } else {
 210    :      // If the directory already exists, return it.
 211  E :      if (ref.offset() != 0) {
 212  i :        LOG(ERROR) << "Existing \"" << block_name << "\" directory is not its "
 213    :                   << "own block.";
 214  i :        return false;
 215    :      }
 216  E :      *directory_block = ref.referenced();
 217    :    }
 218    :  
 219  E :    return true;
 220  E :  }
 221    :  
 222    :  bool ModuleNameMatches(const base::StringPiece& module_name,
 223  E :                         const String& dll_name) {
 224  E :    size_t max_len = dll_name.ElementCount();
 225  E :    if (max_len < module_name.size())
 226  E :      return false;
 227  E :    return base::strncasecmp(dll_name->string, module_name.data(), max_len) == 0;
 228  E :  }
 229    :  
 230    :  bool SymbolNameMatches(const base::StringPiece& symbol_name,
 231  E :                         const ImageImportByName& iibn) {
 232    :    size_t max_len = iibn.block()->data_size() - iibn.offset() -
 233  E :        offsetof(IMAGE_IMPORT_BY_NAME, Name);
 234  E :    if (max_len < symbol_name.size())
 235  E :      return false;
 236  E :    return ::strncmp(iibn->Name, symbol_name.data(), max_len) == 0;
 237  E :  }
 238    :  
 239    :  // Finds or creates an Image Import Descriptor block for the given library.
 240    :  // Returns true on success, false otherwise.
 241    :  bool FindOrAddImageImportDescriptor(bool find_only,
 242    :                                      const char* module_name,
 243    :                                      BlockGraph* block_graph,
 244    :                                      BlockGraph::Block* iida_block,
 245    :                                      BlockGraph::Block* iat_block,
 246    :                                      ImageImportDescriptor* iid,
 247    :                                      bool* added,
 248  E :                                      bool* exists) {
 249  E :    DCHECK(module_name != NULL);
 250  E :    DCHECK(block_graph != NULL);
 251  E :    DCHECK(iida_block != NULL);
 252  E :    DCHECK(iat_block != NULL);
 253  E :    DCHECK(iid != NULL);
 254  E :    DCHECK(added != NULL);
 255  E :    DCHECK(exists != NULL);
 256    :  
 257  E :    *added = false;
 258  E :    *exists = false;
 259    :  
 260  E :    ImageImportDescriptor iida;
 261  E :    if (!iida.Init(0, iida_block)) {
 262  i :      LOG(ERROR) << "Unable to cast Image Import Descriptor.";
 263  i :      return false;
 264    :    }
 265    :  
 266    :    // The array is NULL terminated with a potentially incomplete descriptor so
 267    :    // we can't use ElementCount - 1.
 268  E :    DCHECK_GT(iida_block->size(), 0U);
 269    :    size_t descriptor_count =
 270    :        (common::AlignUp(iida_block->size(), sizeof(IMAGE_IMPORT_DESCRIPTOR)) /
 271  E :            sizeof(IMAGE_IMPORT_DESCRIPTOR)) - 1;
 272    :  
 273  E :    for (size_t iida_index = 0; iida_index < descriptor_count; ++iida_index) {
 274  E :      String dll_name;
 275  E :      if (!iida.Dereference(iida[iida_index].Name, &dll_name)) {
 276  i :        LOG(ERROR) << "Unable to dereference DLL name.";
 277  i :        return false;
 278    :      }
 279    :  
 280  E :      if (ModuleNameMatches(module_name, dll_name)) {
 281    :        // This should never fail, but we sanity check it nonetheless.
 282  E :        bool result = iid->Init(iida.OffsetOf(iida[iida_index]), iida.block());
 283  E :        DCHECK(result);
 284  E :        *exists = true;
 285  E :        return true;
 286    :      }
 287  E :    }
 288    :  
 289    :    // If we get here then the entry doesn't exist. If we've been asked to only
 290    :    // search for it then we can return early.
 291  E :    if (find_only)
 292  E :      return true;
 293    :  
 294    :    // Create room for the new descriptor, which we'll tack on to the end of the
 295    :    // array, but before the NULL terminator. We use 'InsertData' so that all
 296    :    // labels are patched up.
 297  E :    Offset new_iid_offset = descriptor_count * sizeof(IMAGE_IMPORT_DESCRIPTOR);
 298    :    iida_block->InsertData(
 299  E :        new_iid_offset, sizeof(IMAGE_IMPORT_DESCRIPTOR), true);
 300    :    iida_block->SetLabel(
 301    :        new_iid_offset,
 302    :        base::StringPrintf("Image Import Descriptor: %s", module_name),
 303  E :        BlockGraph::DATA_LABEL);
 304    :  
 305    :    // We expect the new entry to be dereferencable using iida[descriptor_count].
 306  E :    DCHECK_GT(iida.ElementCount(), descriptor_count);
 307    :  
 308    :    // Create the various child structures that will be pointed to by the
 309    :    // import descriptor. The INT block and the IAT block are NULL terminated
 310    :    // lists of pointers, and the terminating NULL is allocated. We don't yet
 311    :    // allocate a block to hold the import names, deferring that for later.
 312  E :    BlockGraph::SectionId iida_section_id = iida_block->section();
 313  E :    size_t name_len = strlen(module_name);
 314    :    BlockGraph::Block* int_block = block_graph->AddBlock(
 315    :        BlockGraph::DATA_BLOCK, kPtrSize,
 316  E :        base::StringPrintf("Import Name Table: %s", module_name));
 317    :    BlockGraph::Block* dll_name_block = block_graph->AddBlock(
 318    :        BlockGraph::DATA_BLOCK, name_len + 1,
 319  E :        base::StringPrintf("Import Name: %s", module_name));
 320  E :    if (int_block == NULL || dll_name_block == NULL) {
 321  i :      LOG(ERROR) << "Unable to create blocks for Image Import Descriptor.";
 322  i :      return false;
 323    :    }
 324    :  
 325    :    // NOTE: If PEParser was modified to parse a single INT block, we could be
 326    :    //     extending/reusing it rather than creating a new INT per module.
 327  E :    int_block->set_section(iida_section_id);
 328  E :    int_block->set_attribute(BlockGraph::PE_PARSED);
 329    :    int_block->SetLabel(
 330    :        0,
 331    :        base::StringPrintf("%s INT: NULL entry", module_name),
 332  E :        BlockGraph::DATA_LABEL);
 333  E :    if (int_block->AllocateData(kPtrSize) == NULL) {
 334  i :      LOG(ERROR) << "Failed to allocate block data.";
 335  i :      return false;
 336    :    }
 337    :  
 338    :    // We use the DLL name block and extend it. This keeps things well ordered
 339    :    // when writing back the image using a canonical ordering.
 340  E :    dll_name_block->set_section(iida_section_id);
 341  E :    dll_name_block->set_attribute(BlockGraph::PE_PARSED);
 342  E :    if (dll_name_block->CopyData(name_len + 1, module_name) == NULL) {
 343  i :      LOG(ERROR) << "Failed to copy block data.";
 344  i :      return false;
 345    :    }
 346    :  
 347    :    // Add another NULL entry to the IAT block, but only if it does not already
 348    :    // consist of a single NULL entry (meaning it was just created). We are purely
 349    :    // extending this block, so no need to use the data insertion functions.
 350  E :    Offset iat_offset = 0;
 351  E :    if (iat_block->size() != kPtrSize) {
 352  E :      iat_offset = iat_block->size();
 353  E :      size_t iat_size = iat_offset + kPtrSize;
 354  E :      iat_block->set_size(iat_size);
 355  E :      iat_block->ResizeData(iat_size);
 356  E :      DCHECK_EQ(iat_size, iat_block->size());
 357  E :      DCHECK_EQ(iat_size, iat_block->data_size());
 358    :    }
 359    :  
 360    :    // Add a label for debugging purposes.
 361    :    iat_block->SetLabel(iat_offset,
 362    :                        base::StringPrintf("%s: NULL thunk", module_name),
 363  E :                        BlockGraph::DATA_LABEL);
 364    :  
 365    :    // Hook up these blocks.
 366    :    iida.SetReference(BlockGraph::RELATIVE_REF,
 367  E :                      iida[descriptor_count].OriginalFirstThunk, int_block, 0, 0);
 368    :    iida.SetReference(BlockGraph::RELATIVE_REF,
 369    :                      iida[descriptor_count].FirstThunk, iat_block, iat_offset,
 370  E :                      iat_offset);
 371    :    iida.SetReference(BlockGraph::RELATIVE_REF,
 372  E :                      iida[descriptor_count].Name, dll_name_block, 0, 0);
 373    :  
 374    :    // Finally, return the descriptor.
 375  E :    if (!iid->Init(new_iid_offset, iida_block)) {
 376  i :      LOG(ERROR) << "Unable to cast Image Import Descriptor.";
 377  i :      return false;
 378    :    }
 379    :  
 380  E :    *added = true;
 381  E :    *exists = true;
 382    :  
 383  E :    return true;
 384  E :  }
 385    :  
 386    :  // Searches for the delay-load library with the given module name. Returns true
 387    :  // on success, false otherwise. If found, returns the index. If not found
 388    :  // sets the index to kInvalidIndex.
 389    :  bool FindDelayLoadImportDescriptor(const base::StringPiece& module_name,
 390    :                                     const ImageDelayLoadDescriptor& idld,
 391  E :                                     size_t* index) {
 392  E :    DCHECK_NE(reinterpret_cast<size_t*>(NULL), index);
 393    :  
 394  E :    *index = kInvalidIndex;
 395    :  
 396  E :    for (size_t i = 0; i < idld.ElementCount(); ++i) {
 397  E :      bool zero_data = idld[i].DllNameRVA == 0;
 398  E :      bool has_ref = idld.HasReference(idld[i].DllNameRVA);
 399    :  
 400    :      // Keep an eye out for null termination of the array.
 401  E :      if (zero_data && !has_ref)
 402  E :        return true;
 403    :  
 404    :      // If the data is not zero then we expect there to be a reference.
 405  E :      if (!zero_data && !has_ref) {
 406  i :        LOG(ERROR) << "Expected DllNameRVA reference at index " << i
 407    :                   << " of IMAGE_DELAYLOAD_DESCRIPTOR array.";
 408  i :        return false;
 409    :      }
 410    :  
 411  E :      String dll_name;
 412  E :      if (!idld.Dereference(idld[i].DllNameRVA, &dll_name)) {
 413  i :        LOG(ERROR) << "Failed to dereference DllNameRVA at index " << i
 414    :                   << " of IMAGE_DELAYLOAD_DESCRIPTOR array.";
 415  i :        return false;
 416    :      }
 417    :  
 418  E :      if (ModuleNameMatches(module_name, dll_name)) {
 419  E :        *index = i;
 420  E :        return true;
 421    :      }
 422  E :    }
 423    :  
 424  i :    return true;
 425  E :  }
 426    :  
 427    :  // Finds or adds an imported symbol to the given module (represented by its
 428    :  // import descriptor). Returns true on success, false otherwise. On success
 429    :  // returns a reference to the module's IAT entry. New entries are always added
 430    :  // to the end of the table so as not to invalidate any other unlinked references
 431    :  // (not part of the BlockGraph, so unable to be patched up) into the table.
 432    :  bool FindOrAddImportedSymbol(bool find_only,
 433    :                               const char* symbol_name,
 434    :                               const ImageImportDescriptor& iid,
 435    :                               BlockGraph* block_graph,
 436    :                               BlockGraph::Block* iat_block,
 437    :                               size_t* iat_index,
 438  E :                               bool* added) {
 439  E :    DCHECK(symbol_name != NULL);
 440  E :    DCHECK(block_graph != NULL);
 441  E :    DCHECK(iat_block != NULL);
 442  E :    DCHECK(iat_index != NULL);
 443  E :    DCHECK(added != NULL);
 444    :  
 445  E :    *iat_index = kInvalidIndex;
 446  E :    *added = false;
 447    :  
 448  E :    TypedBlock<IMAGE_IMPORT_BY_NAME*> hna, iat;
 449    :    if (!iid.Dereference(iid->OriginalFirstThunk, &hna) ||
 450  E :        !iid.Dereference(iid->FirstThunk, &iat)) {
 451  i :      LOG(ERROR) << "Unable to dereference OriginalFirstThunk/FirstThunk.";
 452  i :      return false;
 453    :    }
 454    :  
 455    :    // Loop through the existing imports and see if we can't find a match. If so,
 456    :    // we don't need to import the symbol as it is already imported. The array is
 457    :    // NULL terminated so we loop through all elements except for the last one.
 458  E :    size_t i = 0;
 459  E :    for (; i < hna.ElementCount() && i < iat.ElementCount(); ++i) {
 460  E :      ConstTypedBlock<IMAGE_THUNK_DATA32> thunk;
 461  E :      if (!thunk.Init(hna.OffsetOf(hna[i]), hna.block())) {
 462  i :        LOG(ERROR) << "Unable to dereference IMAGE_THUNK_DATA32.";
 463  i :        return false;
 464    :      }
 465    :  
 466    :      // Is this an ordinal import? Skip it, as we have no way of
 467    :      // knowing the actual name of the symbol.
 468  E :      if (IMAGE_SNAP_BY_ORDINAL(thunk->u1.Ordinal))
 469  E :        continue;
 470    :  
 471    :      // Have no reference? Then terminate the iteration.
 472  E :      if (!thunk.HasReference(thunk->u1.AddressOfData)) {
 473    :        // We sanity check that the actual data is null.
 474  E :        DCHECK_EQ(0u, thunk->u1.AddressOfData);
 475  E :        break;
 476    :      }
 477    :  
 478    :      // Otherwise this should point to an IMAGE_IMPORT_BY_NAME structure.
 479  E :      ImageImportByName iibn;
 480  E :      if (!hna.Dereference(hna[i], &iibn)) {
 481  i :        LOG(ERROR) << "Unable to dereference IMAGE_IMPORT_BY_NAME.";
 482  i :        return false;
 483    :      }
 484    :  
 485    :      // Check to see if this symbol matches that of the current image import
 486    :      // by name.
 487  E :      if (SymbolNameMatches(symbol_name, iibn)) {
 488  E :        *iat_index = i;
 489  E :        return true;
 490    :      }
 491  E :    }
 492    :  
 493    :    // If we get here then the entry doesn't exist. If we've been asked to only
 494    :    // search for it then we can return early.
 495  E :    if (find_only)
 496  E :      return true;
 497    :  
 498    :    // Figure out how large the data needs to be to hold the name of this exported
 499    :    // symbol.  The IMAGE_IMPORT_BY_NAME struct has a WORD ordinal and a variable
 500    :    // sized field for the null-terminated function name. Each entry should be
 501    :    // WORD aligned, and will be referenced from the import address table and the
 502    :    // import name table.
 503  E :    size_t symbol_name_len = strlen(symbol_name);
 504    :    size_t iibn_size = sizeof(WORD) + common::AlignUp(symbol_name_len + 1,
 505  E :                                                      sizeof(WORD));
 506    :  
 507    :    // Get the DLL name. We will be inserting the IIBN entry to the block
 508    :    // containing it immediately prior to the DLL name.
 509  E :    String dll_name;
 510  E :    if (!iid.Dereference(iid->Name, &dll_name)) {
 511  i :      LOG(ERROR) << "Unable to dereference DLL name.";
 512  i :      return false;
 513    :    }
 514  E :    Offset iibn_offset = dll_name.offset();
 515  E :    dll_name.block()->InsertData(iibn_offset, iibn_size, true);
 516    :  
 517    :    // Populate the import struct.
 518  E :    TypedBlock<IMAGE_IMPORT_BY_NAME> iibn;
 519  E :    if (!iibn.InitWithSize(iibn_offset, iibn_size, dll_name.block())) {
 520  i :      LOG(ERROR) << "Unable to dereference new IMAGE_IMPORT_BY_NAME.";
 521  i :      return false;
 522    :    }
 523  E :    iibn->Hint = 0;
 524    :    base::strlcpy(reinterpret_cast<char*>(iibn->Name), symbol_name,
 525  E :                  symbol_name_len + 1);
 526    :  
 527    :    // Make room in the INT and the IAT for the new symbol. We place it
 528    :    // after the last entry for this module.
 529  E :    Offset int_offset = hna.OffsetOf(hna[i]);
 530  E :    Offset iat_offset = iat.OffsetOf(iat[i]);
 531    :    // We're pointed at the terminating zero. The position we're pointing at can
 532    :    // be the destination for references (in the normal case where someone is
 533    :    // using the import). However, in the special case where the IAT and the INT
 534    :    // are empty, our slot may also be pointed at by the import descriptor.
 535    :    // If we were to insert data at this position, we'd push the import
 536    :    // descriptor's pointer forward, past our new entry. To avoid this, we insert
 537    :    // the new data after the terminating zero we're pointing at, then usurp the
 538    :    // previously terminating zero for our entry.
 539  E :    hna.block()->InsertData(int_offset + kPtrSize, kPtrSize, true);
 540  E :    iat.block()->InsertData(iat_offset + kPtrSize, kPtrSize, true);
 541    :  
 542    :    // Because of the usurping mentioned above, we manually move any existing
 543    :    // labels.
 544  E :    BlockGraph::Label label;
 545  E :    if (hna.block()->GetLabel(int_offset, &label)) {
 546  E :      hna.block()->RemoveLabel(int_offset);
 547  E :      hna.block()->SetLabel(int_offset + kPtrSize, label);
 548    :    }
 549  E :    if (iat.block()->GetLabel(iat_offset, &label)) {
 550  E :      iat.block()->RemoveLabel(iat_offset);
 551  E :      iat.block()->SetLabel(iat_offset + kPtrSize, label);
 552    :    }
 553    :  
 554    :    // Add the new labels. We have to get the module_name at this point
 555    :    // because it may have been moved with our insertions above.
 556  E :    String module_name;
 557  E :    if (!iid.Dereference(iid->Name, &module_name)) {
 558  i :      LOG(ERROR) << "Unable to dereference import name.";
 559  i :      return false;
 560    :    }
 561    :    hna.block()->SetLabel(
 562    :        int_offset,
 563    :        base::StringPrintf("%s INT: %s", module_name->string, symbol_name),
 564  E :        BlockGraph::DATA_LABEL);
 565    :    iat.block()->SetLabel(
 566    :        iat_offset,
 567    :        base::StringPrintf("%s IAT: %s", module_name->string, symbol_name),
 568  E :        BlockGraph::DATA_LABEL);
 569    :  
 570    :    // Hook up the newly created IMAGE_IMPORT_BY_NAME to both tables.
 571    :    BlockGraph::Reference iibn_ref(BlockGraph::RELATIVE_REF,
 572    :                                   kPtrSize,
 573    :                                   iibn.block(),
 574    :                                   iibn.offset(),
 575  E :                                   iibn.offset());
 576  E :    hna.block()->SetReference(int_offset, iibn_ref);
 577  E :    iat.block()->SetReference(iat_offset, iibn_ref);
 578    :  
 579    :    // Return the reference to the IAT entry for the newly imported symbol.
 580  E :    *iat_index = i;
 581  E :    *added = true;
 582    :  
 583  E :    return true;
 584  E :  }
 585    :  
 586    :  // Looks for the given symbol in the given delay-loaded library descriptor.
 587    :  // Returns true on success, false otherwise. If the symbol was found sets
 588    :  // |found| to true, and return a reference to it via |ref|.
 589    :  bool FindDelayLoadSymbol(const base::StringPiece& symbol_name,
 590    :                           const ImageDelayLoadDescriptor& idld,
 591    :                           size_t module_index,
 592    :                           bool* found,
 593    :                           size_t* index,
 594  E :                           BlockGraph::Reference* ref) {
 595  E :    DCHECK_NE(reinterpret_cast<bool*>(NULL), found);
 596  E :    DCHECK_NE(reinterpret_cast<size_t*>(NULL), index);
 597  E :    DCHECK_NE(reinterpret_cast<BlockGraph::Reference*>(NULL), ref);
 598    :  
 599  E :    *found = false;
 600  E :    *index = kInvalidIndex;
 601    :  
 602  E :    ImageThunkData32 addresses;
 603  E :    ImageThunkData32 names;
 604    :    if (!idld.Dereference(idld[module_index].ImportAddressTableRVA, &addresses) ||
 605  E :        !idld.Dereference(idld[module_index].ImportNameTableRVA, &names)) {
 606  i :      LOG(ERROR) << "Failed to dereference IAT/INT for delay-load library.";
 607  i :      return false;
 608    :    }
 609    :  
 610  E :    size_t count = std::min(addresses.ElementCount(), names.ElementCount());
 611  E :    for (size_t i = 0; i < count; ++i) {
 612    :      // Keep an eye out for zero-terminating IAT entries.
 613  E :      bool zero_data = addresses[i].u1.AddressOfData == 0;
 614  E :      bool has_ref = addresses.HasReference(addresses[i].u1.AddressOfData);
 615  E :      if (zero_data && !has_ref)
 616  i :        break;
 617  E :      if (!zero_data && !has_ref) {
 618  i :        LOG(ERROR) << "Expected reference at offset " << i
 619    :                   << " of delay-load IAT.";
 620  i :        return false;
 621    :      }
 622    :  
 623    :      // Keep an eye out for zero-terminating INT entries.
 624  E :      zero_data = names[i].u1.AddressOfData == 0;
 625  E :      has_ref = names.HasReference(names[i].u1.AddressOfData);
 626  E :      if (zero_data && !has_ref)
 627  i :        break;
 628  E :      if (!zero_data && !has_ref) {
 629  i :        LOG(ERROR) << "Expected reference at offset " << i
 630    :                   << " of delay-load INT.";
 631  i :        return false;
 632    :      }
 633    :  
 634  E :      ImageImportByName iibn;
 635  E :      if (!names.Dereference(names[i].u1.AddressOfData, &iibn)) {
 636  i :        LOG(ERROR) << "Failed to dereference name of entry " << i
 637    :                   << " of delay-load INT.";
 638  i :        return false;
 639    :      }
 640    :  
 641  E :      if (SymbolNameMatches(symbol_name, iibn)) {
 642  E :        Offset offset = addresses.OffsetOf(addresses->u1.Function);
 643    :        *ref = BlockGraph::Reference(BlockGraph::ABSOLUTE_REF,
 644    :                                     BlockGraph::Reference::kMaximumSize,
 645    :                                     addresses.block(),
 646    :                                     offset,
 647  E :                                     offset);
 648  E :        *found = true;
 649  E :        *index = i;
 650  E :        return true;
 651    :      }
 652  i :    }
 653    :  
 654  i :    return true;
 655  E :  }
 656    :  
 657    :  }  // namespace
 658    :  
 659    :  const char PEAddImportsTransform::kTransformName[] = "PEAddImportsTransform";
 660    :  
 661    :  PEAddImportsTransform::PEAddImportsTransform()
 662    :      : image_import_descriptor_block_(NULL),
 663  E :        import_address_table_block_(NULL) {
 664  E :  }
 665    :  
 666    :  bool PEAddImportsTransform::TransformBlockGraph(
 667    :      const TransformPolicyInterface* policy,
 668    :      BlockGraph* block_graph,
 669  E :      BlockGraph::Block* dos_header_block) {
 670  E :    DCHECK_NE(reinterpret_cast<TransformPolicyInterface*>(NULL), policy);
 671  E :    DCHECK_NE(reinterpret_cast<BlockGraph*>(NULL), block_graph);
 672  E :    DCHECK_NE(reinterpret_cast<BlockGraph::Block*>(NULL), dos_header_block);
 673  E :    DCHECK_EQ(BlockGraph::PE_IMAGE, block_graph->image_format());
 674    :  
 675  E :    modules_added_ = 0;
 676  E :    symbols_added_ = 0;
 677    :  
 678  E :    DosHeader dos_header;
 679  E :    NtHeaders nt_headers;
 680    :    if (!dos_header.Init(0, dos_header_block) ||
 681  E :        !dos_header.Dereference(dos_header->e_lfanew, &nt_headers)) {
 682  i :      LOG(ERROR) << "Unable to cast image headers.";
 683  i :      return false;
 684    :    }
 685    :  
 686    :    // Find delay load imports. This is read-only, searching for existing
 687    :    // imports but not injecting new ones.
 688  E :    if (!FindDelayLoadImports(block_graph, nt_headers.block()))
 689  i :      return false;
 690    :  
 691    :    // Before processing regular imports, let's determine if we're on a strictly
 692    :    // exploratory mission. We don't want to add anything if all unresolved
 693    :    // modules/symbols are 'find only'.
 694  E :    bool find_only = true;
 695  E :    for (size_t i = 0; i < imported_modules_.size(); ++i) {
 696  E :      for (size_t j = 0; j < imported_modules_[i]->size(); ++j) {
 697    :        // If the symbol is resolved, we don't care about it. We don't want to
 698    :        // unnecessarily add PE import structures if we're not creating any
 699    :        // imports.
 700  E :        if (imported_modules_[i]->SymbolIsImported(j))
 701  E :          continue;
 702  E :        if (imported_modules_[i]->GetSymbolMode(j) != ImportedModule::kFindOnly) {
 703  E :          find_only = false;
 704  E :          break;
 705    :        }
 706  E :      }
 707  E :    }
 708    :  
 709    :    // Find normal imports. If the symbol is imported as both a delay-load and
 710    :    // a regular import, then this will overwrite it. Thus, regular imports will
 711    :    // be preferred. However, if the symbol was resolved as a delay-load import
 712    :    // then this will not cause it to also be added as a regular import.
 713  E :    if (!FindOrAddImports(find_only, block_graph, nt_headers.block()))
 714  i :      return false;
 715    :  
 716  E :    return true;
 717  E :  }
 718    :  
 719    :  bool PEAddImportsTransform::FindOrAddImports(
 720    :      bool find_only,
 721    :      BlockGraph* block_graph,
 722  E :      BlockGraph::Block* nt_headers_block) {
 723  E :    DCHECK_NE(reinterpret_cast<BlockGraph*>(NULL), block_graph);
 724  E :    DCHECK_NE(reinterpret_cast<BlockGraph::Block*>(NULL), nt_headers_block);
 725    :  
 726  E :    NtHeaders nt_headers;
 727  E :    CHECK(nt_headers.Init(0, nt_headers_block));
 728    :  
 729    :    // Get the import data directory.
 730  E :    image_import_descriptor_block_ = NULL;
 731    :    if (!FindOrAddDataDirectory(find_only,
 732    :                                IMAGE_DIRECTORY_ENTRY_IMPORT,
 733    :                                "Image Import Descriptor Array",
 734    :                                sizeof(IMAGE_IMPORT_DESCRIPTOR),
 735    :                                block_graph,
 736    :                                nt_headers.block(),
 737  E :                                &image_import_descriptor_block_)) {
 738  i :      return false;
 739    :    }
 740  E :    if (image_import_descriptor_block_ == NULL)
 741  i :      return find_only;
 742    :  
 743    :    // Similarly, get the import address table.
 744  E :    import_address_table_block_ = NULL;
 745    :    if (!FindOrAddDataDirectory(find_only,
 746    :                                IMAGE_DIRECTORY_ENTRY_IAT,
 747    :                                "Import Address Table",
 748    :                                kPtrSize,
 749    :                                block_graph,
 750    :                                nt_headers.block(),
 751  E :                                &import_address_table_block_)) {
 752  i :      return false;
 753    :    }
 754  E :    if (import_address_table_block_ == NULL)
 755  i :      return find_only;
 756    :  
 757    :    // Handle each library individually.
 758  E :    for (size_t i = 0; i < imported_modules_.size(); ++i) {
 759  E :      ImportedModule* module = imported_modules_[i];
 760    :  
 761    :      // First find or create an entry for this module in the Image Import
 762    :      // Descriptor Array.
 763  E :      ImageImportDescriptor iid;
 764  E :      bool module_added = false;
 765  E :      bool module_exists = false;
 766    :      if (!FindOrAddImageImportDescriptor(
 767    :              module->mode() == ImportedModule::kFindOnly,
 768    :              module->name().c_str(),
 769    :              block_graph,
 770    :              image_import_descriptor_block_,
 771    :              import_address_table_block_,
 772    :              &iid,
 773    :              &module_added,
 774  E :              &module_exists)) {
 775  i :        LOG(ERROR) << "Failed to find or import module.";
 776  i :        return false;
 777    :      }
 778    :  
 779    :      // If we're fact finding only and the module does not exist then we don't
 780    :      // need to look up its symbols.
 781  E :      if (module->mode() == ImportedModule::kFindOnly && !module_exists) {
 782  E :        DCHECK(!module_added);
 783  E :        continue;
 784    :      }
 785    :  
 786  E :      DCHECK(module_exists);
 787  E :      UpdateModule(true, module_added, module);
 788  E :      modules_added_ += module_added;
 789    :  
 790    :      // Update the version date/time stamp if requested.
 791  E :      if (module->date() != ImportedModule::kInvalidDate)
 792  E :        iid->TimeDateStamp = module->date();
 793    :  
 794    :      // Get a pointer to the import thunks.
 795  E :      ImageThunkData32 thunks;
 796  E :      if (!iid.Dereference(iid->FirstThunk, &thunks)) {
 797  i :        LOG(ERROR) << "Unable to dereference IMAGE_THUNK_DATA32.";
 798  i :        return false;
 799    :      }
 800    :  
 801  E :      for (size_t j = 0; j < module->size(); ++j) {
 802    :        bool symbol_find_only =
 803  E :            module->GetSymbolMode(j) == ImportedModule::kFindOnly;
 804    :  
 805    :        // If the symbol was already resolved as a delay-load import, then
 806    :        // don't allow it to also be added as a normal import.
 807  E :        if (module->SymbolIsImported(j))
 808  i :          symbol_find_only = true;
 809    :  
 810    :        // Now, for each symbol get the offset of the IAT entry. This will create
 811    :        // the entry (and all accompanying structures) if necessary.
 812  E :        size_t symbol_iat_index = kInvalidIndex;
 813  E :        bool symbol_added = false;
 814    :        if (!FindOrAddImportedSymbol(
 815    :                symbol_find_only,
 816    :                module->GetSymbolName(j).c_str(),
 817    :                iid,
 818    :                block_graph,
 819    :                import_address_table_block_,
 820    :                &symbol_iat_index,
 821  E :                &symbol_added)) {
 822  i :          LOG(ERROR) << "Failed to find or import symbol.";
 823  i :          return false;
 824    :        }
 825  E :        symbols_added_ += symbol_added;
 826    :  
 827  E :        if (symbol_iat_index != kInvalidIndex) {
 828    :          Offset offset =
 829  E :              thunks.OffsetOf(thunks[symbol_iat_index].u1.AddressOfData);
 830    :          BlockGraph::Reference ref(BlockGraph::ABSOLUTE_REF, kPtrSize,
 831  E :                                    thunks.block(), offset, offset);
 832    :  
 833  E :          UpdateModuleSymbolInfo(j, true, symbol_added, module);
 834  E :          UpdateModuleSymbolReference(j, ref, true, module);
 835    :        }
 836  E :      }
 837  E :    }
 838    :  
 839    :    // Update the data directory sizes.
 840    :    nt_headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size =
 841  E :        image_import_descriptor_block_->size();
 842    :    nt_headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IAT].Size =
 843  E :        import_address_table_block_->size();
 844    :  
 845  E :    return true;
 846  E :  }
 847    :  
 848    :  bool PEAddImportsTransform::FindDelayLoadImports(
 849  E :      BlockGraph* block_graph, BlockGraph::Block* nt_headers_block) {
 850  E :    DCHECK_NE(reinterpret_cast<BlockGraph*>(NULL), block_graph);
 851  E :    DCHECK_NE(reinterpret_cast<BlockGraph::Block*>(NULL), nt_headers_block);
 852    :  
 853  E :    NtHeaders nt_headers;
 854  E :    CHECK(nt_headers.Init(0, nt_headers_block));
 855    :  
 856    :    // Get the delay-load import data directory.
 857  E :    image_delayload_descriptor_block_ = NULL;
 858    :    if (!FindOrAddDataDirectory(true,
 859    :                                IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT,
 860    :                                "Image Delay Load Descriptor Array",
 861    :                                sizeof(IMAGE_DELAYLOAD_DESCRIPTOR),
 862    :                                block_graph,
 863    :                                nt_headers.block(),
 864  E :                                &image_delayload_descriptor_block_)) {
 865  i :      return false;
 866    :    }
 867  E :    if (image_delayload_descriptor_block_ == NULL)
 868  E :      return true;
 869    :  
 870  E :    ImageDelayLoadDescriptor idld;
 871  E :    if (!idld.Init(0, image_delayload_descriptor_block_)) {
 872  i :      LOG(ERROR) << "Unable to cast IMAGE_DELAYLOAD_DESCRIPTOR.";
 873  i :      return false;
 874    :    }
 875    :  
 876  E :    for (size_t i = 0; i < imported_modules_.size(); ++i) {
 877  E :      ImportedModule* module = imported_modules_[i];
 878    :  
 879    :       // Look for a descriptor corresponding to this module.
 880  E :      size_t module_index = kInvalidIndex;
 881  E :      if (!FindDelayLoadImportDescriptor(module->name(), idld, &module_index))
 882  i :        return false;
 883  E :      if (module_index == kInvalidIndex)
 884  E :        continue;
 885    :  
 886  E :      UpdateModule(true, false, module);
 887    :  
 888    :      // Iterate over the symbols.
 889  E :      for (size_t j = 0; j < module->size(); ++j) {
 890    :        // Don't process symbols that are already imported.
 891  E :        if (module->SymbolIsImported(j))
 892  i :          continue;
 893    :  
 894    :        // Look for a matching symbol.
 895  E :        bool found = false;
 896  E :        size_t index = kInvalidIndex;
 897  E :        BlockGraph::Reference ref;
 898    :        if (!FindDelayLoadSymbol(module->GetSymbolName(j), idld, module_index,
 899  E :                                 &found, &index, &ref)) {
 900  i :          return false;
 901    :        }
 902  E :        if (!found)
 903  i :          continue;
 904    :  
 905    :        // Update the various metadata associated with this symbol.
 906    :        // TODO(chrisha): Currently the import index must be unique. This ensures
 907    :        //     uniqueness for delay-load imports by setting the MSB, and combining
 908    :        //     the module index with the symbol index.
 909  E :        UpdateModuleSymbolInfo(j, true, false, module);
 910  E :        UpdateModuleSymbolReference(j, ref, true, module);
 911  E :      }
 912  E :    }
 913    :  
 914  E :    return true;
 915  E :  }
 916    :  
 917    :  }  // namespace transforms
 918    :  }  // namespace pe

Coverage information generated Thu Mar 26 16:15:41 2015.