1    :  // Copyright 2015 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    :  #include "syzygy/refinery/symbols/dia_symbol_provider.h"
  16    :  
  17    :  #include <string>
  18    :  
  19    :  #include "base/strings/stringprintf.h"
  20    :  #include "syzygy/common/com_utils.h"
  21    :  #include "syzygy/pe/dia_util.h"
  22    :  #include "syzygy/refinery/symbols/symbol_provider_util.h"
  23    :  #include "syzygy/refinery/types/dia_crawler.h"
  24    :  
  25  m :  namespace refinery {
  26    :  
  27  m :  DiaSymbolProvider::DiaSymbolProvider() {
  28  m :  }
  29    :  
  30  m :  DiaSymbolProvider::~DiaSymbolProvider() {
  31  m :  }
  32    :  
  33    :  // TODO(manzagop): revise this function using the code from dia_util.h.
  34  m :  bool DiaSymbolProvider::FindOrCreateDiaSession(
  35  m :      const pe::PEFile::Signature& signature,
  36  m :      base::win::ScopedComPtr<IDiaSession>* session) {
  37  m :    DCHECK(session != nullptr);
  38  m :    *session = nullptr;
  39    :  
  40  m :    base::win::ScopedComPtr<IDiaDataSource> tmp_source;
  41  m :    return GetOrLoad(signature, &tmp_source, session);
  42  m :  }
  43    :  
  44  m :  bool DiaSymbolProvider::GetVFTableRVAs(
  45  m :      const pe::PEFile::Signature& signature,
  46  m :      base::hash_set<RelativeAddress>* vftable_rvas) {
  47  m :    DCHECK(vftable_rvas);
  48    :  
  49  m :    base::win::ScopedComPtr<IDiaDataSource> source;
  50  m :    base::win::ScopedComPtr<IDiaSession> session;
  51  m :    if (!GetOrLoad(signature, &source, &session))
  52  m :      return false;
  53    :  
  54  m :    DiaCrawler crawler;
  55  m :    if (!crawler.InitializeForSession(source, session))
  56  m :      return false;
  57    :  
  58  m :    return crawler.GetVFTableRVAs(vftable_rvas);
  59  m :  }
  60    :  
  61  m :  void DiaSymbolProvider::GetCacheKey(const pe::PEFile::Signature& signature,
  62  m :                                      base::string16* cache_key) {
  63  m :    DCHECK(cache_key);
  64    :    // Note that the cache key does not contain the module's base address.
  65  m :    base::SStringPrintf(cache_key, L"%ls:%d:%d:%d",
  66  m :                        base::FilePath(signature.path).BaseName().value().c_str(),
  67  m :                        signature.module_size, signature.module_checksum,
  68  m :                        signature.module_time_date_stamp);
  69  m :  }
  70    :  
  71  m :  bool DiaSymbolProvider::GetOrLoad(
  72  m :      const pe::PEFile::Signature& signature,
  73  m :      base::win::ScopedComPtr<IDiaDataSource>* source,
  74  m :      base::win::ScopedComPtr<IDiaSession>* session) {
  75  m :    DCHECK(source); DCHECK(session);
  76  m :    *source = nullptr;
  77  m :    *session = nullptr;
  78    :  
  79  m :    base::string16 cache_key;
  80  m :    GetCacheKey(signature, &cache_key);
  81    :  
  82    :    // Look for a pre-existing entry.
  83  m :    auto session_it = pdb_sessions_.find(cache_key);
  84  m :    if (session_it != pdb_sessions_.end()) {
  85  m :      if (session_it->second == nullptr)
  86  m :        return false;  // Negative cache entry.
  87    :  
  88  m :      auto source_it = pdb_sources_.find(cache_key);
  89    :      // Consistency validation: if there's a session, there should be a source.
  90  m :      DCHECK(source_it != pdb_sources_.end());
  91  m :      DCHECK(session_it->second.get() != nullptr);
  92    :  
  93  m :      *source = source_it->second;
  94  m :      *session = session_it->second;
  95  m :      return true;
  96  m :    }
  97    :  
  98    :    // The module is not in the cache. Create negative cache entries, which will
  99    :    // be replaced on success.
 100  m :    pdb_sources_[cache_key] = base::win::ScopedComPtr<IDiaDataSource>();
 101  m :    pdb_sessions_[cache_key] = base::win::ScopedComPtr<IDiaSession>();
 102    :  
 103    :    // Attempt to create a dia session for the module.
 104  m :    base::FilePath pdb_path;
 105  m :    if (!GetPdbPath(signature, &pdb_path))
 106  m :      return false;
 107    :  
 108    :    // Get a source for the pdb.
 109  m :    base::win::ScopedComPtr<IDiaDataSource> pdb_source;
 110  m :    if (!pe::CreateDiaSource(pdb_source.Receive()))
 111  m :      return false;
 112    :  
 113    :    // Get the session.
 114  m :    base::win::ScopedComPtr<IDiaSession> pdb_session;
 115  m :    if (!pe::CreateDiaSession(pdb_path, pdb_source.get(), pdb_session.Receive()))
 116  m :      return false;
 117    :  
 118    :    // Cache source and session.
 119  m :    pdb_sources_[cache_key] = pdb_source;
 120  m :    pdb_sessions_[cache_key] = pdb_session;
 121    :  
 122  m :    *source = pdb_source;
 123  m :    *session = pdb_session;
 124  m :    return true;
 125  m :  }
 126    :  
 127  m :  }  // namespace refinery