From 31ec99d69c5ee32fd7eb213101b5cf9c120d2fc0 Mon Sep 17 00:00:00 2001
From: Alex Cameron <asc@tetsuo.sh>
Date: Wed, 5 Jan 2022 07:22:12 +1100
Subject: [PATCH] sleigh-lift: Add support for loading PSPEC files (#37)

* sleigh-lift: Add support for loading PSPEC files

* sleigh-lift: Remove code to manually switch on `longMode`

* doc: Fix README usage text

* Update main.cpp

Co-authored-by: Peter Goodman <peter@trailofbits.com>
---
 README.md                      |  2 +-
 tools/sleigh-lift/src/main.cpp | 64 +++++++++++++++++++++++++++++-----
 2 files changed, 56 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index 465f727..4349e3e 100644
--- a/README.md
+++ b/README.md
@@ -78,7 +78,7 @@ cmake --build build --target package
 An example program called `sleigh-lift` has been included to demonstrate how to use the SLEIGH API. It takes a hexadecimal string of bytes and either disassembles it or lifts it to p-code. The program can be invoked like so, where the `action` argument must be either `disassemble` or `pcode`:
 
 ```sh
-sleigh-lift [action] [sla_file] [bytes] [-a address] [-p root_sla_dir]
+sleigh-lift [action] [sla_file] [bytes] [-a address] [-p root_sla_dir] [-s pspec_file]
 ```
 
 For example, to disassemble the following byte string:
diff --git a/tools/sleigh-lift/src/main.cpp b/tools/sleigh-lift/src/main.cpp
index e70d8ce..dedf4e7 100644
--- a/tools/sleigh-lift/src/main.cpp
+++ b/tools/sleigh-lift/src/main.cpp
@@ -14,7 +14,7 @@
 
 static void PrintUsage(void) {
   std::cerr << "Usage: sleigh-lift [action] [sla_file] [bytes] [-a address] "
-               "[-p root_sla_dir]"
+               "[-p root_sla_dir] [-s pspec_file]"
             << std::endl;
 }
 
@@ -29,7 +29,7 @@ public:
   }
 
   void loadFill(unsigned char *ptr, int size, const Address &addr) override {
-    uint8_t start = addr.getOffset();
+    uint64_t start = addr.getOffset();
     for (int i = 0; i < size; ++i) {
       uint64_t offset = start + i;
       if (offset >= base_addr) {
@@ -134,11 +134,9 @@ static void PrintPcode(Sleigh &engine, uint64_t addr, size_t len) {
 }
 
 struct LiftArgs {
-  const std::string action;
-  const std::string sla_file_name;
-  const std::string bytes;
+  const std::string action, sla_file_name, bytes;
   const std::optional<uint64_t> addr;
-  const std::optional<std::string> root_sla_dir;
+  const std::optional<std::string> root_sla_dir, pspec_file_name;
 };
 
 std::optional<LiftArgs> ParseArgs(int argc, char *argv[]) {
@@ -159,7 +157,7 @@ std::optional<LiftArgs> ParseArgs(int argc, char *argv[]) {
 
   // Get optional args
   std::optional<uint64_t> addr;
-  std::optional<std::string> root_sla_dir;
+  std::optional<std::string> root_sla_dir, pspec_file_name;
   while (arg_index < argc) {
     const std::string flag = argv[arg_index++];
     if (arg_index == argc) {
@@ -187,13 +185,20 @@ std::optional<LiftArgs> ParseArgs(int argc, char *argv[]) {
         return {};
       }
       root_sla_dir = argv[arg_index++];
+    } else if (flag == "-s") {
+      if (pspec_file_name) {
+        std::cerr << "-s flag provided multiple times" << std::endl;
+        return {};
+      }
+      pspec_file_name = argv[arg_index++];
     } else {
       std::cerr << "Unrecognised optional flag: " << flag << std::endl;
       return {};
     }
   }
-  return LiftArgs{std::move(action), std::move(sla_file_name), std::move(bytes),
-                  addr, std::move(root_sla_dir)};
+  return LiftArgs{std::move(action),       std::move(sla_file_name),
+                  std::move(bytes),        addr,
+                  std::move(root_sla_dir), std::move(pspec_file_name)};
 }
 
 int main(int argc, char *argv[]) {
@@ -220,7 +225,48 @@ int main(int argc, char *argv[]) {
   DocumentStorage storage;
   Element *root = storage.openDocument(*sla_file_path)->getRoot();
   storage.registerTag(root);
+  std::optional<std::filesystem::path> pspec_file_path;
+  if (args->pspec_file_name) {
+    // A PSPEC file was explicitly supplied
+    pspec_file_path = args->root_sla_dir
+                          ? sleigh::FindSpecFile(*args->pspec_file_name,
+                                                 {*args->root_sla_dir})
+                          : sleigh::FindSpecFile(*args->pspec_file_name);
+    if (!pspec_file_path) {
+      std::cerr << "Could not find PSPEC file: " << *args->pspec_file_name
+                << std::endl;
+      return EXIT_FAILURE;
+    }
+  } else {
+    // Otherwise, see if there's a PSPEC file named identically to the SLA file
+    pspec_file_path = *sla_file_path;
+    pspec_file_path->replace_extension(".pspec");
+    if (!std::filesystem::exists(*pspec_file_path)) {
+      // If a file with that extension doesn't exist, don't attempt to load it
+      pspec_file_path = {};
+    }
+  }
+  if (pspec_file_path) {
+    Element *pspec_root = storage.openDocument(*pspec_file_path)->getRoot();
+    storage.registerTag(pspec_root);
+  }
   engine.initialize(storage);
+  engine.allowContextSet(false);
+  // Now that context symbol names are loaded by the translator
+  // we can set the default context
+  // This imitates what is done in
+  //   void Architecture::parseProcessorConfig(DocumentStorage &store)
+  const Element *el = storage.getTag("processor_spec");
+  if (el) {
+    const List &list(el->getChildren());
+    for (List::const_iterator iter = list.begin(); iter != list.end(); ++iter) {
+      const string &elname((*iter)->getName());
+      if (elname == "context_data") {
+        ctx.restoreFromSpec(*iter, &engine);
+        break;
+      }
+    }
+  }
   // In order to parse and validate the byte string properly, we need to get the
   // address size from SLEIGH. Therefore this needs to happen after
   // initialization.