Skip to content

Commit b3a0101

Browse files
miss-islingtonStanFromIrelandrani-pinchuk
authored
[3.13] gh-98896: resource_tracker: use json&base64 to allow arbitrary shared memory names (GH-138473) (GH-142014)
(cherry picked from commit 64d6bde) Co-authored-by: Stan Ulbrych <[email protected]> Co-authored-by: Rani Pinchuk <[email protected]>
1 parent 3bbd669 commit b3a0101

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

@@ -200,6 +203,17 @@ def _launch(self):
200203
finally:
201204
os.close(r)
202205

206+
def _make_probe_message(self):
207+
"""Return a JSON-encoded probe message."""
208+
return (
209+
json.dumps(
210+
{"cmd": "PROBE", "rtype": "noop"},
211+
ensure_ascii=True,
212+
separators=(",", ":"),
213+
)
214+
+ "\n"
215+
).encode("ascii")
216+
203217
def _ensure_running_and_write(self, msg=None):
204218
with self._lock:
205219
if self._lock._recursion_count() > 1:
@@ -211,7 +225,7 @@ def _ensure_running_and_write(self, msg=None):
211225
if self._fd is not None:
212226
# resource tracker was launched before, is it still running?
213227
if msg is None:
214-
to_send = b'PROBE:0:noop\n'
228+
to_send = self._make_probe_message()
215229
else:
216230
to_send = msg
217231
try:
@@ -238,7 +252,7 @@ def _check_alive(self):
238252
try:
239253
# We cannot use send here as it calls ensure_running, creating
240254
# a cycle.
241-
os.write(self._fd, b'PROBE:0:noop\n')
255+
os.write(self._fd, self._make_probe_message())
242256
except OSError:
243257
return False
244258
else:
@@ -257,11 +271,25 @@ def _write(self, msg):
257271
assert nbytes == len(msg), f"{nbytes=} != {len(msg)=}"
258272

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

266294
self._ensure_running_and_write(msg)
267295

@@ -294,7 +322,23 @@ def main(fd):
294322
with open(fd, 'rb') as f:
295323
for line in f:
296324
try:
297-
cmd, name, rtype = line.strip().decode('ascii').split(':')
325+
try:
326+
obj = json.loads(line.decode('ascii'))
327+
except Exception as e:
328+
raise ValueError("malformed resource_tracker message: %r" % (line,)) from e
329+
330+
cmd = obj["cmd"]
331+
rtype = obj["rtype"]
332+
b64 = obj.get("base64_name", "")
333+
334+
if not isinstance(cmd, str) or not isinstance(rtype, str) or not isinstance(b64, str):
335+
raise ValueError("malformed resource_tracker fields: %r" % (obj,))
336+
337+
try:
338+
name = base64.urlsafe_b64decode(b64).decode('utf-8', 'surrogateescape')
339+
except ValueError as e:
340+
raise ValueError("malformed resource_tracker base64_name: %r" % (b64,)) from e
341+
298342
cleanup_func = _CLEANUP_FUNCS.get(rtype, None)
299343
if cleanup_func is None:
300344
raise ValueError(

Lib/test/_test_multiprocessing.py

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6789,3 +6789,47 @@ class SemLock(_multiprocessing.SemLock):
67896789
name = f'test_semlock_subclass-{os.getpid()}'
67906790
s = SemLock(1, 0, 10, name, False)
67916791
_multiprocessing.sem_unlink(name)
6792+
6793+
6794+
@unittest.skipUnless(HAS_SHMEM, "requires multiprocessing.shared_memory")
6795+
class TestSharedMemoryNames(unittest.TestCase):
6796+
def test_that_shared_memory_name_with_colons_has_no_resource_tracker_errors(self):
6797+
# Test script that creates and cleans up shared memory with colon in name
6798+
test_script = textwrap.dedent("""
6799+
import sys
6800+
from multiprocessing import shared_memory
6801+
import time
6802+
6803+
# Test various patterns of colons in names
6804+
test_names = [
6805+
"a:b",
6806+
"a:b:c",
6807+
"test:name:with:many:colons",
6808+
":starts:with:colon",
6809+
"ends:with:colon:",
6810+
"::double::colons::",
6811+
"name\\nwithnewline",
6812+
"name-with-trailing-newline\\n",
6813+
"\\nname-starts-with-newline",
6814+
"colons:and\\nnewlines:mix",
6815+
"multi\\nline\\nname",
6816+
]
6817+
6818+
for name in test_names:
6819+
try:
6820+
shm = shared_memory.SharedMemory(create=True, size=100, name=name)
6821+
shm.buf[:5] = b'hello' # Write something to the shared memory
6822+
shm.close()
6823+
shm.unlink()
6824+
6825+
except Exception as e:
6826+
print(f"Error with name '{name}': {e}", file=sys.stderr)
6827+
sys.exit(1)
6828+
6829+
print("SUCCESS")
6830+
""")
6831+
6832+
rc, out, err = script_helper.assert_python_ok("-c", test_script)
6833+
self.assertIn(b"SUCCESS", out)
6834+
self.assertNotIn(b"traceback", err.lower(), err)
6835+
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)