AK+LibC+LibPthread: Introduce NoAllocationGuard

NoAllocationGuard is an RAII stack guard that prevents allocations
while it exists. This is done through a thread-local global flag which
causes malloc to crash on a VERIFY if it is false. The guard allows for
recursion.

The intended use case for this class is in real-time audio code. In such
code, allocations are really bad, and this is an easy way of dynamically
enforcing the no-allocations rule while giving the user good feedback if
it is violated. Before real-time audio code is executed, e.g. in LibDSP,
a NoAllocationGuard is instantiated. This is not done with this commit,
as currently some code in LibDSP may still incorrectly allocate in real-
time situations.

Other use cases for the Kernel have also been added, so this commit
builds on the previous to add the support both in Userland and in the
Kernel.

AK/NoAllocationGuard.h

@@ -0,0 +1,61 @@
+/*
+ * Copyright (c) 2022, kleines Filmröllchen <malu.bertsch@gmail.com>.
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#pragma once
+
+#include <AK/Forward.h>
+#include <AK/Noncopyable.h>
+
+#if defined(KERNEL)
+#    include <Kernel/Arch/Processor.h>
+#    include <Kernel/Heap/kmalloc.h>
+#else
+#    include <LibC/mallocdefs.h>
+#endif
+
+namespace AK {
+
+class NoAllocationGuard {
+    AK_MAKE_NONCOPYABLE(NoAllocationGuard);
+    AK_MAKE_NONMOVABLE(NoAllocationGuard);
+
+public:
+    NoAllocationGuard()
+        : m_allocation_enabled_previously(get_thread_allocation_state())
+    {
+        set_thread_allocation_state(false);
+    }
+
+    ~NoAllocationGuard()
+    {
+        set_thread_allocation_state(m_allocation_enabled_previously);
+    }
+
+private:
+    static bool get_thread_allocation_state()
+    {
+#if defined(KERNEL)
+        return Processor::current_thread()->get_allocation_enabled();
+#else
+        return s_allocation_enabled;
+#endif
+    }
+
+    static void set_thread_allocation_state(bool value)
+    {
+#if defined(KERNEL)
+        Processor::current_thread()->set_allocation_enabled(value);
+#else
+        s_allocation_enabled = value;
+#endif
+    }
+
+    bool m_allocation_enabled_previously;
+};
+
+}
+
+using AK::NoAllocationGuard;

Userland/Libraries/LibC/malloc.cpp

@@ -198,8 +198,18 @@ enum class CallerWillInitializeMemory {
     Yes,
 };
 
+#ifndef NO_TLS
+// HACK: This is a __thread - marked thread-local variable. If we initialize it globally here, VERY weird errors happen.
+// The initialization happens in __malloc_init() and pthread_create_helper().
+__thread bool s_allocation_enabled;
+#endif
+
 static void* malloc_impl(size_t size, CallerWillInitializeMemory caller_will_initialize_memory)
 {
+#ifndef NO_TLS
+    VERIFY(s_allocation_enabled);
+#endif
+
     if (s_log_malloc)
         dbgln("LibC: malloc({})", size);
 
@@ -330,6 +340,10 @@ static void* malloc_impl(size_t size, CallerWillInitializeMemory caller_will_ini
 
 static void free_impl(void* ptr)
 {
+#ifndef NO_TLS
+    VERIFY(s_allocation_enabled);
+#endif
+
     ScopedValueRollback rollback(errno);
 
     if (!ptr)
@@ -534,6 +548,12 @@ void* realloc(void* ptr, size_t size)
 
 void __malloc_init()
 {
+#ifndef NO_TLS
+    // HACK: This is a __thread - marked thread-local variable. If we initialize it globally, VERY weird errors happen.
+    // Therefore, we need to do the initialization here and in pthread_create_helper().
+    s_allocation_enabled = true;
+#endif
+
     s_in_userspace_emulator = (int)syscall(SC_emuctl, 0) != -ENOSYS;
     if (s_in_userspace_emulator) {
         // Don't bother scrubbing memory if we're running in UE since it

Userland/Libraries/LibC/mallocdefs.h

@@ -19,6 +19,12 @@
 static constexpr unsigned short size_classes[] = { 16, 32, 64, 128, 256, 496, 1008, 2032, 4080, 8176, 16368, 32752, 0 };
 static constexpr size_t num_size_classes = (sizeof(size_classes) / sizeof(unsigned short)) - 1;
 
+#ifndef NO_TLS
+extern "C" {
+extern __thread bool s_allocation_enabled;
+}
+#endif
+
 consteval bool check_size_classes_alignment()
 {
     for (size_t i = 0; i < num_size_classes; i++) {

Userland/Libraries/LibPthread/pthread.cpp

@@ -14,6 +14,7 @@
 #include <bits/pthread_integration.h>
 #include <errno.h>
 #include <limits.h>
+#include <mallocdefs.h>
 #include <pthread.h>
 #include <serenity.h>
 #include <signal.h>
@@ -43,6 +44,9 @@ extern "C" {
 
 static void* pthread_create_helper(void* (*routine)(void*), void* argument, void* stack_location, size_t stack_size)
 {
+    // HACK: This is a __thread - marked thread-local variable. If we initialize it globally, VERY weird errors happen.
+    // Therefore, we need to do the initialization here and in __malloc_init().
+    s_allocation_enabled = true;
     s_stack_location = stack_location;
     s_stack_size = stack_size;
     void* ret_val = routine(argument);