fix(asm): fix memory corruption in waf interface [backport #5260 to 1.9] (#5267)

Backport of #5260 to 1.9

Author: christophe-papazian

ddtrace/appsec/ddwaf/__init__.py

@@ -18,18 +18,17 @@
 
 try:
     from .ddwaf_types import ddwaf_config
-    from .ddwaf_types import ddwaf_context_destroy
-    from .ddwaf_types import ddwaf_context_init
-    from .ddwaf_types import ddwaf_destroy
+    from .ddwaf_types import ddwaf_context_capsule
     from .ddwaf_types import ddwaf_get_version
-    from .ddwaf_types import ddwaf_init
     from .ddwaf_types import ddwaf_object
+    from .ddwaf_types import ddwaf_object_free
     from .ddwaf_types import ddwaf_result
-    from .ddwaf_types import ddwaf_result_free
     from .ddwaf_types import ddwaf_ruleset_info
     from .ddwaf_types import ddwaf_run
-    from .ddwaf_types import ddwaf_update
+    from .ddwaf_types import py_ddwaf_context_init
+    from .ddwaf_types import py_ddwaf_init
     from .ddwaf_types import py_ddwaf_required_addresses
+    from .ddwaf_types import py_ddwaf_update
 
     _DDWAF_LOADED = True
 except OSError:
@@ -82,16 +81,18 @@ def __init__(self, ruleset_map, obfuscation_parameter_key_regexp, obfuscation_pa
                 key_regex=obfuscation_parameter_key_regexp, value_regex=obfuscation_parameter_value_regexp
             )
             self._info = ddwaf_ruleset_info()
-            self._ruleset_map = ddwaf_object.create_without_limits(ruleset_map)
-            self._handle = ddwaf_init(self._ruleset_map, ctypes.byref(config), ctypes.byref(self._info))
-            self._ctx = 0
+            ruleset_map_object = ddwaf_object.create_without_limits(ruleset_map)
+            self._handle = py_ddwaf_init(ruleset_map_object, ctypes.byref(config), ctypes.byref(self._info))
             if not self._handle or self._info.failed:
+                # We keep the handle alive in case of errors, as some valid rules can be loaded
+                # at the same time some invalid ones are rejected
                 LOGGER.error(
                     "DDWAF.__init__: invalid rules\n ruleset: %s\nloaded:%s\nerrors:%s\n",
-                    self._ruleset_map.struct,
+                    ruleset_map_object.struct,
                     self._info.loaded,
                     self.info.errors,
                 )
+            ddwaf_object_free(ctypes.byref(ruleset_map_object))
 
         @property
         def required_data(self):
@@ -110,63 +111,52 @@ def update_rules(self, new_rules):
             # type: (dict[text_type, DDWafRulesType]) -> bool
             """update the rules of the WAF instance. return True if an error occurs."""
             rules = ddwaf_object.create_without_limits(new_rules)
-            result = ddwaf_update(self._handle, rules, ctypes.byref(self._info))
-            if result == 0:
-                LOGGER.error("DDWAF.update_rules: invalid rules")
-                return True
-            else:
+            result = py_ddwaf_update(self._handle, rules, self._info)
+            ddwaf_object_free(rules)
+            if result:
                 LOGGER.debug("DDWAF.update_rules success.\ninfo %s", self.info)
                 self._handle = result
+                return True
+            else:
+                LOGGER.debug("DDWAF.update_rules: keeping the previous handle.")
                 return False
 
         def _at_request_start(self):
-            if self._ctx:
-                ddwaf_context_destroy(self._ctx)
-            self._ctx = ddwaf_context_init(self._handle)
-            if self._ctx == 0:
-                LOGGER.error("DDWaf failure to create the context")
+            # type: () -> ddwaf_context_capsule
+            if self._handle:
+                ctx = py_ddwaf_context_init(self._handle)
+            if not ctx:
+                LOGGER.error("DDWaf._at_request_start: failure to create the context.")
+            return ctx
 
         def _at_request_end(self):
-            if self._ctx:
-                ddwaf_context_destroy(self._ctx)
-                self._ctx = 0
+            # () -> None
+            pass
 
         def run(
             self,  # type: DDWaf
+            ctx,  # type: ddwaf_context_capsule
             data,  # type: DDWafRulesType
             timeout_ms=DEFAULT_DDWAF_TIMEOUT_MS,  # type:int
         ):
             # type: (...) -> DDWaf_result
             start = time.time()
 
-            if self._ctx == 0:
-                LOGGER.warning("DDWaf failsafe to create the context")
-                self._ctx = ddwaf_context_init(self._handle)
-
-            if self._ctx == 0:
-                LOGGER.error("DDWaf failure: no context created")
+            if not ctx:
+                LOGGER.error("DDWaf.run: dry run. no context created.")
                 return DDWaf_result(None, [], 0, (time.time() - start) * 1e6)
 
             result = ddwaf_result()
             wrapper = ddwaf_object(data)
-            error = ddwaf_run(self._ctx, wrapper, ctypes.byref(result), timeout_ms * 1000)
+            error = ddwaf_run(ctx.ctx, wrapper, ctypes.byref(result), timeout_ms * 1000)
             if error < 0:
                 LOGGER.warning("run DDWAF error: %d\ninput %s\nerror %s", error, wrapper.struct, self.info.errors)
-            try:
-                return DDWaf_result(
-                    result.data.decode("UTF-8", errors="ignore") if hasattr(result, "data") and result.data else None,
-                    [result.actions.array[i].decode("UTF-8", errors="ignore") for i in range(result.actions.size)],
-                    result.total_runtime / 1e3,
-                    (time.time() - start) * 1e6,
-                )
-            finally:
-                ddwaf_result_free(ctypes.byref(result))
-
-        def __dealloc__(self):
-            if self._ctx:
-                ddwaf_context_destroy(self._ctx)
-            if self._handle:
-                ddwaf_destroy(self._handle)
+            return DDWaf_result(
+                result.data.decode("UTF-8", errors="ignore") if hasattr(result, "data") and result.data else None,
+                [result.actions.array[i].decode("UTF-8", errors="ignore") for i in range(result.actions.size)],
+                result.total_runtime / 1e3,
+                (time.time() - start) * 1e6,
+            )
 
     def version():
         # type: () -> text_type
@@ -185,13 +175,25 @@ def __init__(self, rules, obfuscation_parameter_key_regexp, obfuscation_paramete
 
         def run(
             self,  # type: DDWaf
+            ctx,  # type: ddwaf_context_capsule
             data,  # type: Union[None, int, text_type, list[Any], dict[text_type, Any]]
             timeout_ms=DEFAULT_DDWAF_TIMEOUT_MS,  # type:int
         ):
             # type: (...) -> DDWaf_result
             LOGGER.warning("DDWaf features disabled. dry run")
             return DDWaf_result(None, [], 0.0, 0.0)
 
+        def update_rules(self, _):
+            # type: (dict[text_type, DDWafRulesType]) -> bool
+            LOGGER.warning("DDWaf features disabled. dry update")
+            return False
+
+        def _at_request_start(self):
+            pass
+
+        def _at_request_end(self):
+            pass
+
     def version():
         # type: () -> text_type
         LOGGER.warning("DDWaf features disabled. null version")

ddtrace/appsec/ddwaf/ddwaf_types.py

@@ -247,6 +247,12 @@ def __repr__(self):
             self.actions,
         )
 
+    def __del__(self):
+        try:
+            ddwaf_result_free(self)
+        except TypeError:
+            pass
+
 
 ddwaf_result_p = ctypes.POINTER(ddwaf_result)
 
@@ -278,7 +284,11 @@ class ddwaf_config_obfuscator(ctypes.Structure):
     ]
 
 
-ddwaf_object_free_fn = ctypes.POINTER(ctypes.CFUNCTYPE(None, ddwaf_object_p))
+ddwaf_object_free_fn = ctypes.CFUNCTYPE(None, ddwaf_object_p)
+ddwaf_object_free = ddwaf_object_free_fn(
+    ("ddwaf_object_free", ddwaf),
+    ((1, "object"),),
+)
 
 
 class ddwaf_config(ctypes.Structure):
@@ -296,7 +306,7 @@ def __init__(
         max_string_length=0,
         key_regex="",
         value_regex="",
-        free_fn=None,
+        free_fn=ddwaf_object_free,
     ):
         # type: (ddwaf_config, int, int, int, unicode, unicode, Optional[Any]) -> None
         self.limits.max_container_size = max_container_size
@@ -310,10 +320,46 @@ def __init__(
 ddwaf_config_p = ctypes.POINTER(ddwaf_config)
 
 
-# TODO MAYBE LATER
 ddwaf_handle = ctypes.c_void_p  # may stay as this because it's mainly an abstract type in the interface
 ddwaf_context = ctypes.c_void_p  # may stay as this because it's mainly an abstract type in the interface
 
+
+class ddwaf_handle_capsule:
+    def __init__(self, handle):
+        # type: (ddwaf_handle) -> None
+        self.handle = handle
+        self.free_fn = ddwaf_destroy
+
+    def __del__(self):
+        if self.handle:
+            try:
+                self.free_fn(self.handle)
+            except TypeError:
+                pass
+            self.handle = None
+
+    def __bool__(self):
+        return bool(self.handle)
+
+
+class ddwaf_context_capsule:
+    def __init__(self, ctx):
+        # type: (ddwaf_context) -> None
+        self.ctx = ctx
+        self.free_fn = ddwaf_context_destroy
+
+    def __del__(self):
+        if self.ctx:
+            try:
+                self.free_fn(self.ctx)
+            except TypeError:
+                pass
+            self.ctx = None
+
+    def __bool__(self):
+        return bool(self.ctx)
+
+
 ddwaf_log_cb = ctypes.POINTER(
     ctypes.CFUNCTYPE(
         None, ctypes.c_int, ctypes.c_char_p, ctypes.c_char_p, ctypes.c_uint, ctypes.c_char_p, ctypes.c_uint64
@@ -334,6 +380,12 @@ def __init__(
     ),
 )
 
+
+def py_ddwaf_init(ruleset_map, config, info):
+    # type: (ddwaf_object, Any, Any) -> ddwaf_handle_capsule
+    return ddwaf_handle_capsule(ddwaf_init(ruleset_map, config, info))
+
+
 ddwaf_update = ctypes.CFUNCTYPE(ddwaf_handle, ddwaf_handle, ddwaf_object_p, ddwaf_ruleset_info_p)(
     ("ddwaf_update", ddwaf),
     (
@@ -343,6 +395,12 @@ def __init__(
     ),
 )
 
+
+def py_ddwaf_update(handle, ruleset_map, info):
+    # type: (ddwaf_handle_capsule, ddwaf_object, Any) -> ddwaf_handle_capsule
+    return ddwaf_handle_capsule(ddwaf_update(handle.handle, ruleset_map, ctypes.byref(info)))
+
+
 ddwaf_destroy = ctypes.CFUNCTYPE(None, ddwaf_handle)(
     ("ddwaf_destroy", ddwaf),
     ((1, "handle"),),
@@ -366,9 +424,9 @@ def __init__(
 
 
 def py_ddwaf_required_addresses(handle):
-    # type: (ctypes.c_void_p) -> list[unicode]
+    # type: (ddwaf_handle_capsule) -> list[unicode]
     size = ctypes.c_uint32()
-    obj = ddwaf_required_addresses(handle, ctypes.byref(size))
+    obj = ddwaf_required_addresses(handle.handle, ctypes.byref(size))
     return [obj[i].decode("UTF-8") for i in range(size.value)]
 
 
@@ -377,6 +435,12 @@ def py_ddwaf_required_addresses(handle):
     ((1, "handle"),),
 )
 
+
+def py_ddwaf_context_init(handle):
+    # type: (ddwaf_handle_capsule) -> ddwaf_context_capsule
+    return ddwaf_context_capsule(ddwaf_context_init(handle.handle))
+
+
 ddwaf_run = ctypes.CFUNCTYPE(ctypes.c_int, ddwaf_context, ddwaf_object_p, ddwaf_result_p, ctypes.c_uint64)(
     ("ddwaf_run", ddwaf), ((1, "context"), (1, "data"), (1, "result"), (1, "timeout"))
 )
@@ -470,10 +534,6 @@ def py_ddwaf_required_addresses(handle):
 # ddwaf_object_get_index
 # ddwaf_object_get_bool https://github.com/DataDog/libddwaf/commit/7dc68dacd972ae2e2a3c03a69116909c98dbd9cb
 
-ddwaf_object_free = ctypes.CFUNCTYPE(None, ddwaf_object_p)(
-    ("ddwaf_object_free", ddwaf),
-    ((1, "object"),),
-)
 
 ddwaf_get_version = ctypes.CFUNCTYPE(ctypes.c_char_p)(
     ("ddwaf_get_version", ddwaf),

ddtrace/appsec/processor.py

@@ -42,8 +42,7 @@
     from typing import Tuple
     from typing import Union
 
-    from ddtrace.appsec.ddwaf import DDWaf_result
-    from ddtrace.appsec.ddwaf.ddwaf_types import DDWafRulesType
+    from ddtrace.appsec.ddwaf.ddwaf_types import ddwaf_context_capsule
     from ddtrace.span import Span
 
 
@@ -195,7 +194,7 @@ def on_span_start(self, span):
 
         if span.span_type != SpanTypes.WEB:
             return
-        self._ddwaf._at_request_start()
+        ctx = self._ddwaf._at_request_start()
 
         peer_ip = _asm_request_context.get_ip()
         headers = _asm_request_context.get_headers()
@@ -205,7 +204,7 @@ def on_span_start(self, span):
         span.set_tag_str(RUNTIME_FAMILY, "python")
 
         def waf_callable(custom_data=None):
-            return self._waf_action(span._local_root or span, custom_data)
+            return self._waf_action(span._local_root or span, ctx, custom_data)
 
         _asm_request_context.set_waf_callback(waf_callable)
 
@@ -228,8 +227,8 @@ def waf_callable(custom_data=None):
                 # _asm_request_context.call_callback()
                 _asm_request_context.call_waf_callback({"REQUEST_HTTP_IP": None})
 
-    def _waf_action(self, span, custom_data=None):
-        # type: (Span, dict[str, Any] | None) -> None
+    def _waf_action(self, span, ctx, custom_data=None):
+        # type: (Span, ddwaf_context_capsule, dict[str, Any] | None) -> None
         """
         Call the `WAF` with the given parameters. If `custom_data_names` is specified as
         a list of `(WAF_NAME, WAF_STR)` tuples specifying what values of the `WAF_DATA_NAMES`
@@ -267,7 +266,7 @@ def _waf_action(self, span, custom_data=None):
                     log.debug("[action] WAF missing value %s", SPAN_DATA_NAMES[key])
 
         log.debug("[DDAS-001-00] Executing ASM In-App WAF")
-        waf_results = self._ddwaf.run(data, self._waf_timeout)
+        waf_results = self._ddwaf.run(ctx, data, self._waf_timeout)
         if waf_results and waf_results.data:
             log.debug("[DDAS-011-00] ASM In-App WAF returned: %s", waf_results.data)
 
@@ -342,15 +341,6 @@ def _is_needed(self, address):
         # type: (str) -> bool
         return address in self._addresses_to_keep
 
-    def _run_ddwaf(self, data):
-        # type: (DDWafRulesType) -> DDWaf_result | None
-        try:
-            return self._ddwaf.run(data, self._waf_timeout)  # res is a serialized json
-        except Exception:
-            log.warning("Error executing ASM In-App WAF: ", exc_info=True)
-
-        return None
-
     def on_span_finish(self, span):
         # type: (Span) -> None
 
@@ -359,5 +349,5 @@ def on_span_finish(self, span):
         # this call is only necessary for tests or frameworks that are not using blocking
         if span.get_tag(APPSEC.JSON) is None:
             log.debug("metrics waf call")
-            self._waf_action(span)
+            _asm_request_context.call_waf_callback()
         self._ddwaf._at_request_end()

releasenotes/notes/fix-memory-leaks-and-segfaults-9fbdc8d66b32fb4f.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    ASM: fix memory leaks and memory corruption in the interface between ASM and the WAF library

tests/appsec/test_processor.py

@@ -473,7 +473,8 @@ def test_ddwaf_run():
             "server.request.cookies": {"attack": "1' or '1' = '1'"},
             "server.response.headers.no_cookies": {"content-type": "text/html; charset=utf-8", "content-length": "207"},
         }
-        res = _ddwaf.run(data, DEFAULT.WAF_TIMEOUT)  # res is a serialized json
+        ctx = _ddwaf._at_request_start()
+        res = _ddwaf.run(ctx, data, DEFAULT.WAF_TIMEOUT)  # res is a serialized json
         assert res.data.startswith('[{"rule":{"id":"crs-942-100"')
         assert res.runtime > 0
         assert res.total_runtime > 0