Code and data segment

3 files changed, 130 insertions, 57 deletions

diff --git a/elf.cpp b/elf.cpp
index 2923392..0ad7a04 100644
--- a/elf.cpp
+++ b/elf.cpp
@@ -2,50 +2,82 @@
 
 #include "file.h"
 
+#include <iostream>
+
 // TODO
 //#include <boost/endian/conversion.hpp>
 
 #include <cstring>
 
-// We need at least 4096 bytes aligment. Otherwise we get segfaults at startup
-#define PAD 0x1000
+#define PAD
 
 // Helper Functions
 namespace {
- void AddFileHeader(std::vector<uint8_t>& data, const std::vector<uint8_t>& code)
+ // On amd64, we need at least 4096 bytes aligment. Otherwise we get segfaults at startup
+ const size_t PAGE_SIZE = 0x1000;
+
+ // Helper function: returns size with full pages: 0, 4096, ...
+ size_t paged_size(const std::vector<uint8_t>& data)
+ {
+  return (data.size() + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
+ }
+
+ void AddFileHeader(std::vector<uint8_t>& elf, const std::vector<uint8_t>& code, const std::vector<uint8_t>& data)
  {
   Elf::FileHeader fh;
 
 #ifdef PAD
-  fh.e_shoff = PAD + code.size() + 17;
+  fh.e_shoff = PAGE_SIZE + paged_size(code) + data.size() + 23;
 #else
-  fh.e_shoff = sizeof(fh) + sizeof(Elf::ProgramHeader) + code.size() + 17;
+  fh.e_shoff = sizeof(fh) + sizeof(Elf::ProgramHeader) * 2 + code.size() + data.size() + 23;
 #endif
 
-  size_t old_size {data.size()};
-  data.resize(old_size + sizeof(fh));
-  std::memcpy(data.data() + old_size, reinterpret_cast<void*>(&fh), sizeof(fh));
+  size_t old_size {elf.size()};
+  elf.resize(old_size + sizeof(fh));
+  std::memcpy(elf.data() + old_size, reinterpret_cast<void*>(&fh), sizeof(fh));
  }
 
- void AddProgramHeader(std::vector<uint8_t>& data, const std::vector<uint8_t>& code)
+ void AddProgramHeaderCode(std::vector<uint8_t>& elf, const std::vector<uint8_t>& code, const std::vector<uint8_t>& data)
  {
   Elf::ProgramHeader ph;
 
 #ifdef PAD
-  ph.p_offset = PAD;
-  ph.p_align = PAD;
+  ph.p_offset = PAGE_SIZE;
+  ph.p_align = PAGE_SIZE;
 #else
-  ph.p_offset = sizeof(Elf::FileHeader) + sizeof(Elf::ProgramHeader);
+  ph.p_offset = sizeof(Elf::FileHeader) + sizeof(Elf::ProgramHeader) * 2;
 #endif
   ph.p_filesz = code.size();
   ph.p_memsz = code.size();
 
-  size_t old_size {data.size()};
-  data.resize(old_size + sizeof(ph));
-  std::memcpy(data.data() + old_size, reinterpret_cast<void*>(&ph), sizeof(ph));
+  size_t old_size {elf.size()};
+  elf.resize(old_size + sizeof(ph));
+  std::memcpy(elf.data() + old_size, reinterpret_cast<void*>(&ph), sizeof(ph));
+ }
+
+ void AddProgramHeaderData(std::vector<uint8_t>& elf, const std::vector<uint8_t>& code, const std::vector<uint8_t>& data)
+ {
+  Elf::ProgramHeader ph;
+
+#ifdef PAD
+  ph.p_offset = PAGE_SIZE + paged_size(code);
+  ph.p_align = PAGE_SIZE;
+#else
+  ph.p_offset = sizeof(Elf::FileHeader) + sizeof(Elf::ProgramHeader) * 2 + code.size();
+#endif
+  ph.p_filesz = data.size();
+  ph.p_memsz = data.size();
+
+  ph.p_vaddr = 0x401000 + paged_size(code);
+  ph.p_paddr = 0x401000 + paged_size(code);
+  ph.p_flags = 6; // RW
+
+  size_t old_size {elf.size()};
+  elf.resize(old_size + sizeof(ph));
+  std::memcpy(elf.data() + old_size, reinterpret_cast<void*>(&ph), sizeof(ph));
  }
 
- void AddSectionHeaderNull(std::vector<uint8_t>& data, const std::vector<uint8_t>& code)
+ void AddSectionHeaderNull(std::vector<uint8_t>& elf, const std::vector<uint8_t>& code, const std::vector<uint8_t>& data)
  {
   Elf::SectionHeader sh;
 
@@ -56,12 +88,12 @@ namespace {
   sh.sh_addr = 0;
   sh.sh_offset = 0;
 
-  size_t old_size {data.size()};
-  data.resize(old_size + sizeof(sh));
-  std::memcpy(data.data() + old_size, reinterpret_cast<void*>(&sh), sizeof(sh));
+  size_t old_size {elf.size()};
+  elf.resize(old_size + sizeof(sh));
+  std::memcpy(elf.data() + old_size, reinterpret_cast<void*>(&sh), sizeof(sh));
  }
  
- void AddSectionHeaderText(std::vector<uint8_t>& data, const std::vector<uint8_t>& code)
+ void AddSectionHeaderText(std::vector<uint8_t>& elf, const std::vector<uint8_t>& code, const std::vector<uint8_t>& data)
  {
   Elf::SectionHeader sh;
 
@@ -69,73 +101,100 @@ namespace {
   sh.sh_size = code.size();
   sh.sh_type = 1; // program
 #ifdef PAD
-  sh.sh_offset = PAD;
+  sh.sh_offset = PAGE_SIZE;
 #else
-  sh.sh_offset = sizeof(Elf::FileHeader) + sizeof(Elf::ProgramHeader);
+  sh.sh_offset = sizeof(Elf::FileHeader) + sizeof(Elf::ProgramHeader) * 2;
 #endif
   sh.sh_addralign = 1;
 
-  size_t old_size {data.size()};
-  data.resize(old_size + sizeof(sh));
-  std::memcpy(data.data() + old_size, reinterpret_cast<void*>(&sh), sizeof(sh));
+  size_t old_size {elf.size()};
+  elf.resize(old_size + sizeof(sh));
+  std::memcpy(elf.data() + old_size, reinterpret_cast<void*>(&sh), sizeof(sh));
  }
  
- void AddSectionHeaderSectionNames(std::vector<uint8_t>& data, const std::vector<uint8_t>& code)
+ void AddSectionHeaderData(std::vector<uint8_t>& elf, const std::vector<uint8_t>& code, const std::vector<uint8_t>& data)
  {
   Elf::SectionHeader sh;
 
   sh.sh_name = 7; // offset in section
-  sh.sh_size = 17;
+  sh.sh_size = data.size();
+  sh.sh_type = 1; // program (also data)
+#ifdef PAD
+  sh.sh_offset = PAGE_SIZE + paged_size(code);
+#else
+  sh.sh_offset = sizeof(Elf::FileHeader) + sizeof(Elf::ProgramHeader) * 2 + code.size();
+#endif
+  sh.sh_addralign = 1;
+
+  size_t old_size {elf.size()};
+  elf.resize(old_size + sizeof(sh));
+  std::memcpy(elf.data() + old_size, reinterpret_cast<void*>(&sh), sizeof(sh));
+ }
+ 
+ void AddSectionHeaderSectionNames(std::vector<uint8_t>& elf, const std::vector<uint8_t>& code, const std::vector<uint8_t>& data)
+ {
+  Elf::SectionHeader sh;
+
+  sh.sh_name = 13; // offset in section
+  sh.sh_size = 23;
   sh.sh_type = 3; // section names
   sh.sh_flags = 0;
   sh.sh_addr = 0;
 #ifdef PAD
-  sh.sh_offset = PAD + code.size();
+  sh.sh_offset = PAGE_SIZE + paged_size(code) + data.size();
 #else
-  sh.sh_offset = sizeof(Elf::FileHeader) + sizeof(Elf::ProgramHeader) + code.size();
+  sh.sh_offset = sizeof(Elf::FileHeader) + sizeof(Elf::ProgramHeader) * 2 + code.size() + data.size();
 #endif
   sh.sh_addralign = 1;
 
-  size_t old_size {data.size()};
-  data.resize(old_size + sizeof(sh));
-  std::memcpy(data.data() + old_size, reinterpret_cast<void*>(&sh), sizeof(sh));
+  size_t old_size {elf.size()};
+  elf.resize(old_size + sizeof(sh));
+  std::memcpy(elf.data() + old_size, reinterpret_cast<void*>(&sh), sizeof(sh));
  }
 
- void PadUpTo(std::vector<uint8_t>& data, size_t size)
+ void PadUpTo(std::vector<uint8_t>& elf, size_t size)
  {
-  if (data.size() < size) {
-   data.resize(size);
+  if (elf.size() < size) {
+   elf.resize(size);
   } else
    throw std::runtime_error("Padding not possible. Too many bytes already.");
  }
 }
 
-void Elf::Write(const std::filesystem::path& path, const std::vector<uint8_t>& code)
+void Elf::Write(const std::filesystem::path& path, const std::vector<uint8_t>& code, const std::vector<uint8_t>& data)
 {
- std::vector<uint8_t> data;
+ std::vector<uint8_t> elf;
 
- AddFileHeader(data, code);
- AddProgramHeader(data, code);
+ AddFileHeader(elf, code, data);
+ AddProgramHeaderCode(elf, code, data);
+ AddProgramHeaderData(elf, code, data);
 
 #ifdef PAD
- PadUpTo(data, PAD);
+ PadUpTo(elf, PAGE_SIZE);
 #endif
 
- data.insert(data.end(), code.begin(), code.end());
+ elf.insert(elf.end(), code.begin(), code.end());
+
+#ifdef PAD
+ PadUpTo(elf, PAGE_SIZE + paged_size(code));
+#endif
  
- std::string section_names("\0.text\0.shstrtab\0", 17);
- if (section_names.size() != 17)
+ elf.insert(elf.end(), data.begin(), data.end());
+
+ std::string section_names("\0.text\0.data\0.shstrtab\0", 23);
+ if (section_names.size() != 23)
   throw std::runtime_error("Bad size");
 
- size_t old_size {data.size()};
- data.resize(old_size + section_names.size());
- std::memcpy(data.data() + old_size, section_names.data(), section_names.size());
+ size_t old_size {elf.size()};
+ elf.resize(old_size + section_names.size());
+ std::memcpy(elf.data() + old_size, section_names.data(), section_names.size());
 
- AddSectionHeaderNull(data, code);
- AddSectionHeaderText(data, code);
- AddSectionHeaderSectionNames(data, code);
+ AddSectionHeaderNull(elf, code, data);
+ AddSectionHeaderText(elf, code, data);
+ AddSectionHeaderData(elf, code, data);
+ AddSectionHeaderSectionNames(elf, code, data);
 
- File::setFile(path, data);
+ File::setFile(path, elf);
 }
 
 std::vector<uint8_t> Elf::Read(const std::filesystem::path& path)
diff --git a/elf.h b/elf.h
index abcbca0..260f86f 100644
--- a/elf.h
+++ b/elf.h
@@ -42,10 +42,10 @@ struct FileHeader
  uint32_t e_flags{};
  uint16_t e_ehsize{0x40}; // size of this header
  uint16_t e_phentsize{0x38}; // size of program header
- uint16_t e_phnum{1}; // 1 program header
+ uint16_t e_phnum{2}; // program headers: .text, .data
  uint16_t e_shentsize{0x40}; // size of section header
- uint16_t e_shnum{3}; // number of section headers
- uint16_t e_shstrndx{2}; // 1st section header contains section names
+ uint16_t e_shnum{4}; // number of section headers
+ uint16_t e_shstrndx{3}; // 1st section header contains section names
 };
 
 // Segment
@@ -76,7 +76,7 @@ struct SectionHeader
 };
 #pragma pack(pop)
 
-void Write(const std::filesystem::path& path, const std::vector<uint8_t>&);
+void Write(const std::filesystem::path& path, const std::vector<uint8_t>& code, const std::vector<uint8_t>& data);
 std::vector<uint8_t> Read(const std::filesystem::path& path);
 }
 
diff --git a/test-elf.cpp b/test-elf.cpp
index d3fa9a9..46ef267 100644
--- a/test-elf.cpp
+++ b/test-elf.cpp
@@ -35,7 +35,7 @@ protected:
  }
  void TearDown(){
   std::error_code ec;
-  //fs::remove(TempFilename(), ec);
+  fs::remove(TempFilename(), ec);
  }
 };
 
@@ -44,13 +44,27 @@ TEST_F(ElfTest, read) {
 }
 #endif
 
-TEST_F(ElfTest, write) {
+TEST_F(ElfTest, write_code) {
  Elf::Write(TempFilename(),
             {
             0x48, 0xc7, 0xc0, 0x3c, 0x00, 0x00, 0x00, // mov    $0x3c,%rax     # syscall 60
             0x48, 0x31, 0xff,                         // xor    %rdi,%rdi      # exit code 0
             0x0f, 0x05,                               // syscall
-            });
+            }, {});
+
+ ASSERT_TRUE(fs::exists(TempFilename()));
+ ASSERT_GT(fs::file_size(TempFilename()), 0);
+}
+
+TEST_F(ElfTest, write_code_data) {
+ Elf::Write(TempFilename(),
+            {
+            0x48, 0xc7, 0xc0, 0x3c, 0x00, 0x00, 0x00, // mov    $0x3c,%rax     # syscall 60
+            0x48, 0x8b, 0x3c, 0x25, 0x00, 0x20, 0x40, // mov    0x402000,%rdi  # use value from data segment as exit code
+            0,
+            0x0f, 0x05,                               // syscall
+            },
+            {1, 0, 0, 0, 0, 0, 0, 0});
 
  ASSERT_TRUE(fs::exists(TempFilename()));
  ASSERT_GT(fs::file_size(TempFilename()), 0);