Add support for adding LC_UNIXTHREAD commands in a MachO (include min…

…or fixes)

include/LIEF/MachO/Builder.hpp

@@ -103,6 +103,9 @@ class LIEF_API Builder {
   template<class T>
   void build(DyldEnvironment* de);
 
+  template<class T>
+  void build(ThreadCommand* tc);
+
   template <typename T>
   void build_segments(void);
 

include/LIEF/MachO/ThreadCommand.hpp

@@ -34,7 +34,8 @@ class LIEF_API ThreadCommand : public LoadCommand {
   friend class BinaryParser;
   public:
     ThreadCommand(void);
-    ThreadCommand(const thread_command *cmd);
+    ThreadCommand(const thread_command *cmd, CPU_TYPES arch=CPU_TYPES::CPU_TYPE_ANY);
+    ThreadCommand(uint32_t flavor, uint32_t count, CPU_TYPES arch=CPU_TYPES::CPU_TYPE_ANY);
 
     ThreadCommand& operator=(const ThreadCommand& copy);
     ThreadCommand(const ThreadCommand& copy);
@@ -43,8 +44,9 @@ class LIEF_API ThreadCommand : public LoadCommand {
 
     virtual ~ThreadCommand(void);
 
-    uint32_t flavor(void) const;
-    uint32_t count(void) const;
+    uint32_t  flavor(void) const;
+    uint32_t  count(void) const;
+    CPU_TYPES architecture(void) const;
 
     const std::vector<uint8_t>& state(void) const;
     std::vector<uint8_t>& state(void);
@@ -54,6 +56,7 @@ class LIEF_API ThreadCommand : public LoadCommand {
     void state(const std::vector<uint8_t>& state);
     void flavor(uint32_t flavor);
     void count(uint32_t count);
+    void architecture(CPU_TYPES arch);
 
     bool operator==(const ThreadCommand& rhs) const;
     bool operator!=(const ThreadCommand& rhs) const;

src/MachO/Binary.cpp

@@ -536,6 +536,16 @@ void Binary::shift_command(size_t width, size_t from_offset) {
     if (symcmd.strings_offset() > from_offset) {
       symcmd.strings_offset(symcmd.strings_offset() + width);
     }
+
+    for (Symbol* s : this->symbols_) {
+      static constexpr size_t N_TYPE = 0x0e;
+      if (static_cast<N_LIST_TYPES>(s->type() & N_TYPE) == N_LIST_TYPES::N_SECT) {
+        uint64_t value = s->value();
+        if (value > from_offset) {
+          s->value(value + width);
+        }
+      }
+    }
   }
 
   // Data In Code
@@ -868,15 +878,11 @@ bool Binary::remove(const LoadCommand& command) {
 
 
 bool Binary::remove(LOAD_COMMAND_TYPES type) {
-  if (not this->has(type)) {
-    return false;
-  }
-
-  bool done = false;
-  while (not done) {
-    done = not this->remove(this->get(type));
+  bool removed = false;
+  while (this->has(type)) {
+    removed = this->remove(this->get(type));
   }
-  return true;
+  return removed;
 }
 
 bool Binary::remove_command(size_t index) {

src/MachO/Builder.cpp

@@ -153,6 +153,11 @@ void Builder::build(void) {
       continue;
     }
 
+    if (cmd->is<ThreadCommand>()) {
+      this->build<T>(cmd->as<ThreadCommand>());
+      continue;
+    }
+
     if (cmd->is<BuildVersion>()) {
       this->build<T>(cmd->as<BuildVersion>());
       continue;

src/MachO/Builder.tcc

@@ -938,6 +938,56 @@ void Builder::build(DyldEnvironment* de) {
   de->originalData_.insert(std::end(de->originalData_), padding, 0);
 }
 
+template<class T>
+void Builder::build(ThreadCommand* tc) {
+  thread_command raw_cmd;
+
+  const std::vector<uint8_t>& state = tc->state();
+
+  const uint32_t raw_size = sizeof(thread_command) + state.size();
+  const uint32_t size_needed = align(raw_size, sizeof(typename T::uint));
+  const uint32_t padding = size_needed - raw_size;
+
+  if (tc->originalData_.size() != size_needed or
+      tc->size() != size_needed) {
+    LOG(WARNING) << "Not enough spaces to rebuild thread command: Skip!";
+    LOG(WARNING) << std::hex << tc->originalData_.size() << " vs " << size_needed;
+  }
+
+  const uint32_t state_size_needed = tc->count() * sizeof(uint32_t);
+  if (state.size() < state_size_needed) {
+    LOG(WARNING) << "Not enough data in state to rebuild thread command: Skip!";
+    LOG(WARNING) << std::hex << state.size() << " vs " << state_size_needed;
+  }
+
+  std::fill(
+    reinterpret_cast<uint8_t*>(&raw_cmd),
+    reinterpret_cast<uint8_t*>(&raw_cmd) + sizeof(dylinker_command),
+    0);
+
+  raw_cmd.cmd      = static_cast<uint32_t>(tc->command());
+  raw_cmd.cmdsize  = static_cast<uint32_t>(size_needed);
+  raw_cmd.flavor   = static_cast<uint32_t>(tc->flavor());
+  raw_cmd.count    = static_cast<uint32_t>(tc->count());
+
+  tc->size_ = size_needed;
+  tc->originalData_.clear();
+
+  // Write Header
+  std::move(
+    reinterpret_cast<uint8_t*>(&raw_cmd),
+    reinterpret_cast<uint8_t*>(&raw_cmd) + sizeof(raw_cmd),
+    std::back_inserter(tc->originalData_));
+
+  // Write state
+  std::move(
+    std::begin(state),
+    std::end(state),
+    std::back_inserter(tc->originalData_));
+  tc->originalData_.push_back(0);
+  tc->originalData_.insert(std::end(tc->originalData_), padding, 0);
+}
+
 
 template<class T>
 void Builder::build(BuildVersion* bv) {

src/MachO/ThreadCommand.cpp

@@ -29,14 +29,24 @@ ThreadCommand& ThreadCommand::operator=(const ThreadCommand&) = default;
 ThreadCommand::ThreadCommand(const ThreadCommand&) = default;
 ThreadCommand::~ThreadCommand(void) = default;
 
-ThreadCommand::ThreadCommand(const thread_command *cmd) :
+ThreadCommand::ThreadCommand(const thread_command *cmd, CPU_TYPES arch) :
   LoadCommand::LoadCommand{static_cast<LOAD_COMMAND_TYPES>(cmd->cmd), cmd->cmdsize},
   flavor_{cmd->flavor},
   count_{cmd->count},
-  architecture_{CPU_TYPES::CPU_TYPE_ANY},
+  architecture_{arch},
   state_{}
 {}
 
+ThreadCommand::ThreadCommand(uint32_t flavor, uint32_t count, CPU_TYPES arch) :
+  LoadCommand::LoadCommand{static_cast<LOAD_COMMAND_TYPES>(LIEF::MachO::LOAD_COMMAND_TYPES::LC_UNIXTHREAD),
+                           static_cast<uint32_t>(sizeof(LIEF::MachO::thread_command) + count * sizeof(uint32_t))},
+  flavor_{flavor},
+  count_{count},
+  architecture_{arch},
+  state_(this->size() - sizeof(LIEF::MachO::thread_command), 0)
+{
+}
+
 ThreadCommand* ThreadCommand::clone(void) const {
   return new ThreadCommand(*this);
 }
@@ -50,6 +60,10 @@ uint32_t ThreadCommand::count(void) const {
   return this->count_;
 }
 
+CPU_TYPES ThreadCommand::architecture(void) const {
+  return this->architecture_;
+}
+
 const std::vector<uint8_t>& ThreadCommand::state(void) const {
   return this->state_;
 }
@@ -104,6 +118,10 @@ void ThreadCommand::count(uint32_t count) {
   this->count_ = count;
 }
 
+void ThreadCommand::architecture(CPU_TYPES arch) {
+  this->architecture_ = arch;
+}
+
 void ThreadCommand::accept(Visitor& visitor) const {
   visitor.visit(*this);
 }