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