1 files changed, 54 insertions, 19 deletions

diff --git a/run/runner.cpp b/run/runner.cpp
index 8f19bfb..d8b9cad 100644
--- a/run/runner.cpp
+++ b/run/runner.cpp
@@ -4,40 +4,77 @@
 #include <cstring>
 #include <cassert>
 #include <csignal>
+#include <cmath>
 
 #include <iostream>
+#include <memory>
 
 #include <sys/socket.h>
 #include <sys/un.h>
 #include <sys/mman.h>
 
+#include <asm/ldt.h>
+
 #include "common.hpp"
 #include "dos_emu.hpp"
 
 int loader_socket;
 
+class mmap_deleter{
+    size_t size;
+
+public:
+    mmap_deleter(size_t size)
+    : size(size)
+    {}
+
+    void operator()(void* ptr) const {
+        munmap(ptr, size);
+    }
+};
+
 enum ErrorCode {
     ok,
     socket_creation_failed,
     connection_error
 };
 
+using mmap_pointer_t = std::unique_ptr<void, mmap_deleter>;
+
 static ErrorCode connect_to_loader();
-static BinaryLoadMessage receive_message();
-static void load_binary(BinaryLoadMessage const& message);
+
+template<typename T>
+static T receive_message();
+static mmap_pointer_t map_memory(BinarySizeMessage const& message);
+static void load_binary(mmap_pointer_t& binary, uint32_t binary_size);
 static void register_signal_handlers();
 
 int main(int, char*[]) {
     if(connect_to_loader() != ok) {
         return -1;
     }
-    BinaryLoadMessage message = receive_message();
-    load_binary(message);
+    auto size_message = receive_message<BinarySizeMessage>();
+    auto binary = map_memory(size_message);
+
+    BinaryBaseOffsetMessage offset_message
+    {
+        .offset = (uint32_t)binary.get()
+    };
+
+    send(loader_socket, &offset_message, sizeof offset_message, 0);
+
+    auto binary_load_message = receive_message<BinaryLoadMessage>();
+    load_binary(binary, size_message.binary_size);
+
     register_signal_handlers();
 
-    asm ("JMP *%0 \n\t"
+    std::cout << std::hex << "Binary base: " << offset_message.offset << std::endl;
+
+    asm (
+        "jmp *%0        \n\t"
         :
-        : "r" (message.eip));
+        : "r" (binary_load_message.eip)
+    );
 }
 
 static ErrorCode connect_to_loader() {
@@ -58,28 +95,26 @@ static ErrorCode connect_to_loader() {
     return ok;
 }
 
-static BinaryLoadMessage receive_message() {
-    BinaryLoadMessage message;
+template <typename T>
+static T receive_message() {
+    T message;
     uint32_t bytes_received = recv(loader_socket, &message, sizeof message, 0);
-    assert((bytes_received == sizeof message) && "Didn't receive full BinaryLoadMessage");
+    assert((bytes_received == sizeof message) && "Didn't receive message");
     return message;
 }
 
-static void load_binary(BinaryLoadMessage const& message) {
-    uint8_t* binary = static_cast<uint8_t*>(mmap((void*)message.base_address, message.binary_size - message.base_address, PROT_READ | PROT_WRITE | PROT_EXEC, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0));
+static mmap_pointer_t map_memory(BinarySizeMessage const& message) {
+    void* binary = mmap(NULL, message.binary_size, PROT_READ | PROT_WRITE | PROT_EXEC, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
     if(binary == MAP_FAILED) {
         std::cerr << "error during mmap: " << errno << " " << strerror(errno) << std::endl;
         exit(1);
     }
+    return mmap_pointer_t(binary, message.binary_size);
+}
 
-    assert((binary == (void*)message.base_address) && "binary couldn't be loaded on the prefered address.");
-
-    uint8_t x[message.base_address];
-    recv(loader_socket, x, sizeof(x), 0);
-
-    size_t bytes_received = recv(loader_socket, binary, message.binary_size - message.base_address, 0);
-
-    assert((bytes_received == message.binary_size - (size_t)message.base_address) && "bytes_received did not match expected");
+static void load_binary(mmap_pointer_t& binary, uint32_t binary_size) {
+    size_t bytes_received = recv(loader_socket, binary.get(), binary_size, 0);
+    assert((bytes_received == binary_size) && "bytes_received did not match expected");
 }
 
 static void register_signal_handlers() {