Skip to content

Commit 0f9acfd

Browse files
kesmit13claude
kesmit13
and
claude
committed
Address PR #121 review comments: refcount leaks, type shadowing, protocol hardening
Fix PyDict_SetItem key reference leaks in load_rowdat_1_numpy by creating temporary key objects and decrementing after use (7 call sites on hot path). Remove NEWDECIMAL (246) from string_types so decimal_types handler is reachable, returning decimal.Decimal instead of strings. Fix _pack_time to use integer arithmetic instead of float total_seconds(). Reject datetime.time UDF annotations with a clear TypeError (timedelta required). Normalize VECTOR element_type to uppercase before SQL emission. Add recvmsg partial-read check in plugin handshake to prevent protocol desync. Validate socket path before unlink in _bind_socket to prevent arbitrary file deletion. Use private module namespace for dynamic UDF registration instead of __main__. Broaden lazy import exception handling to catch OSError for WASM/WASI environments where optional deps may not raise ImportError. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 79f8238 commit 0f9acfd

9 files changed

Lines changed: 83 additions & 26 deletions

File tree

‎accel.c‎

Lines changed: 55 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -2570,7 +2570,13 @@ static PyObject *load_rowdat_1_numpy(PyObject *self, PyObject *args, PyObject *k
25702570
// Create dict for strings/blobs
25712571
py_objs = PyDict_New();
25722572
if (!py_objs) goto error;
2573-
CHECKRC(PyDict_SetItem(py_objs, PyLong_FromUnsignedLongLong(0), Py_None));
2573+
{
2574+
PyObject *py_key = PyLong_FromUnsignedLongLong(0);
2575+
if (!py_key) goto error;
2576+
int rc = PyDict_SetItem(py_objs, py_key, Py_None);
2577+
Py_DECREF(py_key);
2578+
CHECKRC(rc);
2579+
}
25742580

25752581
// Build output arrays
25762582
j = 0;
@@ -2677,8 +2683,14 @@ static PyObject *load_rowdat_1_numpy(PyObject *self, PyObject *args, PyObject *k
26772683
if (!py_dec) goto error;
26782684
u64 = (uint64_t)py_dec;
26792685
memcpy(out_cols[i] + j * 8, &u64, 8);
2680-
CHECKRC(PyDict_SetItem(py_objs, PyLong_FromUnsignedLongLong(u64), py_dec));
2681-
Py_CLEAR(py_dec);
2686+
{
2687+
PyObject *py_key = PyLong_FromUnsignedLongLong(u64);
2688+
if (!py_key) { Py_CLEAR(py_dec); goto error; }
2689+
int rc = PyDict_SetItem(py_objs, py_key, py_dec);
2690+
Py_DECREF(py_key);
2691+
Py_CLEAR(py_dec);
2692+
CHECKRC(rc);
2693+
}
26822694
}
26832695
break;
26842696

@@ -2693,8 +2705,14 @@ static PyObject *load_rowdat_1_numpy(PyObject *self, PyObject *args, PyObject *k
26932705
if (!py_dt) goto error;
26942706
u64 = (uint64_t)py_dt;
26952707
memcpy(out_cols[i] + j * 8, &u64, 8);
2696-
CHECKRC(PyDict_SetItem(py_objs, PyLong_FromUnsignedLongLong(u64), py_dt));
2697-
Py_CLEAR(py_dt);
2708+
{
2709+
PyObject *py_key = PyLong_FromUnsignedLongLong(u64);
2710+
if (!py_key) { Py_CLEAR(py_dt); goto error; }
2711+
int rc = PyDict_SetItem(py_objs, py_key, py_dt);
2712+
Py_DECREF(py_key);
2713+
Py_CLEAR(py_dt);
2714+
CHECKRC(rc);
2715+
}
26982716
}
26992717
break;
27002718
}
@@ -2709,8 +2727,14 @@ static PyObject *load_rowdat_1_numpy(PyObject *self, PyObject *args, PyObject *k
27092727
if (!py_td) goto error;
27102728
u64 = (uint64_t)py_td;
27112729
memcpy(out_cols[i] + j * 8, &u64, 8);
2712-
CHECKRC(PyDict_SetItem(py_objs, PyLong_FromUnsignedLongLong(u64), py_td));
2713-
Py_CLEAR(py_td);
2730+
{
2731+
PyObject *py_key = PyLong_FromUnsignedLongLong(u64);
2732+
if (!py_key) { Py_CLEAR(py_td); goto error; }
2733+
int rc = PyDict_SetItem(py_objs, py_key, py_td);
2734+
Py_DECREF(py_key);
2735+
Py_CLEAR(py_td);
2736+
CHECKRC(rc);
2737+
}
27142738
}
27152739
break;
27162740
}
@@ -2726,8 +2750,14 @@ static PyObject *load_rowdat_1_numpy(PyObject *self, PyObject *args, PyObject *k
27262750
if (!py_dt) goto error;
27272751
u64 = (uint64_t)py_dt;
27282752
memcpy(out_cols[i] + j * 8, &u64, 8);
2729-
CHECKRC(PyDict_SetItem(py_objs, PyLong_FromUnsignedLongLong(u64), py_dt));
2730-
Py_CLEAR(py_dt);
2753+
{
2754+
PyObject *py_key = PyLong_FromUnsignedLongLong(u64);
2755+
if (!py_key) { Py_CLEAR(py_dt); goto error; }
2756+
int rc = PyDict_SetItem(py_objs, py_key, py_dt);
2757+
Py_DECREF(py_key);
2758+
Py_CLEAR(py_dt);
2759+
CHECKRC(rc);
2760+
}
27312761
}
27322762
break;
27332763
}
@@ -2749,8 +2779,14 @@ static PyObject *load_rowdat_1_numpy(PyObject *self, PyObject *args, PyObject *k
27492779
if (!py_str) goto error;
27502780
u64 = (uint64_t)py_str;
27512781
memcpy(out_cols[i] + j * 8, &u64, 8);
2752-
CHECKRC(PyDict_SetItem(py_objs, PyLong_FromUnsignedLongLong(u64), py_str));
2753-
Py_CLEAR(py_str);
2782+
{
2783+
PyObject *py_key = PyLong_FromUnsignedLongLong(u64);
2784+
if (!py_key) { Py_CLEAR(py_str); goto error; }
2785+
int rc = PyDict_SetItem(py_objs, py_key, py_str);
2786+
Py_DECREF(py_key);
2787+
Py_CLEAR(py_str);
2788+
CHECKRC(rc);
2789+
}
27542790
}
27552791
break;
27562792

@@ -2765,8 +2801,14 @@ static PyObject *load_rowdat_1_numpy(PyObject *self, PyObject *args, PyObject *k
27652801
if (!py_blob) goto error;
27662802
u64 = (uint64_t)py_blob;
27672803
memcpy(out_cols[i] + j * 8, &u64, 8);
2768-
CHECKRC(PyDict_SetItem(py_objs, PyLong_FromUnsignedLongLong(u64), py_blob));
2769-
Py_CLEAR(py_blob);
2804+
{
2805+
PyObject *py_key = PyLong_FromUnsignedLongLong(u64);
2806+
if (!py_key) { Py_CLEAR(py_blob); goto error; }
2807+
int rc = PyDict_SetItem(py_objs, py_key, py_blob);
2808+
Py_DECREF(py_key);
2809+
Py_CLEAR(py_blob);
2810+
CHECKRC(rc);
2811+
}
27702812
}
27712813
break;
27722814

‎singlestoredb/functions/dtypes.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1783,7 +1783,8 @@ def VECTOR(
17831783
SQLString
17841784
17851785
"""
1786-
if element_type.upper() not in (F16, F32, F64, I8, I16, I32, I64):
1786+
element_type = element_type.upper()
1787+
if element_type not in (F16, F32, F64, I8, I16, I32, I64):
17871788
raise ValueError(f'unsupported element type: {element_type}')
17881789
out = f'VECTOR({int(length)}, {element_type})'
17891790
out = SQLString(

‎singlestoredb/functions/ext/plugin/connection.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,13 @@ def _handle_connection_inner(
9797
socket.CMSG_LEN(2 * fd_model.itemsize),
9898
)
9999

100+
if len(msg) != namelen:
101+
logger.warning(
102+
f'Short read on function name: expected {namelen}, '
103+
f'got {len(msg)}',
104+
)
105+
return
106+
100107
# Validate ancdata and extract FDs
101108
received_fds: list[int] = []
102109
try:

‎singlestoredb/functions/ext/plugin/registry.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -380,7 +380,7 @@ def create_function(
380380

381381
full_code = self._build_python_code(sig, code)
382382

383-
name = '__main__'
383+
name = 'singlestoredb.functions.ext.plugin._dynamic'
384384
compiled = compile(full_code, f'<{name}>', 'exec')
385385

386386
if name in sys.modules:

‎singlestoredb/functions/ext/plugin/server.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import select
1414
import signal
1515
import socket
16+
import stat
1617
import struct
1718
import sys
1819
import threading
@@ -201,6 +202,11 @@ def _bind_socket(self) -> socket.socket:
201202
"""Create, bind, and listen on the Unix domain socket."""
202203
sock_path = self.config['socket']
203204
if os.path.exists(sock_path):
205+
mode = os.stat(sock_path).st_mode
206+
if not stat.S_ISSOCK(mode):
207+
raise RuntimeError(
208+
f'Path exists but is not a socket: {sock_path}',
209+
)
204210
os.unlink(sock_path)
205211

206212
server_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)

‎singlestoredb/functions/ext/rowdat_1.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@
8585
ft.TINY, -ft.TINY, ft.SHORT, -ft.SHORT, ft.INT24, -ft.INT24,
8686
ft.LONG, -ft.LONG, ft.LONGLONG, -ft.LONGLONG,
8787
])
88-
string_types = set([15, 245, 246, 247, 248, 249, 250, 251, 252, 253, 254])
88+
string_types = set([15, 245, 247, 248, 249, 250, 251, 252, 253, 254])
8989
binary_types = set([-x for x in string_types])
9090
datetime_types = {ft.DATETIME, ft.TIMESTAMP}
9191
date_types = {ft.DATE}
@@ -139,7 +139,7 @@ def _unpack_date(v: int) -> _dt.date:
139139

140140
def _pack_time(td: _dt.timedelta) -> int:
141141
"""Pack a timedelta into int64 per rowdat_1 spec."""
142-
total_us = int(td.total_seconds() * 1_000_000)
142+
total_us = td.days * 86_400_000_000 + td.seconds * 1_000_000 + td.microseconds
143143
sign = -1 if total_us < 0 else 1
144144
total_us = abs(total_us)
145145
us = total_us % 1_000_000

‎singlestoredb/functions/signature.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -433,7 +433,10 @@ def normalize_dtype(dtype: Any) -> str:
433433
if issubclass(dtype, datetime.date):
434434
return 'date'
435435
if issubclass(dtype, datetime.time):
436-
return 'time'
436+
raise TypeError(
437+
'datetime.time is not supported for UDF annotations; '
438+
'use datetime.timedelta instead',
439+
)
437440
if issubclass(dtype, datetime.timedelta):
438441
return 'time'
439442

‎singlestoredb/tests/test_ext_func_data.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1385,10 +1385,8 @@ def test_decimal_pass_through(self):
13851385
colspec=colspec, returns=[NEWDECIMAL],
13861386
data=input_data, func=lambda x: x,
13871387
)
1388-
# accel reads as Decimal, writes back as length-prefixed string;
1389-
# _load treats 246 as string_types, so we get a string back
13901388
_, rows = rowdat_1._load([('r0', NEWDECIMAL)], output_data)
1391-
assert rows[0][0] == '123.456'
1389+
assert rows[0][0] == decimal.Decimal('123.456')
13921390

13931391
def test_decimal_negative(self):
13941392
colspec = [('d', NEWDECIMAL)]
@@ -1402,7 +1400,7 @@ def test_decimal_negative(self):
14021400
data=input_data, func=lambda x: x,
14031401
)
14041402
_, rows = rowdat_1._load([('r0', NEWDECIMAL)], output_data)
1405-
assert rows[0][0] == '-99.99'
1403+
assert rows[0][0] == decimal.Decimal('-99.99')
14061404

14071405
def test_decimal_null(self):
14081406
colspec = [('d', NEWDECIMAL)]

‎singlestoredb/utils/_lazy_import.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ def get_numpy() -> Optional[Any]:
1111
"""Return numpy module or None if not installed."""
1212
try:
1313
return importlib.import_module('numpy')
14-
except ImportError:
14+
except (ImportError, OSError):
1515
return None
1616

1717

@@ -20,7 +20,7 @@ def get_pandas() -> Optional[Any]:
2020
"""Return pandas module or None if not installed."""
2121
try:
2222
return importlib.import_module('pandas')
23-
except ImportError:
23+
except (ImportError, OSError):
2424
return None
2525

2626

@@ -29,7 +29,7 @@ def get_polars() -> Optional[Any]:
2929
"""Return polars module or None if not installed."""
3030
try:
3131
return importlib.import_module('polars')
32-
except ImportError:
32+
except (ImportError, OSError):
3333
return None
3434

3535

@@ -38,5 +38,5 @@ def get_pyarrow() -> Optional[Any]:
3838
"""Return pyarrow module or None if not installed."""
3939
try:
4040
return importlib.import_module('pyarrow')
41-
except ImportError:
41+
except (ImportError, OSError):
4242
return None

0 commit comments

Comments
 (0)