initial commit

Signed-off-by: Toni Uhlig <matzeton@googlemail.com>

1 files changed, 165 insertions, 0 deletions

diff --git a/aoe2hd/src/CodeInjector.cpp b/aoe2hd/src/CodeInjector.cpp
new file mode 100755
index 0000000..b62f969
--- /dev/null
+++ b/aoe2hd/src/CodeInjector.cpp
@@ -0,0 +1,165 @@
+#include <assert.h>
+
+#include <sstream>
+#include <iomanip>
+#include <algorithm>
+
+#include "CodeInjector.h"
+
+CodeInjector::CodeInjector(const native_data& nd)
+    : nd(nd)
+{
+    assert(nd.alloc_fn && nd.write_fn);
+}
+
+CodeInjector::~CodeInjector()
+{
+}
+
+bool CodeInjector::allocCodeSegment(const std::string& name, unsigned long siz)
+{
+    if (codeSegExists(name))
+        return false;
+    code_seg seg = {0};
+    seg.siz = siz;
+    seg.addr = nd.alloc_fn(&nd, siz);
+    if (!seg.addr)
+        return false;
+    code_map[name] = seg;
+    return true;
+}
+
+bool CodeInjector::addCode(const std::string& name, const std::string& code_name,
+                           const std::vector<unsigned char>& code)
+{
+    assert(code.size());
+    if (!codeSegExists(name)
+            || codeBinExists(name, code_name))
+        return false;
+    auto cave = findCodeCave(name, code.size());
+    if (!cave)
+        return false;
+    code_bin bin = {0};
+    bin.addr = cave;
+    bin.siz = code.size();
+    if (!nd.write_fn(&nd, bin.addr, &code[0], bin.siz))
+        return false;
+    code_map[name].children[code_name] = bin;
+    return true;
+}
+
+bool CodeInjector::addCode(const std::string& name, const std::string& code_name,
+                           unsigned long siz)
+{
+    assert(siz);
+    std::vector<unsigned char> code(siz, 0x90);
+    return addCode(name, code_name, code);
+}
+
+bool CodeInjector::setCode(const std::string& name, const std::string& code_name,
+                           const std::vector<unsigned char>& code,
+                           unsigned long offset)
+{
+    assert(code.size());
+    if (!codeSegExists(name)
+            || !codeBinExists(name, code_name))
+        return false;
+    code_bin bin = {0};
+    if (!getCodeBin(name, code_name, &bin))
+        return false;
+    assert(bin.addr && bin.siz);
+    if (bin.addr + offset + code.size() > bin.addr + bin.siz)
+        return false;
+    return nd.write_fn(&nd, bin.addr + offset, &code[0], code.size());
+}
+
+bool CodeInjector::delCode(const std::string& name, const std::string& code_name)
+{
+    if (!codeBinExists(name, code_name))
+        return false;
+    code_map[name].children[code_name];
+    return code_map[name].children.erase(code_name) > 0;
+}
+unsigned long CodeInjector::getCodeAddr(const std::string& name, const std::string& code_name)
+{
+    assert(codeBinExists(name, code_name));
+    assert(code_map[name].children[code_name].addr);
+    return code_map[name].children[code_name].addr;
+}
+
+bool CodeInjector::getCodeSeg(const std::string& name, code_seg *seg)
+{
+    assert(seg);
+    if (!codeSegExists(name))
+        return false;
+    *seg = code_map[name];
+    return true;
+}
+
+bool CodeInjector::getCodeBin(const std::string& name, const std::string& code_name, code_bin *bin)
+{
+    assert(bin);
+    if (!codeBinExists(name, code_name))
+        return false;
+    *bin = code_map[name].children[code_name];
+    return true;
+}
+
+std::string CodeInjector::toString()
+{
+    std::stringstream out;
+    for (auto& code : code_map)
+    {
+        out << std::setw(16) << code.first << "[ "
+            << std::setw(8) << std::hex << code.second.addr << " , "
+            << std::setw(8) << std::hex << code.second.siz
+            << " ]" << std::endl;
+        for (auto& child : code.second.children)
+        {
+            out << std::setw(18) << child.first << "[ "
+                << std::setw(8) << std::hex << child.second.addr << " , "
+                << std::setw(6) << std::hex << child.second.siz
+                << " ]" << std::endl;
+        }
+    }
+    return out.str();
+}
+
+bool CodeInjector::codeBinExists(const std::string& name, const std::string& code_name)
+{
+    if (!codeSegExists(name))
+        return false;
+    return code_map[name].children.find(code_name)
+           != code_map[name].children.end();
+}
+
+std::vector<code_bin> CodeInjector::convertCodeSegChildren(const std::string& name)
+{
+    auto cs = code_map[name].children;
+    std::vector<code_bin> ret;
+    ret.reserve(cs.size());
+    std::for_each(cs.begin(), cs.end(),  [&ret](std::pair<const std::string, code_bin> element)
+    {
+        ret.push_back(element.second);
+    });
+    return ret;
+}
+
+unsigned long CodeInjector::findCodeCave(const std::string& name, unsigned long siz)
+{
+    auto& cs = code_map[name];
+    auto cl = convertCodeSegChildren(name);
+    std::sort(cl.begin(), cl.end());
+    unsigned long end_addr = cs.addr;
+    for (auto& el : cl)
+    {
+        if (el.addr >= end_addr + siz)
+        {
+            return end_addr;
+        }
+        end_addr = el.addr + el.siz;
+    }
+    if (end_addr <= cs.addr + cs.siz)
+        return end_addr;
+    return 0;
+}