LibELF: Consolidate main executable loading a bit

Merge the load_elf() and commit_elf() functions into a single
load_main_executable() function that takes care of both things.

Also split "stage 3" into two separate stages, keeping the lazy
relocations in stage 3, and adding a stage 4 for calling library
initialization functions.

We also make sure to map the main executable before dealing with
any of its dependencies, to ensure that non-PIE executables get
loaded at their desired address.

3 files changed, 70 insertions, 54 deletions

diff --git a/Userland/Libraries/LibELF/DynamicLinker.cpp b/Userland/Libraries/LibELF/DynamicLinker.cpp
index 827dd9de12..89fa43ac17 100644
--- a/Userland/Libraries/LibELF/DynamicLinker.cpp
+++ b/Userland/Libraries/LibELF/DynamicLinker.cpp
@@ -31,6 +31,7 @@
 #include <AK/HashTable.h>
 #include <AK/LexicalPath.h>
 #include <AK/LogStream.h>
+#include <AK/NonnullRefPtrVector.h>
 #include <AK/ScopeGuard.h>
 #include <LibC/mman.h>
 #include <LibC/unistd.h>
@@ -50,7 +51,7 @@ namespace {
 HashMap<String, NonnullRefPtr<ELF::DynamicLoader>> g_loaders;
 Vector<NonnullRefPtr<ELF::DynamicObject>> g_global_objects;
 
-using MainFunction = int (*)(int, char**, char**);
+using EntryPointFunction = int (*)(int, char**, char**);
 using LibCExitFunction = void (*)(int);
 
 size_t g_current_tls_offset = 0;
@@ -180,62 +181,69 @@ static void initialize_libc(DynamicObject& libc)
 }
 
 template<typename Callback>
-static void for_each_dependency_of_impl(const String& name, HashTable<String>& seen_names, Callback callback)
+static void for_each_dependency_of(const String& name, HashTable<String>& seen_names, Callback callback)
 {
     if (seen_names.contains(name))
         return;
     seen_names.set(name);
 
     for (const auto& needed_name : get_dependencies(name))
-        for_each_dependency_of_impl(get_library_name(needed_name), seen_names, callback);
+        for_each_dependency_of(get_library_name(needed_name), seen_names, callback);
 
     callback(*g_loaders.get(name).value());
 }
 
-template<typename Callback>
-static void for_each_dependency_of(const String& name, Callback callback)
+static NonnullRefPtrVector<DynamicLoader> collect_loaders_for_executable(const String& name)
 {
     HashTable<String> seen_names;
-    for_each_dependency_of_impl(name, seen_names, move(callback));
+    NonnullRefPtrVector<DynamicLoader> loaders;
+    for_each_dependency_of(name, seen_names, [&](auto& loader) {
+        loaders.append(loader);
+    });
+    return loaders;
 }
 
-static void load_elf(const String& name)
+static NonnullRefPtr<DynamicLoader> load_main_executable(const String& name)
 {
-    for_each_dependency_of(name, [](auto& loader) {
+    // NOTE: We always map the main executable first, since it may require
+    //       placement at a specific address.
+    auto& main_executable_loader = *g_loaders.get(name).value();
+    auto main_executable_object = main_executable_loader.map();
+    g_global_objects.append(*main_executable_object);
+
+    auto loaders = collect_loaders_for_executable(name);
+
+    for (auto& loader : loaders) {
         auto dynamic_object = loader.map();
-        VERIFY(dynamic_object);
-        g_global_objects.append(*dynamic_object);
-    });
-    for_each_dependency_of(name, [](auto& loader) {
+        if (dynamic_object)
+            g_global_objects.append(*dynamic_object);
+    }
+
+    for (auto& loader : loaders) {
         bool success = loader.link(RTLD_GLOBAL | RTLD_LAZY, g_total_tls_size);
         VERIFY(success);
-    });
-}
-
-static NonnullRefPtr<DynamicLoader> commit_elf(const String& name)
-{
-    auto loader = g_loaders.get(name).value();
-    for (const auto& needed_name : get_dependencies(name)) {
-        String library_name = get_library_name(needed_name);
-        if (g_loaders.contains(library_name)) {
-            commit_elf(library_name);
-        }
     }
 
-    auto object = loader->load_stage_3(RTLD_GLOBAL | RTLD_LAZY, g_total_tls_size);
-    VERIFY(object);
+    for (auto& loader : loaders) {
+        auto object = loader.load_stage_3(RTLD_GLOBAL | RTLD_LAZY, g_total_tls_size);
+        VERIFY(object);
 
-    if (name == "libsystem.so") {
-        if (syscall(SC_msyscall, object->base_address().as_ptr())) {
-            VERIFY_NOT_REACHED();
+        if (loader.filename() == "libsystem.so") {
+            if (syscall(SC_msyscall, object->base_address().as_ptr())) {
+                VERIFY_NOT_REACHED();
+            }
+        }
+
+        if (loader.filename() == "libc.so") {
+            initialize_libc(*object);
         }
     }
 
-    if (name == "libc.so") {
-        initialize_libc(*object);
+    for (auto& loader : loaders) {
+        loader.load_stage_4();
     }
-    g_loaders.remove(name);
-    return loader;
+
+    return main_executable_loader;
 }
 
 static void read_environment_variables()
@@ -265,33 +273,26 @@ void ELF::DynamicLinker::linker_main(String&& main_program_name, int main_progra
 
     allocate_tls();
 
-    load_elf(main_program_name);
-
-    // NOTE: We put this in a RefPtr instead of a NonnullRefPtr so we can release it later.
-    RefPtr main_program_lib = commit_elf(main_program_name);
+    auto entry_point_function = [&main_program_name] {
+        auto main_executable_loader = load_main_executable(main_program_name);
+        auto entry_point = main_executable_loader->image().entry();
+        if (main_executable_loader->is_dynamic())
+            entry_point = entry_point.offset(main_executable_loader->text_segment_load_address().get());
+        return (EntryPointFunction)(entry_point.as_ptr());
+    }();
 
-    FlatPtr entry_point = reinterpret_cast<FlatPtr>(main_program_lib->image().entry().as_ptr());
-    if (main_program_lib->is_dynamic())
-        entry_point += reinterpret_cast<FlatPtr>(main_program_lib->text_segment_load_address().as_ptr());
-
-    dbgln_if(DYNAMIC_LOAD_DEBUG, "entry point: {:p}", (void*)entry_point);
     g_loaders.clear();
 
-    MainFunction main_function = (MainFunction)(entry_point);
-    dbgln_if(DYNAMIC_LOAD_DEBUG, "jumping to main program entry point: {:p}", main_function);
-    if (g_do_breakpoint_trap_before_entry) {
-        asm("int3");
-    }
-
-    // Unmap the main executable and release our related resources.
-    main_program_lib = nullptr;
-
     int rc = syscall(SC_msyscall, nullptr);
     if (rc < 0) {
         VERIFY_NOT_REACHED();
     }
 
-    rc = main_function(argc, argv, envp);
+    dbgln_if(DYNAMIC_LOAD_DEBUG, "Jumping to entry point: {:p}", entry_point_function);
+    if (g_do_breakpoint_trap_before_entry) {
+        asm("int3");
+    }
+    rc = entry_point_function(argc, argv, envp);
     dbgln_if(DYNAMIC_LOAD_DEBUG, "rc: {}", rc);
     if (g_libc_exit != nullptr) {
         g_libc_exit(rc);
diff --git a/Userland/Libraries/LibELF/DynamicLoader.cpp b/Userland/Libraries/LibELF/DynamicLoader.cpp
index 880bdd420d..682da8642b 100644
--- a/Userland/Libraries/LibELF/DynamicLoader.cpp
+++ b/Userland/Libraries/LibELF/DynamicLoader.cpp
@@ -154,7 +154,10 @@ void* DynamicLoader::symbol_for_name(const StringView& name)
 
 RefPtr<DynamicObject> DynamicLoader::map()
 {
-    VERIFY(!m_dynamic_object);
+    if (m_dynamic_object) {
+        // Already mapped.
+        return nullptr;
+    }
 
     if (!m_valid) {
         dbgln("DynamicLoader::map failed: image is invalid");
@@ -245,10 +248,14 @@ RefPtr<DynamicObject> DynamicLoader::load_stage_3(unsigned flags, size_t total_t
 #endif
     }
 
-    call_object_init_functions();
     return m_dynamic_object;
 }
 
+void DynamicLoader::load_stage_4()
+{
+    call_object_init_functions();
+}
+
 void DynamicLoader::do_lazy_relocations(size_t total_tls_size)
 {
     for (const auto& relocation : m_unresolved_relocations) {
diff --git a/Userland/Libraries/LibELF/DynamicLoader.h b/Userland/Libraries/LibELF/DynamicLoader.h
index cc85270c2b..4e63af6366 100644
--- a/Userland/Libraries/LibELF/DynamicLoader.h
+++ b/Userland/Libraries/LibELF/DynamicLoader.h
@@ -43,6 +43,8 @@ public:
     static RefPtr<DynamicLoader> try_create(int fd, String filename);
     ~DynamicLoader();
 
+    const String& filename() const { return m_filename; }
+
     bool is_valid() const { return m_valid; }
 
     // Load a full ELF image from file into the current process and create an DynamicObject
@@ -55,8 +57,12 @@ public:
     // Stage 2 of loading: dynamic object loading and primary relocations
     bool load_stage_2(unsigned flags, size_t total_tls_size);
 
-    // Stage 3 of loading: lazy relocations and initializers
+    // Stage 3 of loading: lazy relocations
     RefPtr<DynamicObject> load_stage_3(unsigned flags, size_t total_tls_size);
+
+    // Stage 4 of loading: initializers
+    void load_stage_4();
+
     // Intended for use by dlsym or other internal methods
     void* symbol_for_name(const StringView&);
 
@@ -111,6 +117,8 @@ private:
     // Stage 3
     void do_lazy_relocations(size_t total_tls_size);
     void setup_plt_trampoline();
+
+    // Stage 4
     void call_object_init_functions();
 
     bool validate();