Implement some functions in NativeSession.

Summary:
This change implements readFromExe, and calculating VA and RVA, which
are some of the functionalities that will be used for native PDB reading
for llvm symbolizer.

bug: https://bugs.llvm.org/show_bug.cgi?id=41795

Reviewers: hans, amccarth, rnk

Subscribers: hiraditya, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D78128

llvm/include/llvm/DebugInfo/PDB/Native/NativeSession.h

@@ -26,15 +26,23 @@ class PDBFile;
 class NativeExeSymbol;
 
 class NativeSession : public IPDBSession {
+  struct PdbSearchOptions {
+    StringRef ExePath;
+    // FIXME: Add other PDB search options (_NT_SYMBOL_PATH, symsrv)
+  };
+
 public:
   NativeSession(std::unique_ptr<PDBFile> PdbFile,
                 std::unique_ptr<BumpPtrAllocator> Allocator);
   ~NativeSession() override;
 
   static Error createFromPdb(std::unique_ptr<MemoryBuffer> MB,
                              std::unique_ptr<IPDBSession> &Session);
+  static Error createFromPdbPath(StringRef PdbPath,
+                                 std::unique_ptr<IPDBSession> &Session);
   static Error createFromExe(StringRef Path,
                              std::unique_ptr<IPDBSession> &Session);
+  static Expected<std::string> searchForPdb(const PdbSearchOptions &Opts);
 
   uint64_t getLoadAddress() const override;
   bool setLoadAddress(uint64_t Address) override;
@@ -109,6 +117,7 @@ class NativeSession : public IPDBSession {
 
   SymbolCache Cache;
   SymIndexId ExeSymbol = 0;
+  uint64_t LoadAddress = 0;
 };
 } // namespace pdb
 } // namespace llvm

llvm/lib/DebugInfo/PDB/Native/NativeSession.cpp

@@ -12,6 +12,7 @@
 #include "llvm/DebugInfo/CodeView/TypeIndex.h"
 #include "llvm/DebugInfo/PDB/IPDBEnumChildren.h"
 #include "llvm/DebugInfo/PDB/IPDBSourceFile.h"
+#include "llvm/DebugInfo/PDB/Native/DbiStream.h"
 #include "llvm/DebugInfo/PDB/Native/NativeCompilandSymbol.h"
 #include "llvm/DebugInfo/PDB/Native/NativeEnumInjectedSources.h"
 #include "llvm/DebugInfo/PDB/Native/NativeEnumTypes.h"
@@ -29,7 +30,10 @@
 #include "llvm/Support/BinaryByteStream.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorOr.h"
+#include "llvm/Support/FileSystem.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Path.h"
+#include "llvm/object/COFF.h"
 
 #include <algorithm>
 #include <cassert>
@@ -75,14 +79,112 @@ Error NativeSession::createFromPdb(std::unique_ptr<MemoryBuffer> Buffer,
   return Error::success();
 }
 
-Error NativeSession::createFromExe(StringRef Path,
+static Expected<std::unique_ptr<PDBFile>>
+loadPdbFile(StringRef PdbPath, std::unique_ptr<BumpPtrAllocator> &Allocator) {
+  ErrorOr<std::unique_ptr<MemoryBuffer>> ErrorOrBuffer =
+      MemoryBuffer::getFile(PdbPath, /*FileSize=*/-1,
+                            /*RequiresNullTerminator=*/false);
+  if (!ErrorOrBuffer)
+    return make_error<RawError>(ErrorOrBuffer.getError());
+  std::unique_ptr<llvm::MemoryBuffer> Buffer = std::move(*ErrorOrBuffer);
+
+  PdbPath = Buffer->getBufferIdentifier();
+  file_magic Magic;
+  auto EC = identify_magic(PdbPath, Magic);
+  if (EC || Magic != file_magic::pdb)
+    return make_error<RawError>(EC);
+
+  auto Stream = std::make_unique<MemoryBufferByteStream>(std::move(Buffer),
+                                                         llvm::support::little);
+
+  auto File = std::make_unique<PDBFile>(PdbPath, std::move(Stream), *Allocator);
+  if (auto EC = File->parseFileHeaders())
+    return std::move(EC);
+
+  if (auto EC = File->parseStreamData())
+    return std::move(EC);
+
+  return File;
+}
+
+Error NativeSession::createFromPdbPath(StringRef PdbPath,
+                                       std::unique_ptr<IPDBSession> &Session) {
+  auto Allocator = std::make_unique<BumpPtrAllocator>();
+  auto PdbFile = loadPdbFile(PdbPath, Allocator);
+  if (!PdbFile)
+    return PdbFile.takeError();
+
+  Session = std::make_unique<NativeSession>(std::move(PdbFile.get()),
+                                            std::move(Allocator));
+  return Error::success();
+}
+
+static Expected<std::string> getPdbPathFromExe(StringRef ExePath) {
+  Expected<object::OwningBinary<object::Binary>> BinaryFile =
+      object::createBinary(ExePath);
+  if (!BinaryFile)
+    return BinaryFile.takeError();
+
+  const object::COFFObjectFile *ObjFile =
+      dyn_cast<object::COFFObjectFile>(BinaryFile->getBinary());
+  if (!ObjFile)
+    return make_error<RawError>(raw_error_code::invalid_format);
+
+  StringRef PdbPath;
+  const llvm::codeview::DebugInfo *PdbInfo = nullptr;
+  if (auto EC = ObjFile->getDebugPDBInfo(PdbInfo, PdbPath))
+    return make_error<RawError>(EC);
+
+  return std::string(PdbPath);
+}
+
+Error NativeSession::createFromExe(StringRef ExePath,
                                    std::unique_ptr<IPDBSession> &Session) {
-  return make_error<RawError>(raw_error_code::feature_unsupported);
+  Expected<std::string> PdbPath = getPdbPathFromExe(ExePath);
+  if (!PdbPath)
+    return PdbPath.takeError();
+
+  file_magic Magic;
+  auto EC = identify_magic(PdbPath.get(), Magic);
+  if (EC || Magic != file_magic::pdb)
+    return make_error<RawError>(EC);
+
+  auto Allocator = std::make_unique<BumpPtrAllocator>();
+  auto File = loadPdbFile(PdbPath.get(), Allocator);
+  if (!File)
+    return File.takeError();
+
+  Session = std::make_unique<NativeSession>(std::move(File.get()),
+                                            std::move(Allocator));
+
+  return Error::success();
 }
 
-uint64_t NativeSession::getLoadAddress() const { return 0; }
+Expected<std::string>
+NativeSession::searchForPdb(const PdbSearchOptions &Opts) {
+  Expected<std::string> PathOrErr = getPdbPathFromExe(Opts.ExePath);
+  if (!PathOrErr)
+    return PathOrErr.takeError();
+  StringRef PdbName = sys::path::filename(PathOrErr.get());
 
-bool NativeSession::setLoadAddress(uint64_t Address) { return false; }
+  // Check if pdb exists in the executable directory.
+  SmallString<128> PdbPath = StringRef(Opts.ExePath);
+  sys::path::remove_filename(PdbPath);
+  sys::path::append(PdbPath, PdbName);
+
+  auto Allocator = std::make_unique<BumpPtrAllocator>();
+  if (loadPdbFile(PdbPath, Allocator))
+    return std::string(PdbPath);
+
+  return make_error<RawError>("PDB not found");
+}
+
+uint64_t NativeSession::getLoadAddress() const { return LoadAddress; }
+
+bool NativeSession::setLoadAddress(uint64_t Address) {
+  LoadAddress = Address;
+  return true;
+}
 
 std::unique_ptr<PDBSymbolExe> NativeSession::getGlobalScope() {
   return PDBSymbol::createAs<PDBSymbolExe>(*this, getNativeGlobalScope());
@@ -95,12 +197,30 @@ NativeSession::getSymbolById(SymIndexId SymbolId) const {
 
 bool NativeSession::addressForVA(uint64_t VA, uint32_t &Section,
                                  uint32_t &Offset) const {
-  return false;
+  uint32_t RVA = VA - getLoadAddress();
+  return addressForRVA(RVA, Section, Offset);
 }
 
-bool NativeSession::addressForRVA(uint32_t VA, uint32_t &Section,
+bool NativeSession::addressForRVA(uint32_t RVA, uint32_t &Section,
                                   uint32_t &Offset) const {
-  return false;
+  auto Dbi = Pdb->getPDBDbiStream();
+  if (!Dbi)
+    return false;
+
+  Section = 0;
+  Offset = 0;
+
+  if ((int32_t)RVA < 0)
+    return true;
+
+  Offset = RVA;
+  for (; Section < Dbi->getSectionHeaders().size(); ++Section) {
+    auto &Sec = Dbi->getSectionHeaders()[Section];
+    if (RVA < Sec.VirtualAddress)
+      return true;
+    Offset = RVA - Sec.VirtualAddress;
+  }
+  return true;
 }
 
 std::unique_ptr<PDBSymbol>

llvm/lib/DebugInfo/PDB/PDB.cpp

@@ -23,15 +23,8 @@ using namespace llvm::pdb;
 Error llvm::pdb::loadDataForPDB(PDB_ReaderType Type, StringRef Path,
                                 std::unique_ptr<IPDBSession> &Session) {
   // Create the correct concrete instance type based on the value of Type.
-  if (Type == PDB_ReaderType::Native) {
-    ErrorOr<std::unique_ptr<MemoryBuffer>> ErrorOrBuffer =
-        MemoryBuffer::getFileOrSTDIN(Path, /*FileSize=*/-1,
-                                     /*RequiresNullTerminator=*/false);
-    if (!ErrorOrBuffer)
-      return errorCodeToError(ErrorOrBuffer.getError());
-
-    return NativeSession::createFromPdb(std::move(*ErrorOrBuffer), Session);
-  }
+  if (Type == PDB_ReaderType::Native)
+    return NativeSession::createFromPdbPath(Path, Session);
 
 #if LLVM_ENABLE_DIA_SDK
   return DIASession::createFromPdb(Path, Session);
@@ -43,8 +36,17 @@ Error llvm::pdb::loadDataForPDB(PDB_ReaderType Type, StringRef Path,
 Error llvm::pdb::loadDataForEXE(PDB_ReaderType Type, StringRef Path,
                                 std::unique_ptr<IPDBSession> &Session) {
   // Create the correct concrete instance type based on the value of Type.
-  if (Type == PDB_ReaderType::Native)
-    return NativeSession::createFromExe(Path, Session);
+  if (Type == PDB_ReaderType::Native) {
+    if (auto Err = NativeSession::createFromExe(Path, Session)) {
+      consumeError(std::move(Err));
+
+      Expected<std::string> PdbPath = NativeSession::searchForPdb({Path});
+      if (!PdbPath)
+        return PdbPath.takeError();
+      return NativeSession::createFromPdbPath(PdbPath.get(), Session);
+    }
+    return Error::success();
+  }
 
 #if LLVM_ENABLE_DIA_SDK
   return DIASession::createFromExe(Path, Session);

llvm/unittests/DebugInfo/PDB/CMakeLists.txt

@@ -6,6 +6,7 @@ set(LLVM_LINK_COMPONENTS
 
 add_llvm_unittest_with_input_files(DebugInfoPDBTests
   HashTableTest.cpp
+  NativeSessionTest.cpp
   NativeSymbolReuseTest.cpp
   StringTableBuilderTest.cpp
   PDBApiTest.cpp

llvm/unittests/DebugInfo/PDB/Inputs/SimpleTest.cpp

@@ -0,0 +1,4 @@
+// Compile with "cl /c /Zi SimpleTest.cpp"
+// Link with "link SimpleTest.obj /debug /nodefaultlib /entry:main"
+
+int main() { return 0; }

llvm/unittests/DebugInfo/PDB/NativeSessionTest.cpp

@@ -0,0 +1,93 @@
+//===- llvm/unittest/DebugInfo/PDB/NativeSessionTest.cpp ------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/DebugInfo/PDB/Native/NativeSession.h"
+#include "llvm/DebugInfo/PDB/IPDBSession.h"
+#include "llvm/DebugInfo/PDB/PDB.h"
+#include "llvm/Support/Path.h"
+
+#include "llvm/Testing/Support/Error.h"
+
+#include "gtest/gtest.h"
+
+#include <vector>
+
+using namespace llvm;
+using namespace llvm::pdb;
+
+extern const char *TestMainArgv0;
+
+static std::string getExePath() {
+  SmallString<128> InputsDir = unittest::getInputFileDirectory(TestMainArgv0);
+  llvm::sys::path::append(InputsDir, "SimpleTest.exe");
+  return std::string(InputsDir);
+}
+
+TEST(NativeSessionTest, TestCreateFromExe) {
+  std::unique_ptr<IPDBSession> S;
+  Error E = pdb::loadDataForEXE(PDB_ReaderType::Native, getExePath(), S);
+  ASSERT_THAT_ERROR(std::move(E), Succeeded());
+}
+
+TEST(NativeSessionTest, TestSetLoadAddress) {
+  std::unique_ptr<IPDBSession> S;
+  Error E = pdb::loadDataForEXE(PDB_ReaderType::Native, getExePath(), S);
+  ASSERT_THAT_ERROR(std::move(E), Succeeded());
+
+  S->setLoadAddress(123);
+  EXPECT_EQ(S->getLoadAddress(), 123U);
+}
+
+TEST(NativeSessionTest, TestAddressForVA) {
+  std::unique_ptr<IPDBSession> S;
+  Error E = pdb::loadDataForEXE(PDB_ReaderType::Native, getExePath(), S);
+  ASSERT_THAT_ERROR(std::move(E), Succeeded());
+
+  uint64_t LoadAddr = S->getLoadAddress();
+  uint32_t Section;
+  uint32_t Offset;
+  ASSERT_TRUE(S->addressForVA(LoadAddr + 5000, Section, Offset));
+  EXPECT_EQ(1U, Section);
+  EXPECT_EQ(904U, Offset);
+
+  ASSERT_TRUE(S->addressForVA(-1, Section, Offset));
+  EXPECT_EQ(0U, Section);
+  EXPECT_EQ(0U, Offset);
+
+  ASSERT_TRUE(S->addressForVA(4, Section, Offset));
+  EXPECT_EQ(0U, Section);
+  EXPECT_EQ(4U, Offset);
+
+  ASSERT_TRUE(S->addressForVA(LoadAddr + 100000, Section, Offset));
+  EXPECT_EQ(3U, Section);
+  EXPECT_EQ(83616U, Offset);
+}
+
+TEST(NativeSessionTest, TestAddressForRVA) {
+  std::unique_ptr<IPDBSession> S;
+  Error E = pdb::loadDataForEXE(PDB_ReaderType::Native, getExePath(), S);
+  ASSERT_THAT_ERROR(std::move(E), Succeeded());
+
+  uint32_t Section;
+  uint32_t Offset;
+  ASSERT_TRUE(S->addressForVA(5000, Section, Offset));
+  EXPECT_EQ(1U, Section);
+  EXPECT_EQ(904U, Offset);
+
+  ASSERT_TRUE(S->addressForVA(-1, Section, Offset));
+  EXPECT_EQ(0U, Section);
+  EXPECT_EQ(0U, Offset);
+
+  ASSERT_TRUE(S->addressForVA(4, Section, Offset));
+  EXPECT_EQ(0U, Section);
+  EXPECT_EQ(4U, Offset);
+
+  ASSERT_TRUE(S->addressForVA(100000, Section, Offset));
+  EXPECT_EQ(3U, Section);
+  EXPECT_EQ(83616U, Offset);
+}