Add more atom-related syscalls (#211)

While trying to run some executables I noticed handlers for
`NtUserRegisterClassExWOW`/`NtUserUnregisterClass` were necessary, so
this PR aims to add them. Also, while working on this PR I noticed that
atoms should be ref-counted, so this PR also does that.

I found a previous abandoned PR that aimed to add it
(https://github.com/momo5502/emulator/pull/101), so I took some of the
feedback from that and added a common function in process_context to
handle the atoms, I hope that's ok.

src/windows-emulator/process_context.cpp

@@ -236,3 +236,79 @@ handle process_context::create_thread(memory_manager& memory, const uint64_t sta
     this->callbacks_->on_create_thread(h, *thr);
     return h;
 }
+
+uint16_t process_context::add_or_find_atom(std::u16string name)
+{
+    uint16_t index = 0;
+    if (!atoms.empty())
+    {
+        auto i = atoms.end();
+        --i;
+        index = i->first + 1;
+    }
+
+    std::optional<uint16_t> last_entry{};
+    for (auto& entry : atoms)
+    {
+        if (entry.second.name == name)
+        {
+            entry.second.ref_count++;
+            return entry.first;
+        }
+
+        if (entry.first > 0)
+        {
+            if (!last_entry)
+            {
+                index = 0;
+            }
+            else
+            {
+                const auto diff = entry.first - *last_entry;
+                if (diff > 1)
+                {
+                    index = *last_entry + 1;
+                }
+            }
+        }
+
+        last_entry = entry.first;
+    }
+
+    atoms[index] = {std::move(name), 1};
+
+    return index;
+}
+
+bool process_context::delete_atom(const std::u16string& name)
+{
+    for (auto it = atoms.begin(); it != atoms.end(); ++it)
+    {
+        if (it->second.name == name)
+        {
+            if (--it->second.ref_count == 0)
+            {
+                atoms.erase(it);
+            }
+            return true;
+        }
+    }
+
+    return false;
+}
+
+bool process_context::delete_atom(uint16_t atom_id)
+{
+    const auto it = atoms.find(atom_id);
+    if (it == atoms.end())
+    {
+        return false;
+    }
+
+    if (--it->second.ref_count == 0)
+    {
+        atoms.erase(it);
+    }
+
+    return true;
+}

src/windows-emulator/process_context.hpp

@@ -36,6 +36,20 @@ struct process_context
         utils::optional_function<void(handle h, emulator_thread& thr)> on_thread_terminated{};
     };
 
+    struct atom_entry
+    {
+        std::u16string name;
+        uint32_t ref_count = 0;
+
+        atom_entry(std::u16string n, uint32_t count)
+            : name(std::move(n)),
+              ref_count(count)
+        {
+        }
+
+        atom_entry() = default;
+    };
+
     process_context(x64_emulator& emu, memory_manager& memory, utils::clock& clock, callbacks& cb)
         : callbacks_(&cb),
           base_allocator(emu),
@@ -51,6 +65,10 @@ struct process_context
     handle create_thread(memory_manager& memory, uint64_t start_address, uint64_t argument, uint64_t stack_size,
                          bool suspended);
 
+    uint16_t add_or_find_atom(std::u16string name);
+    bool delete_atom(const std::u16string& name);
+    bool delete_atom(uint16_t atom_id);
+
     void serialize(utils::buffer_serializer& buffer) const;
     void deserialize(utils::buffer_deserializer& buffer);
 
@@ -86,7 +104,7 @@ struct process_context
     handle_store<handle_types::port, port> ports{};
     handle_store<handle_types::mutant, mutant> mutants{};
     handle_store<handle_types::registry, registry_key, 2> registry_keys{};
-    std::map<uint16_t, std::u16string> atoms{};
+    std::map<uint16_t, atom_entry> atoms{};
 
     std::vector<std::byte> default_register_set{};
 

src/windows-emulator/syscalls.cpp

@@ -530,47 +530,21 @@ namespace syscalls
 
         c.emu.read_memory(atom_name, name.data(), length);
 
-        uint16_t index = 0;
-        if (!c.proc.atoms.empty())
-        {
-            auto i = c.proc.atoms.end();
-            --i;
-            index = i->first + 1;
-        }
-
-        std::optional<uint16_t> last_entry{};
-        for (auto& entry : c.proc.atoms)
-        {
-            if (entry.second == name)
-            {
-                if (atom)
-                {
-                    atom.write(entry.first);
-                    return STATUS_SUCCESS;
-                }
-            }
-
-            if (entry.first > 0)
-            {
-                if (!last_entry)
-                {
-                    index = 0;
-                }
-                else
-                {
-                    const auto diff = entry.first - *last_entry;
-                    if (diff > 1)
-                    {
-                        index = *last_entry + 1;
-                    }
-                }
-            }
-
-            last_entry = entry.first;
-        }
-
-        c.proc.atoms[index] = std::move(name);
+        uint16_t index = c.proc.add_or_find_atom(name);
         atom.write(index);
+
+        return STATUS_SUCCESS;
+    }
+
+    NTSTATUS handle_NtAddAtom(const syscall_context& c, const uint64_t atom_name, const ULONG length,
+                              const emulator_object<RTL_ATOM> atom)
+    {
+        return handle_NtAddAtomEx(c, atom_name, length, atom, 0);
+    }
+
+    NTSTATUS handle_NtDeleteAtom(const syscall_context& c, const RTL_ATOM atom)
+    {
+        c.proc.delete_atom(atom);
         return STATUS_SUCCESS;
     }
 
@@ -653,6 +627,33 @@ namespace syscalls
     {
         return 0;
     }
+
+    template <typename Traits>
+    struct CLSMENUNAME
+    {
+        EMULATOR_CAST(typename Traits::PVOID, char*) pszClientAnsiMenuName;
+        EMULATOR_CAST(typename Traits::PVOID, char16_t*) pwszClientUnicodeMenuName;
+        EMULATOR_CAST(typename Traits::PVOID, UNICODE_STRING*) pusMenuName;
+    };
+
+    NTSTATUS handle_NtUserRegisterClassExWOW(
+        const syscall_context& c, const emulator_pointer /*wnd_class_ex*/,
+        const emulator_object<UNICODE_STRING<EmulatorTraits<Emu64>>> class_name,
+        const emulator_object<UNICODE_STRING<EmulatorTraits<Emu64>>> /*class_version*/,
+        const emulator_object<CLSMENUNAME<EmulatorTraits<Emu64>>> /*class_menu_name*/, const DWORD /*function_id*/,
+        const DWORD /*flags*/, const emulator_pointer /*wow*/)
+    {
+        uint16_t index = c.proc.add_or_find_atom(read_unicode_string(c.emu, class_name));
+        return index;
+    }
+
+    NTSTATUS handle_NtUserUnregisterClass(const syscall_context& c,
+                                          const emulator_object<UNICODE_STRING<EmulatorTraits<Emu64>>> class_name,
+                                          const emulator_pointer /*instance*/,
+                                          const emulator_object<CLSMENUNAME<EmulatorTraits<Emu64>>> /*class_menu_name*/)
+    {
+        return c.proc.delete_atom(read_unicode_string(c.emu, class_name));
+    }
 }
 
 void syscall_dispatcher::add_handlers(std::map<std::string, syscall_handler>& handler_mapping)
@@ -716,6 +717,8 @@ void syscall_dispatcher::add_handlers(std::map<std::string, syscall_handler>& ha
     add_handler(NtQueryInformationToken);
     add_handler(NtDxgkIsFeatureEnabled);
     add_handler(NtAddAtomEx);
+    add_handler(NtAddAtom);
+    add_handler(NtDeleteAtom);
     add_handler(NtInitializeNlsFiles);
     add_handler(NtUnmapViewOfSection);
     add_handler(NtUnmapViewOfSectionEx);
@@ -801,6 +804,8 @@ void syscall_dispatcher::add_handlers(std::map<std::string, syscall_handler>& ha
     add_handler(NtQueryFullAttributesFile);
     add_handler(NtFlushBuffersFile);
     add_handler(NtUserGetProcessWindowStation);
+    add_handler(NtUserRegisterClassExWOW);
+    add_handler(NtUserUnregisterClass);
 
 #undef add_handler
 }