WIP

include/LIEF/iostream.hpp

@@ -46,8 +46,14 @@ class vector_iostream {
     return write(sp.data(), sp.size());
   }
 
-  vector_iostream& write(std::vector<uint8_t> s);
-  vector_iostream& write(const std::string& s);
+  vector_iostream& write(std::vector<uint8_t> s) {
+    return write(s.data(), s.size());
+  }
+
+  vector_iostream& write(const std::string& s) {
+    return write(reinterpret_cast<const uint8_t*>(s.c_str()), s.size() + 1);
+  }
+
   vector_iostream& write(size_t count, uint8_t value) {
     raw_.insert(std::end(raw_), count, value);
     current_pos_ += count;

src/ELF/Builder.cpp

@@ -86,6 +86,12 @@ bool Builder::should_swap() const {
 
 
 void Builder::build() {
+  const Header::CLASS elf_class = binary_->type();
+  if (elf_class != Header::CLASS::ELF32 && elf_class != Header::CLASS::ELF64) {
+    LIEF_ERR("Invalid ELF class");
+    return;
+  }
+
   auto res = binary_->type() == Header::CLASS::ELF32 ?
              build<details::ELF32>() : build<details::ELF64>();
   if (!res) {

src/ELF/Section.cpp

@@ -199,6 +199,11 @@ span<const uint8_t> Section::content() const {
   }
   const std::vector<uint8_t>& binary_content = datahandler_->content();
   DataHandler::Node& node = res.value();
+  auto end_offset = (int64_t)node.offset() + (int64_t)node.size();
+  if (end_offset <= 0 || end_offset > (int64_t)binary_content.size()) {
+    return {};
+  }
+
   const uint8_t* ptr = binary_content.data() + node.offset();
   return {ptr, ptr + node.size()};
 }
@@ -248,6 +253,12 @@ void Section::content(const std::vector<uint8_t>& data) {
               data.size(), name(), node.size());
   }
 
+  auto max_offset = (int64_t)node.offset() + (int64_t)data.size();
+  if (max_offset < 0 || max_offset > (int64_t)binary_content.size()) {
+    LIEF_ERR("Write out of range");
+    return;
+  }
+
   size(data.size());
 
   std::copy(std::begin(data), std::end(data),
@@ -292,6 +303,12 @@ void Section::content(std::vector<uint8_t>&& data) {
 
   size(data.size());
 
+  auto max_offset = (int64_t)node.offset() + (int64_t)data.size();
+  if (max_offset < 0 || max_offset > (int64_t)binary_content.size()) {
+    LIEF_ERR("Write out of range");
+    return;
+  }
+
   std::move(std::begin(data), std::end(data),
             std::begin(binary_content) + node.offset());
 }

src/ELF/Segment.cpp

@@ -331,6 +331,12 @@ void Segment::content(std::vector<uint8_t> content) {
                 content.size(), to_string(type()), virtual_size(), node.size());
   }
 
+  auto max_offset = (int64_t)node.offset() + (int64_t)content.size();
+  if (max_offset < 0 || max_offset > (int64_t)binary_content.size()) {
+    LIEF_ERR("Write out of range");
+    return;
+  }
+
   physical_size(node.size());
 
   std::move(std::begin(content), std::end(content),

src/iostream.cpp

@@ -54,44 +54,12 @@ vector_iostream& vector_iostream::write(const uint8_t* s, std::streamsize n) {
   if (raw_.size() < (pos + n)) {
     raw_.resize(pos + n);
   }
-
-  auto it = std::begin(raw_);
-  std::advance(it, pos);
-  std::copy(s, s + n, it);
-
+  std::copy(s, s + n, raw_.data() + pos);
   current_pos_ += n;
-  return *this;
-}
-
-vector_iostream& vector_iostream::write(std::vector<uint8_t> s) {
-  const auto pos = static_cast<size_t>(tellp());
-  if (raw_.size() < (pos + s.size())) {
-    raw_.resize(pos + s.size());
-  }
-
-  auto it = std::begin(raw_);
-  std::advance(it, pos);
-  std::move(std::begin(s), std::end(s), it);
 
-  current_pos_ += s.size();
   return *this;
 }
 
-vector_iostream& vector_iostream::write(const std::string& s) {
-  const auto pos = static_cast<size_t>(tellp());
-  if (raw_.size() < (pos + s.size() + 1)) {
-    raw_.resize(pos + s.size() + 1);
-  }
-
-  auto it = std::begin(raw_);
-  std::advance(it, pos);
-  std::copy(std::begin(s), std::end(s), it);
-
-  current_pos_ += s.size() + 1;
-  return *this;
-}
-
-
 vector_iostream& vector_iostream::write_uleb128(uint64_t value) {
   uint8_t byte;
   do {