Fix crash in setting new resource table #1

Author: dominicletz

lib/libpe.ex

@@ -263,17 +263,26 @@ defmodule LibPE do
         LibPE.ResourceTable.parse(virtual_data, virtual_address)
 
       nil ->
-        nil
+        %LibPE.ResourceTable{}
     end
   end
 
-  def set_resources(%LibPE{coff_sections: sections} = pe, resources = %LibPE.ResourceTable{}) do
+  defp ensure_resource_section(%LibPE{coff_sections: sections} = pe) do
+    idx = Enum.find_index(sections, fn %LibPE.Section{name: name} -> name == ".rsrc" end)
+
+    if idx == nil do
+      %LibPE{pe | coff_sections: sections ++ [%LibPE.Section{name: ".rsrc"}]}
+    else
+      pe
+    end
+  end
+
+  def set_resources(pe, resources = %LibPE.ResourceTable{}) do
     # need to ensure that the virtual_address is up-to-date
-    pe = update_layout(pe)
+    %LibPE{coff_sections: sections} = pe = update_layout(ensure_resource_section(pe))
 
     # now fetching and setting the resource
     idx = Enum.find_index(sections, fn %LibPE.Section{name: name} -> name == ".rsrc" end)
-
     section = %LibPE.Section{virtual_address: virtual_address} = Enum.at(sections, idx)
     data = LibPE.ResourceTable.encode(resources, virtual_address)
     section = %LibPE.Section{section | virtual_data: data}

lib/libpe/section.ex

@@ -2,21 +2,19 @@ defmodule LibPE.Section do
   @moduledoc false
   alias LibPE.Section
 
-  defstruct [
-    :name,
-    :padding,
-    :virtual_data,
-    :virtual_size,
-    :virtual_address,
-    :raw_data,
-    :size_of_raw_data,
-    :pointer_to_raw_data,
-    :pointer_to_relocations,
-    :pointer_to_linenumbers,
-    :number_of_relocations,
-    :number_of_linenumbers,
-    :flags
-  ]
+  defstruct name: nil,
+            padding: "\0",
+            virtual_data: "",
+            virtual_size: 0,
+            virtual_address: 0,
+            raw_data: "",
+            size_of_raw_data: 0,
+            pointer_to_raw_data: 0,
+            pointer_to_relocations: 0,
+            pointer_to_linenumbers: 0,
+            number_of_relocations: 0,
+            number_of_linenumbers: 0,
+            flags: 0
 
   def parse(rest, number, full_image) do
     List.duplicate(nil, number)

test/hello.exe

@@ -44,6 +44,24 @@ defmodule LibPETest do
     end
   end
 
+  test "set icon" do
+    {:ok, pe} = LibPE.parse_file("test/hello.exe")
+
+    resource_table = LibPE.get_resources(pe)
+
+    data = File.read!("test/logo.ico")
+    type = LibPE.ResourceTypes.encode("RT_ICON")
+    resource_table = LibPE.ResourceTable.set_resource(resource_table, type, data)
+
+    raw =
+      LibPE.set_resources(pe, resource_table)
+      |> LibPE.update_layout()
+      |> LibPE.update_checksum()
+      |> LibPE.encode()
+
+    File.write!("test/hello-out.exe", raw)
+  end
+
   # defp tip(rsrc) do
   #   clean_data(hd(rsrc.entries))
   # end