Address PR #121 review comments: memory safety, correctness, hardening

accel.c:
- Replace empty TODO type stubs with NotImplementedError raises
- Add CHECK_REMAINING macro for bounds checking on buffer reads
- Replace unaligned pointer-cast reads with memcpy for WASM/ARM safety
- Fix double-decref in output error paths (set to NULL before goto)
- Fix Py_None reference leak by removing pre-switch INCREF
- Fix MYSQL_TYPE_NULL consuming an extra byte from next column
- Add PyErr_Format in default switch cases
- Add PyErr_Occurred() checks after PyLong/PyFloat conversions

Python:
- Align list/tuple multi-return handling in registry.py with C path
- Add _write_all_fd helper for partial os.write() handling
- Harden handshake recvmsg: name length bound, ancdata validation,
  MSG_CTRUNC check, FD cleanup on error
- Wrap get_context('fork') with platform safety error
- Narrow events.py exception catch to (ImportError, OSError)
- Fix _iquery DataFrame check ordering (check before list())
- Expand setblocking(False) warning comment
- Update WIT and wasm.py docstrings for code parameter

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: kesmit13

accel.c

@@ -1166,12 +1166,13 @@ def _iquery(
             cur.execute(oper, params)
             if not re.match(r'^\s*(select|show|call|echo)\s+', oper, flags=re.I):
                 return []
-            out = list(cur.fetchall())
+            raw = cur.fetchall()
+            if hasattr(raw, 'to_dict') and callable(raw.to_dict):
+                return raw.to_dict(orient='records')
+            out = list(raw)
             if not out:
                 return []
-            if hasattr(out, 'to_dict') and callable(getattr(out, 'to_dict')):
-                out = out.to_dict(orient='records')
-            elif isinstance(out[0], (tuple, list)):
+            if isinstance(out[0], (tuple, list)):
                 if cur.description:
                     names = [x[0] for x in cur.description]
                     if fix_names:

singlestoredb/connection.py

@@ -41,6 +41,9 @@
 # Pre-pack the status OK header prefix to avoid per-request struct.pack
 _STATUS_OK_PREFIX = struct.pack('<Q', STATUS_OK)
 
+# Maximum function name length to prevent resource exhaustion
+_MAX_FUNCTION_NAME_LEN = 4096
+
 # Enable per-request timing via environment variable
 _PROFILE = os.environ.get('SINGLESTOREDB_UDF_PROFILE', '') == '1'
 
@@ -83,18 +86,56 @@ def _handle_connection_inner(
         logger.warning(f'Unsupported protocol version: {version}')
         return
 
+    if namelen > _MAX_FUNCTION_NAME_LEN:
+        logger.warning(f'Function name too long: {namelen}')
+        return
+
     # Receive function name + 2 FDs via SCM_RIGHTS
     fd_model = array.array('i', [0, 0])
     msg, ancdata, flags, addr = conn.recvmsg(
         namelen,
         socket.CMSG_LEN(2 * fd_model.itemsize),
     )
-    if len(ancdata) != 1:
-        logger.warning(f'Expected 1 ancdata, got {len(ancdata)}')
-        return
 
-    function_name = msg.decode('utf8')
-    input_fd, output_fd = struct.unpack('<ii', ancdata[0][2])
+    # Validate ancdata and extract FDs
+    received_fds: list[int] = []
+    try:
+        if len(ancdata) != 1:
+            logger.warning(f'Expected 1 ancdata, got {len(ancdata)}')
+            return
+
+        level, type_, fd_data = ancdata[0]
+        if level != socket.SOL_SOCKET or type_ != socket.SCM_RIGHTS:
+            logger.warning(
+                f'Unexpected ancdata level={level} type={type_}',
+            )
+            return
+
+        if flags & getattr(socket, 'MSG_CTRUNC', 0):
+            logger.warning('Ancillary data was truncated (MSG_CTRUNC)')
+            return
+
+        fd_array = array.array('i')
+        fd_array.frombytes(fd_data)
+        received_fds = list(fd_array)
+
+        if len(received_fds) != 2:
+            logger.warning(
+                f'Expected 2 FDs, got {len(received_fds)}',
+            )
+            return
+
+        function_name = msg.decode('utf8')
+        input_fd, output_fd = received_fds[0], received_fds[1]
+        # Clear so finally doesn't close FDs we're handing off
+        received_fds = []
+    finally:
+        # Close any received FDs if we're returning early
+        for fd in received_fds:
+            try:
+                os.close(fd)
+            except OSError:
+                pass
 
     # --- Control signal path ---
     if function_name.startswith('@@'):
@@ -160,7 +201,7 @@ def _handle_control_signal(
             else:
                 os.ftruncate(output_fd, max(_MIN_OUTPUT_SIZE, response_size))
                 os.lseek(output_fd, 0, os.SEEK_SET)
-                os.write(output_fd, response_bytes)
+                _write_all_fd(output_fd, response_bytes)
 
             # Send [status=200, size]
             conn.sendall(struct.pack('<QQ', STATUS_OK, response_size))
@@ -298,7 +339,7 @@ def _handle_udf_loop(
                         os.ftruncate(output_fd, needed)
                         current_output_size = needed
                     os.lseek(output_fd, 0, os.SEEK_SET)
-                    os.write(output_fd, output_data)
+                    _write_all_fd(output_fd, output_data)
                 if profile:
                     t_mmap_write += time.monotonic() - t0
 
@@ -359,3 +400,17 @@ def _recv_exact_py(sock: socket.socket, n: int) -> bytes | None:
             return None
         pos += nbytes
     return bytes(buf)
+
+
+def _write_all_fd(fd: int, data: bytes) -> None:
+    """Write all bytes to a file descriptor, handling partial writes."""
+    view = memoryview(data)
+    written = 0
+    while written < len(data):
+        try:
+            n = os.write(fd, view[written:])
+        except InterruptedError:
+            continue
+        if n == 0:
+            raise RuntimeError('short write to output fd')
+        written += n

singlestoredb/functions/ext/collocated/connection.py

@@ -458,7 +458,7 @@ def call_function(
         results = []
         for row in rows:
             result = func(*row)
-            if not isinstance(result, tuple):
+            if not isinstance(result, (tuple, list)):
                 result = [result]
             results.append(list(result))
         return bytes(_dump_rowdat_1(return_types, row_ids, results))

singlestoredb/functions/ext/collocated/registry.py

@@ -24,6 +24,7 @@
 from typing import Optional
 from typing import Tuple
 
+from .connection import _write_all_fd
 from .connection import handle_connection
 from .registry import FunctionRegistry
 
@@ -61,7 +62,7 @@ def _write_pipe_message(fd: int, payload: bytes) -> None:
     Wire format: [u32 LE length][payload].
     """
     header = struct.pack('<I', len(payload))
-    os.write(fd, header + payload)
+    _write_all_fd(fd, header + payload)
 
 
 class SharedRegistry:
@@ -252,7 +253,14 @@ def _run_process_mode(
         kills and re-forks all workers so every worker has the updated
         registry state.
         """
-        ctx = multiprocessing.get_context('fork')
+        try:
+            ctx = multiprocessing.get_context('fork')
+        except ValueError:
+            raise RuntimeError(
+                "Process mode requires 'fork' multiprocessing context, "
+                'which is not available on this platform. '
+                "Use process_mode='thread' instead.",
+            )
         # workers[wid] = (process, pipe_read_fd)
         workers: Dict[
             int,
@@ -407,11 +415,13 @@ def _worker_signal_handler(
             signal.signal(signal.SIGTERM, _worker_signal_handler)
             signal.signal(signal.SIGINT, signal.SIG_IGN)
 
-            # Set non-blocking so accept() raises BlockingIOError
-            # instead of blocking when another worker wins the race.
-            # O_NONBLOCK is on the open file description (shared across
-            # forked processes), but that's fine — all workers want
-            # non-blocking accept and the parent doesn't call accept.
+            # WARNING: setblocking(False) sets O_NONBLOCK on the open
+            # file description, which is shared across all forked
+            # processes. This is intentional here — all workers need
+            # non-blocking accept() to handle the thundering-herd race,
+            # and the parent process never calls accept() on this
+            # socket. Do NOT add blocking operations on this socket
+            # in the parent process after workers are forked.
             server_sock.setblocking(False)
 
             registry = self.shared_registry.get_thread_local_registry()

singlestoredb/functions/ext/collocated/server.py

@@ -60,7 +60,11 @@ def create_function(
         code: str,
         replace: bool,
     ) -> None:
-        """Register a function from its signature and Python source code."""
+        """Register a function from its signature and function body.
+
+        The ``code`` parameter should contain the function body, not a
+        full ``def`` statement or ``@udf``-decorated source.
+        """
         try:
             _registry.create_function(signature, code, replace)
         except Exception as e:

singlestoredb/functions/ext/collocated/wasm.py

@@ -7,7 +7,7 @@
 try:
     from IPython import get_ipython
     has_ipython = True
-except Exception:
+except (ImportError, OSError):
     has_ipython = False
 
 

singlestoredb/utils/events.py

@@ -14,9 +14,9 @@ interface function-handler {
     /// args_data_format, returns_data_format, function_type, doc
     describe-functions: func() -> result<string, string>;
 
-    /// Register a function from its signature and Python source code.
+    /// Register a function from its signature and source code.
     /// `signature` is a JSON object matching the describe-functions element schema.
-    /// `code` is the Python source containing the @udf-decorated function.
+    /// `code` is the function body (not a full `def` statement).
     /// `replace` controls whether an existing function of the same name is overwritten.
     create-function: func(signature: string, code: string, replace: bool) -> result<_, string>;
 }