Skip to content

Commit 64d6bde

Browse files
[3.14] gh-98896: resource_tracker: use json&base64 to allow arbitrary shared memory names (GH-138473) (GH-141922)
Co-authored-by: Rani Pinchuk <33353578+rani-pinchuk@users.noreply.github.com>
1 parent ad60d89 commit 64d6bde

File tree

3 files changed

+98
-8
lines changed

3 files changed

+98
-8
lines changed

Lib/multiprocessing/resource_tracker.py

Lines changed: 52 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,16 @@
1515
# this resource tracker process, "killall python" would probably leave unlinked
1616
# resources.
1717

18+
import base64
1819
import os
1920
import signal
2021
import sys
2122
import threading
2223
import warnings
2324
from collections import deque
2425

26+
import json
27+
2528
from . import spawn
2629
from . import util
2730

@@ -196,6 +199,17 @@ def _launch(self):
196199
finally:
197200
os.close(r)
198201

202+
def _make_probe_message(self):
203+
"""Return a JSON-encoded probe message."""
204+
return (
205+
json.dumps(
206+
{"cmd": "PROBE", "rtype": "noop"},
207+
ensure_ascii=True,
208+
separators=(",", ":"),
209+
)
210+
+ "\n"
211+
).encode("ascii")
212+
199213
def _ensure_running_and_write(self, msg=None):
200214
with self._lock:
201215
if self._lock._recursion_count() > 1:
@@ -207,7 +221,7 @@ def _ensure_running_and_write(self, msg=None):
207221
if self._fd is not None:
208222
# resource tracker was launched before, is it still running?
209223
if msg is None:
210-
to_send = b'PROBE:0:noop\n'
224+
to_send = self._make_probe_message()
211225
else:
212226
to_send = msg
213227
try:
@@ -234,7 +248,7 @@ def _check_alive(self):
234248
try:
235249
# We cannot use send here as it calls ensure_running, creating
236250
# a cycle.
237-
os.write(self._fd, b'PROBE:0:noop\n')
251+
os.write(self._fd, self._make_probe_message())
238252
except OSError:
239253
return False
240254
else:
@@ -253,11 +267,25 @@ def _write(self, msg):
253267
assert nbytes == len(msg), f"{nbytes=} != {len(msg)=}"
254268

255269
def _send(self, cmd, name, rtype):
256-
msg = f"{cmd}:{name}:{rtype}\n".encode("ascii")
257-
if len(msg) > 512:
258-
# posix guarantees that writes to a pipe of less than PIPE_BUF
259-
# bytes are atomic, and that PIPE_BUF >= 512
260-
raise ValueError('msg too long')
270+
# POSIX guarantees that writes to a pipe of less than PIPE_BUF (512 on Linux)
271+
# bytes are atomic. Therefore, we want the message to be shorter than 512 bytes.
272+
# POSIX shm_open() and sem_open() require the name, including its leading slash,
273+
# to be at most NAME_MAX bytes (255 on Linux)
274+
# With json.dump(..., ensure_ascii=True) every non-ASCII byte becomes a 6-char
275+
# escape like \uDC80.
276+
# As we want the overall message to be kept atomic and therefore smaller than 512,
277+
# we encode encode the raw name bytes with URL-safe Base64 - so a 255 long name
278+
# will not exceed 340 bytes.
279+
b = name.encode('utf-8', 'surrogateescape')
280+
if len(b) > 255:
281+
raise ValueError('shared memory name too long (max 255 bytes)')
282+
b64 = base64.urlsafe_b64encode(b).decode('ascii')
283+
284+
payload = {"cmd": cmd, "rtype": rtype, "base64_name": b64}
285+
msg = (json.dumps(payload, ensure_ascii=True, separators=(",", ":")) + "\n").encode("ascii")
286+
287+
# The entire JSON message is guaranteed < PIPE_BUF (512 bytes) by construction.
288+
assert len(msg) <= 512, f"internal error: message too long ({len(msg)} bytes)"
261289

262290
self._ensure_running_and_write(msg)
263291

@@ -290,7 +318,23 @@ def main(fd):
290318
with open(fd, 'rb') as f:
291319
for line in f:
292320
try:
293-
cmd, name, rtype = line.strip().decode('ascii').split(':')
321+
try:
322+
obj = json.loads(line.decode('ascii'))
323+
except Exception as e:
324+
raise ValueError("malformed resource_tracker message: %r" % (line,)) from e
325+
326+
cmd = obj["cmd"]
327+
rtype = obj["rtype"]
328+
b64 = obj.get("base64_name", "")
329+
330+
if not isinstance(cmd, str) or not isinstance(rtype, str) or not isinstance(b64, str):
331+
raise ValueError("malformed resource_tracker fields: %r" % (obj,))
332+
333+
try:
334+
name = base64.urlsafe_b64decode(b64).decode('utf-8', 'surrogateescape')
335+
except ValueError as e:
336+
raise ValueError("malformed resource_tracker base64_name: %r" % (b64,)) from e
337+
294338
cleanup_func = _CLEANUP_FUNCS.get(rtype, None)
295339
if cleanup_func is None:
296340
raise ValueError(

Lib/test/_test_multiprocessing.py

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7144,3 +7144,47 @@ class SemLock(_multiprocessing.SemLock):
71447144
name = f'test_semlock_subclass-{os.getpid()}'
71457145
s = SemLock(1, 0, 10, name, False)
71467146
_multiprocessing.sem_unlink(name)
7147+
7148+
7149+
@unittest.skipUnless(HAS_SHMEM, "requires multiprocessing.shared_memory")
7150+
class TestSharedMemoryNames(unittest.TestCase):
7151+
def test_that_shared_memory_name_with_colons_has_no_resource_tracker_errors(self):
7152+
# Test script that creates and cleans up shared memory with colon in name
7153+
test_script = textwrap.dedent("""
7154+
import sys
7155+
from multiprocessing import shared_memory
7156+
import time
7157+
7158+
# Test various patterns of colons in names
7159+
test_names = [
7160+
"a:b",
7161+
"a:b:c",
7162+
"test:name:with:many:colons",
7163+
":starts:with:colon",
7164+
"ends:with:colon:",
7165+
"::double::colons::",
7166+
"name\\nwithnewline",
7167+
"name-with-trailing-newline\\n",
7168+
"\\nname-starts-with-newline",
7169+
"colons:and\\nnewlines:mix",
7170+
"multi\\nline\\nname",
7171+
]
7172+
7173+
for name in test_names:
7174+
try:
7175+
shm = shared_memory.SharedMemory(create=True, size=100, name=name)
7176+
shm.buf[:5] = b'hello' # Write something to the shared memory
7177+
shm.close()
7178+
shm.unlink()
7179+
7180+
except Exception as e:
7181+
print(f"Error with name '{name}': {e}", file=sys.stderr)
7182+
sys.exit(1)
7183+
7184+
print("SUCCESS")
7185+
""")
7186+
7187+
rc, out, err = script_helper.assert_python_ok("-c", test_script)
7188+
self.assertIn(b"SUCCESS", out)
7189+
self.assertNotIn(b"traceback", err.lower(), err)
7190+
self.assertNotIn(b"resource_tracker.py", err, err)
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Fix a failure in multiprocessing resource_tracker when SharedMemory names contain colons.
2+
Patch by Rani Pinchuk.

0 commit comments

Comments
 (0)