From 4e4600475bd73733db9498b15e2568e54ba36ed5 Mon Sep 17 00:00:00 2001
From: "Jonathan B. Coe" <jonathanbcoe@gmail.com>
Date: Sat, 21 Mar 2026 23:00:00 +0000
Subject: [PATCH] Allow sanitizers to be enabled in CMakeLists

---
 ASAN.md                  | 61 ++++++++++++++++++++++++++++++++++++++++
 CMakeLists.txt           | 30 ++++++++++++++------
 nle/tests/test_system.py | 18 ++++++++++++
 3 files changed, 100 insertions(+), 9 deletions(-)
 create mode 100644 ASAN.md

diff --git a/ASAN.md b/ASAN.md
new file mode 100644
index 000000000..e38a61e8e
--- /dev/null
+++ b/ASAN.md
@@ -0,0 +1,61 @@
+# Address Sanitizer (ASAN) Guide
+
+This project supports Address Sanitizer (ASAN) to help detect memory corruption,
+use-after-free, and buffer overflows in the C/C++ NetHack engine and its Python
+extensions.
+
+## Enabling ASAN
+
+ASAN is integrated into the CMake build system and can be enabled by editing
+`pyproject.toml`.
+
+### Current Configuration
+
+```toml
+[tool.scikit-build]
+cmake.build-type = "Release"
+cmake.args = ["-DHACKDIR=nle/nethackdir", "-DPYTHON_PACKAGE_NAME=nle"]
+```
+
+To enable ASAN, add the cmake argument `-DENABLE_ASAN=On` and switch
+`cmake.build-type` to `Debug`.
+
+## Running Tests with ASAN
+
+Because the Python interpreter itself is not built with ASAN, you must preload
+the ASAN runtime library when running tests.
+
+```bash
+LD_PRELOAD=$(gcc -print-file-name=libasan.so):$(gcc -print-file-name=libstdc++.so) \
+ASAN_OPTIONS=detect_leaks=0 \
+uv run pytest
+```
+
+_Note: Preloading `libstdc++.so` may be necessary on some platforms (like
+aarch64 Linux) to avoid crashes when C++ exceptions are thrown._
+
+### Why `detect_leaks=0`?
+
+We disable the LeakSanitizer (`detect_leaks=0`) for several reasons:
+
+1. Python Shutdown: CPython does not free all memory at exit (e.g., global
+   singletons, interned strings). This is intentional for performance but is
+   flagged as a "leak" by ASAN.
+2. Pytest State: `pytest` keeps tracebacks, local variables, and fixture data in
+   memory until the end of the session to generate reports.
+3. Standard Interpreter: Since we are running a sanitized C extension inside a
+   non-sanitized Python interpreter, the leak detector cannot accurately track
+   the ownership boundary between the two.
+
+Disabling leak detection still allows ASAN to catch critical memory corruption
+errors (Buffer Overflows, Use-After-Free, etc.) as they happen.
+
+## Other Sanitizers
+
+The build system also supports:
+
+- Thread Sanitizer (TSAN): Use `-DENABLE_TSAN=ON`.
+- Undefined Behavior Sanitizer (UBSAN): Use `-DENABLE_UBSAN=ON`.
+
+To use these, update `pyproject.toml` accordingly and preload the corresponding
+library (e.g., `libtsan.so`).
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 48a45e4b8..0337a7eb5 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -19,16 +19,28 @@ if(CMAKE_BUILD_TYPE MATCHES Debug)
   set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -g")
   set(CMAKE_XCODE_ATTRIBUTE_DEBUG_INFORMATION_FORMAT "dwarf-with-dsym")
 
-  if(0)
-    # address sanitizer.
-    set(CMAKE_CXX_FLAGS_DEBUG
-        "${CMAKE_CXX_FLAGS_DEBUG} -fno-omit-frame-pointer -fsanitize=address")
-    set(CMAKE_C_FLAGS_DEBUG
-        "${CMAKE_C_FLAGS_DEBUG} -fno-omit-frame-pointer -fsanitize=address")
-    set(CMAKE_LINKER_FLAGS_DEBUG
-        "${CMAKE_LINKER_FLAGS_DEBUG} -fno-omit-frame-pointer -fsanitize=address"
-    )
+  option(ENABLE_ASAN "Enable Address Sanitizer" OFF)
+  option(ENABLE_TSAN "Enable Thread Sanitizer" OFF)
+  option(ENABLE_UBSAN "Enable Undefined Behavior Sanitizer" OFF)
+
+  if(ENABLE_ASAN)
+    message(STATUS "Enabling Address Sanitizer")
+    add_compile_options(-fsanitize=address -fno-omit-frame-pointer)
+    add_link_options(-fsanitize=address)
   endif()
+
+  if(ENABLE_TSAN)
+    message(STATUS "Enabling Thread Sanitizer")
+    add_compile_options(-fsanitize=thread)
+    add_link_options(-fsanitize=thread)
+  endif()
+
+  if(ENABLE_UBSAN)
+    message(STATUS "Enabling Undefined Behavior Sanitizer")
+    add_compile_options(-fsanitize=undefined)
+    add_link_options(-fsanitize=undefined)
+  endif()
+
   if(MSVC)
     add_compile_options(/W4)
   else()
diff --git a/nle/tests/test_system.py b/nle/tests/test_system.py
index 9fe5e60be..9cc1de4f4 100644
--- a/nle/tests/test_system.py
+++ b/nle/tests/test_system.py
@@ -1,4 +1,7 @@
 # Copyright (c) Facebook, Inc. and its affiliates.
+
+import ctypes
+import functools
 import multiprocessing as mp
 import queue
 import random
@@ -19,17 +22,32 @@ def new_env_one_step():
     return terminated
 
 
+@functools.cache
+def is_asan():
+    """Checks if the process is running with ASAN.
+
+    See if the __asan_init symbol is present in the current process.
+    """
+
+    current_process = ctypes.CDLL(None)
+    return hasattr(current_process, "__asan_init")
+
+
 @pytest.mark.parametrize(
     "ctx", [mp.get_context(m) for m in START_METHODS], ids=START_METHODS
 )
 class TestEnvSubprocess:
     def test_env_in_subprocess(self, ctx):
+        if ctx.get_start_method() == "spawn" and is_asan():
+            pytest.skip("ASAN crashes on spawn on this environment")
         p = ctx.Process(target=new_env_one_step)
         p.start()
         p.join()
         assert p.exitcode == 0
 
     def test_env_before_and_in_subprocess(self, ctx):
+        if ctx.get_start_method() == "spawn" and is_asan():
+            pytest.skip("ASAN crashes on spawn on this environment")
         new_env_one_step()
         p = ctx.Process(target=new_env_one_step)
         p.start()