From 7f4e5872ba33432c8b714bff0672fd7380c0158c Mon Sep 17 00:00:00 2001 From: hoon <230467962+atc722@users.noreply.github.com> Date: Sat, 22 Aug 2026 17:01:12 +0900 Subject: [PATCH] feat(protocol): add shared MID exchange registry --- smartthings_local/protocol/dtls_session.py | 229 ++++++++++---- tests/test_dtls_session_mid_registry.py | 339 +++++++++++++++++++++ tests/test_dtls_session_reader_death.py | 20 ++ 3 files changed, 534 insertions(+), 54 deletions(-) create mode 100644 tests/test_dtls_session_mid_registry.py diff --git a/smartthings_local/protocol/dtls_session.py b/smartthings_local/protocol/dtls_session.py index b90487c..8b0ff3b 100644 --- a/smartthings_local/protocol/dtls_session.py +++ b/smartthings_local/protocol/dtls_session.py @@ -31,6 +31,7 @@ import socket import threading import time +from dataclasses import dataclass from OpenSSL import SSL @@ -108,6 +109,14 @@ class _EtagChanged(Exception): transfer, so the blocks in hand are from two different versions.""" +@dataclass(slots=True) +class _MidExchange: + """One pending request, indexed independently by token and MID.""" + + pending: tuple[threading.Event, dict] + acknowledged: bool = False + + # ICMP errors a connected UDP socket surfaces on the next recv. On these # appliances they show up while the device is rebooting, while it holds an # orphaned association, or across a router blip, and the next datagram @@ -270,10 +279,11 @@ def __init__(self, host, port, cert_path=None, key_path=None, *, self.endpoint = None self._send_lock = threading.Lock() - # Guards the MID/token counters and _pending. The refetch worker - # makes the session its own second concurrent get() caller, so - # two threads can mint tokens at once; without this they can - # collide and one transfer silently absorbs the other's blocks. + # Guards the MID/token counters and pending-request registries. + # The refetch worker makes the session its own second concurrent + # get() caller, so two threads can mint tokens at once; without + # this they can collide and one transfer silently absorbs the + # other's blocks. self._state_lock = threading.Lock() # Randomize MID and token counter starting points so reconnects # don't reuse identifiers from previous sessions — Samsung's @@ -288,6 +298,10 @@ def __init__(self, host, port, cert_path=None, key_path=None, *, self._observe_tok_counter = 0x40 + (os.urandom(1)[0] & 0xBF) # token (bytes) → (Event, container_dict) self._pending = {} + # request MID (int) → _MidExchange. Empty ACK and RST frames carry + # no token, so request lifecycle state must also be reachable by the + # MID that was registered before send. + self._pending_mids = {} # token (bytes) → href (str) self._observe_tokens = {} @@ -500,24 +514,68 @@ def close(self): self.sock.close() except Exception: pass - with self._state_lock: - pending = list(self._pending.items()) - self._pending.clear() - for tok, (ev, container) in pending: - container.setdefault('err', SessionClosedError()) - ev.set() - self._observe_tokens.clear() + # Publish the closed state before draining pending requests. A caller + # that passed its entry check just before close() will then fail the + # post-registration liveness check instead of registering after the + # drain and waiting against a session that can no longer respond. self.sock = None self.conn = None self.dest = None self.endpoint = None + self._close_pending_requests() + self._observe_tokens.clear() # ---- send / receive plumbing ------------------------------------- + def _next_available_mid_locked(self): + """Return the next MID that is not owned by a live request.""" + for _ in range(0x10000): + self._mid = (self._mid + 1) & 0xFFFF + if self._mid not in self._pending_mids: + return self._mid + raise SessionError() + def _next_mid(self): with self._state_lock: - self._mid = (self._mid + 1) & 0xFFFF - return self._mid + return self._next_available_mid_locked() + + def _register_pending_request(self, tok, ev, container): + """Atomically allocate a MID and index one request by MID and token.""" + with self._state_lock: + if tok in self._pending: + raise SessionError() + mid = self._next_available_mid_locked() + pending = (ev, container) + exchange = _MidExchange(pending) + self._pending[tok] = pending + self._pending_mids[mid] = exchange + return mid, exchange + + def _unregister_pending_request(self, tok, mid, exchange): + """Remove only the exact request registered under both indices.""" + with self._state_lock: + if self._pending.get(tok) is exchange.pending: + self._pending.pop(tok, None) + if self._pending_mids.get(mid) is exchange: + self._pending_mids.pop(mid, None) + + def _close_pending_requests(self): + """Fail, unregister, and wake every request that can no longer finish.""" + with self._state_lock: + pending_by_id = { + id(record): record for record in self._pending.values() + } + pending_by_id.update( + (id(exchange.pending), exchange.pending) + for exchange in self._pending_mids.values() + ) + pending = list(pending_by_id.values()) + for _ev, container in pending: + container.setdefault('err', SessionClosedError()) + self._pending.clear() + self._pending_mids.clear() + for ev, _container in pending: + ev.set() def _next_tok(self): with self._state_lock: @@ -640,11 +698,7 @@ def _reader_loop(self): # Reader no longer owns the socket — callers must fail fast. self._reader_running.clear() # Make sure pending waiters don't hang if the reader dies. - with self._state_lock: - pending = list(self._pending.items()) - for tok, (ev, container) in pending: - container.setdefault('err', SessionClosedError()) - ev.set() + self._close_pending_requests() # Nothing will answer a refetch now either. with self._refetch_cond: self._refetch_pending.clear() @@ -663,6 +717,17 @@ def _dispatch_coap(self, datagram): kind, fmt_code(code), mid, tok.hex() or '-', len(ropts), len(payload)) + # RFC 7252 empty messages are exactly the four-byte v1 header with + # TKL=0 and code=0. parse_coap() intentionally stays a lightweight + # general parser, so validate the raw shape before a control frame is + # allowed to mutate a MID-indexed exchange. + bare_control = ( + len(datagram) == 4 + and datagram[0] >> 6 == 1 + and datagram[0] & 0x0F == 0 + and code == 0 + ) + # ACK back any CON from the device to suppress retransmits. # RFC 7252 §4.2 — ACK is a bare frame (token len 0, code 0). if mt == TYPE_CON: @@ -674,19 +739,44 @@ def _dispatch_coap(self, datagram): # Empty ACK with no options & no payload = "separate response # coming" — used by Samsung's RT-OCF for the larger reads. Stop # the retransmit timer on the client side and wait for the CON. - if mt == TYPE_ACK and code == 0 and not payload and not ropts: + if mt == TYPE_ACK and code == 0: + if not bare_control: + return + with self._state_lock: + exchange = self._pending_mids.get(mid) + if exchange is not None: + exchange.acknowledged = True + ev, _container = exchange.pending + if exchange is not None: + ev.set() + return + + # A reset rejects the matching exchange. Like an empty ACK it has no + # response token, so surface it through the request's MID registry. + if mt == TYPE_RST: + if not bare_control: + return + with self._state_lock: + exchange = self._pending_mids.get(mid) + if exchange is not None: + ev, container = exchange.pending + if 'code' not in container and 'err' not in container: + container['err'] = SessionError() + if exchange is not None: + ev.set() return # Pending one-shot? Resolve and return. with self._state_lock: rec = self._pending.get(tok) + if rec is not None: + ev, container = rec + container['code'] = code + container['mtype'] = mt + container['mid'] = mid + container['options'] = ropts + container['payload'] = payload if rec is not None: - ev, container = rec - container['code'] = code - container['mtype'] = mt - container['mid'] = mid - container['options'] = ropts - container['payload'] = payload ev.set() return @@ -935,10 +1025,12 @@ def _exchange_block(self, tok, path_segs, query, num, szx, deadline): for attempt in range(_BLOCK_MAX_ATTEMPTS): ev = threading.Event() container = {} - with self._state_lock: - self._pending[tok] = (ev, container) + mid, exchange = self._register_pending_request(tok, ev, container) try: - mid = self._next_mid() + # Close the reader-death registration race: after this request + # is visible to reader-finally, recheck that the reader still + # owns the session before sending. + self._check_live() opts = [(URI_PATH, s.encode()) for s in path_segs] for q in query: opts.append((URI_QUERY, q.encode())) @@ -948,19 +1040,33 @@ def _exchange_block(self, tok, path_segs, query, num, szx, deadline): self._send_dgram( build_coap(TYPE_CON, METHOD_GET, mid, tok, opts)) while True: - per_wait = min(_BLOCK_ACK_TIMEOUT, - max(0.1, deadline - time.time())) + with self._state_lock: + if ('err' in container + or ('code' in container + and self._block_num_matches( + container, num))): + return container + if 'code' in container: + logger.debug( + "GET %s /%s block %d: stale block, " + "still waiting", + self.host, '/'.join(path_segs), num, + ) + container.clear() + acknowledged = exchange.acknowledged + ev.clear() + remaining = deadline - time.time() + if remaining <= 0: + if acknowledged: + raise SessionTimeoutError() + break + per_wait = ( + remaining if acknowledged + else min(_BLOCK_ACK_TIMEOUT, max(0.1, remaining)) + ) if not self._wait_for_block(ev, per_wait): - break # attempt timed out - if 'err' in container or self._block_num_matches( - container, num): - return container - logger.debug( - "GET %s /%s block %d: stale block, still waiting", - self.host, '/'.join(path_segs), num) - ev.clear() - container.clear() - if deadline - time.time() <= 0: + if acknowledged: + raise SessionTimeoutError() break remaining = deadline - time.time() if remaining <= 0 or attempt == _BLOCK_MAX_ATTEMPTS - 1: @@ -975,8 +1081,7 @@ def _exchange_block(self, tok, path_segs, query, num, szx, deadline): attempt + 1, _BLOCK_MAX_ATTEMPTS, ) finally: - with self._state_lock: - self._pending.pop(tok, None) + self._unregister_pending_request(tok, mid, exchange) raise SessionTimeoutError() def _wait_for_block(self, ev, per_wait): @@ -1015,28 +1120,44 @@ def post(self, path_segs, body_cbor, timeout=8.0): (code, payload_bytes). body_cbor must already be encoded.""" self._check_live() tok = self._next_tok() - mid = self._next_mid() opts = [(URI_PATH, s.encode()) for s in path_segs] opts.append((CONTENT_FORMAT, CF_CBOR)) opts.append((ACCEPT, CF_CBOR)) - datagram = build_coap(TYPE_CON, METHOD_POST, mid, tok, opts, - body_cbor) ev = threading.Event() container = {} - with self._state_lock: - self._pending[tok] = (ev, container) + mid, exchange = self._register_pending_request(tok, ev, container) try: self.pace() + # The reader can exit between the entry liveness check and the + # registration above. Recheck after registration so its teardown + # cannot miss this waiter. self._check_live() + datagram = build_coap(TYPE_CON, METHOD_POST, mid, tok, opts, + body_cbor) self._send_dgram(datagram) - if not ev.wait(timeout): - raise SessionTimeoutError() - if 'err' in container: - raise container['err'] - return container['code'], container['payload'] + deadline = time.time() + timeout + while True: + with self._state_lock: + error = container.get('err') + has_response = 'code' in container + response = ( + (container['code'], container['payload']) + if has_response else None + ) + if error is None and not has_response: + ev.clear() + if error is not None: + raise error + if has_response: + return response + remaining = deadline - time.time() + if remaining <= 0 or not ev.wait(remaining): + raise SessionTimeoutError() + # An empty ACK only stops retransmission. POST does not retry + # today, but keeping the acknowledged exchange pending here is + # the common contract the write retry path will build on. finally: - with self._state_lock: - self._pending.pop(tok, None) + self._unregister_pending_request(tok, mid, exchange) def ping(self): """RFC 7252 §4.4 CoAP Ping — empty CON, no token, no payload. diff --git a/tests/test_dtls_session_mid_registry.py b/tests/test_dtls_session_mid_registry.py new file mode 100644 index 0000000..e99be02 --- /dev/null +++ b/tests/test_dtls_session_mid_registry.py @@ -0,0 +1,339 @@ +"""Shared MID lifecycle for token-correlated CoAP requests.""" + +from __future__ import annotations + +import threading + +import pytest + +from smartthings_local.errors import ( + EndpointError, + SessionClosedError, + SessionError, + SessionTimeoutError, +) +from smartthings_local.protocol import dtls_session +from smartthings_local.protocol.coap import ( + TYPE_ACK, + TYPE_CON, + TYPE_NON, + TYPE_RST, + build_coap, + parse_coap, +) +from smartthings_local.protocol.dtls_session import DtlsCoapSession + + +class _NullAuth: + def configure_context(self, _context): + return None + + +def _session(): + session = DtlsCoapSession( + "device.example", + 5684, + auth=_NullAuth(), + rate_limit_rps=1_000_000, + ) + session.conn = object() + return session + + +def _request(session, operation, *, timeout=1.0): + if operation == "get": + return session.get(["device", "0"], timeout=timeout) + return session.post(["mode", "vs", "0"], b"payload", timeout=timeout) + + +@pytest.mark.parametrize("operation", ["get", "post"]) +def test_request_is_indexed_by_token_and_mid_before_send(operation): + session = _session() + + def send(datagram): + mtype, _code, mid, token, _options, _payload = parse_coap(datagram) + assert mtype == TYPE_CON + exchange = session._pending_mids[mid] + assert session._pending[token] is exchange.pending + session._dispatch_coap( + build_coap(TYPE_ACK, 0x45, mid, token, [], b"ok") + ) + + session._send_dgram = send + + assert _request(session, operation) == (0x45, b"ok") + assert session._pending == {} + assert session._pending_mids == {} + + +@pytest.mark.parametrize("operation", ["get", "post"]) +def test_empty_ack_keeps_request_pending_for_separate_response(operation): + session = _session() + + def send(datagram): + mtype, code, mid, token, _options, _payload = parse_coap(datagram) + if mtype == TYPE_ACK and code == 0: + return + + exchange = session._pending_mids[mid] + session._dispatch_coap( + build_coap(TYPE_ACK, 0, mid, b"", []) + ) + + assert exchange.acknowledged is True + assert session._pending_mids[mid] is exchange + assert session._pending[token] is exchange.pending + assert exchange.pending[1] == {} + + session._dispatch_coap( + build_coap(TYPE_CON, 0x45, mid + 1, token, [], b"ok") + ) + + session._send_dgram = send + + assert _request(session, operation) == (0x45, b"ok") + assert session._pending == {} + assert session._pending_mids == {} + + +def test_empty_ack_stops_get_retransmission(monkeypatch): + session = _session() + requests = [] + waits = [] + + monkeypatch.setattr(dtls_session, "_BLOCK_ACK_TIMEOUT", 0.01) + + def send(datagram): + mtype, code, mid, token, _options, _payload = parse_coap(datagram) + if mtype == TYPE_ACK and code == 0: + return + requests.append((mid, token)) + session._dispatch_coap(build_coap(TYPE_ACK, 0, mid, b"", [])) + + def wait_for_block(_event, timeout): + waits.append(timeout) + mid, token = requests[-1] + session._dispatch_coap( + build_coap(TYPE_CON, 0x45, mid + 1, token, [], b"ok") + ) + return True + + session._send_dgram = send + session._wait_for_block = wait_for_block + + assert session.get(["device", "0"], timeout=1.0) == (0x45, b"ok") + assert len(requests) == 1 + assert waits and waits[0] > dtls_session._BLOCK_ACK_TIMEOUT + + +@pytest.mark.parametrize("operation", ["get", "post"]) +def test_matching_bare_rst_fails_request(operation): + session = _session() + + def send(datagram): + _mtype, _code, mid, _token, _options, _payload = parse_coap(datagram) + session._dispatch_coap(build_coap(TYPE_RST, 0, mid, b"", [])) + + session._send_dgram = send + + with pytest.raises(SessionError) as raised: + _request(session, operation) + + assert type(raised.value) is SessionError + assert session._pending == {} + assert session._pending_mids == {} + + +@pytest.mark.parametrize("mtype", [TYPE_ACK, TYPE_RST]) +def test_unknown_bare_ack_or_rst_is_ignored(mtype): + session = _session() + event = threading.Event() + container = {} + mid, exchange = session._register_pending_request( + b"token", event, container + ) + try: + session._dispatch_coap( + build_coap(mtype, 0, (mid + 1) & 0xFFFF, b"", []) + ) + + assert not event.is_set() + assert exchange.acknowledged is False + assert container == {} + finally: + session._unregister_pending_request(b"token", mid, exchange) + + +@pytest.mark.parametrize("mtype", [TYPE_ACK, TYPE_RST]) +@pytest.mark.parametrize( + "malformed", + ["token", "wrong_version", "empty_payload_marker"], +) +def test_malformed_empty_control_message_is_ignored(mtype, malformed): + session = _session() + event = threading.Event() + container = {} + mid, exchange = session._register_pending_request( + b"token", event, container + ) + try: + if malformed == "token": + datagram = build_coap(mtype, 0, mid, b"token", []) + elif malformed == "wrong_version": + datagram = bytes([ + (2 << 6) | (mtype << 4), + 0, + mid >> 8, + mid & 0xFF, + ]) + else: + datagram = build_coap(mtype, 0, mid, b"", []) + b"\xFF" + session._dispatch_coap(datagram) + + assert not event.is_set() + assert exchange.acknowledged is False + assert container == {} + finally: + session._unregister_pending_request(b"token", mid, exchange) + + +@pytest.mark.parametrize("operation", ["get", "post"]) +@pytest.mark.parametrize("outcome", ["send_failure", "timeout"]) +def test_failure_paths_unregister_both_indices(operation, outcome): + session = _session() + + if outcome == "send_failure": + def fail_send(_datagram): + raise EndpointError() + + session._send_dgram = fail_send + expected = EndpointError + timeout = 1.0 + else: + session._send_dgram = lambda _datagram: None + if operation == "get": + session._wait_for_block = lambda _event, _timeout: False + expected = SessionTimeoutError + timeout = 0.0 + + with pytest.raises(expected): + _request(session, operation, timeout=timeout) + + assert session._pending == {} + assert session._pending_mids == {} + + +def test_close_drains_both_indices_and_wakes_waiters(): + session = _session() + pending = [] + for token in (b"get", b"post"): + event = threading.Event() + container = {} + session._register_pending_request(token, event, container) + pending.append((event, container)) + + session.close() + + assert session._pending == {} + assert session._pending_mids == {} + assert all(event.is_set() for event, _container in pending) + assert all( + isinstance(container.get("err"), SessionClosedError) + for _event, container in pending + ) + + +@pytest.mark.parametrize("operation", ["get", "post"]) +def test_request_cannot_register_after_close_drains(operation): + session = _session() + sends = [] + register_entered = threading.Event() + resume_registration = threading.Event() + drained = threading.Event() + release_close = threading.Event() + original_register = session._register_pending_request + original_drain = session._close_pending_requests + outcome = {} + session._send_dgram = sends.append + + def pause_before_registration(token, event, container): + register_entered.set() + resume_registration.wait(1.0) + return original_register(token, event, container) + + def pause_after_drain(): + original_drain() + drained.set() + release_close.wait(1.0) + + def request(): + try: + _request(session, operation, timeout=0.0) + except Exception as error: + outcome["error"] = error + + session._register_pending_request = pause_before_registration + session._close_pending_requests = pause_after_drain + request_thread = threading.Thread(target=request) + close_thread = threading.Thread(target=session.close) + request_thread.start() + assert register_entered.wait(1.0) + close_thread.start() + assert drained.wait(1.0) + try: + resume_registration.set() + request_thread.join(1.0) + assert not request_thread.is_alive() + assert isinstance(outcome.get("error"), SessionClosedError) + assert sends == [] + assert session._pending == {} + assert session._pending_mids == {} + finally: + resume_registration.set() + release_close.set() + request_thread.join(1.0) + close_thread.join(1.0) + assert not request_thread.is_alive() + assert not close_thread.is_alive() + + +def test_mid_allocation_skips_a_live_exchange_across_wrap(): + session = _session() + session._mid = 0xFFFF + first_mid, first = session._register_pending_request( + b"first", threading.Event(), {} + ) + session._mid = 0xFFFF + second_mid, second = session._register_pending_request( + b"second", threading.Event(), {} + ) + try: + assert first_mid == 0 + assert second_mid == 1 + assert session._pending_mids[first_mid] is first + assert session._pending_mids[second_mid] is second + finally: + session._unregister_pending_request(b"first", first_mid, first) + session._unregister_pending_request(b"second", second_mid, second) + + +def test_non_control_response_still_resolves_by_token(): + session = _session() + event = threading.Event() + container = {} + request_mid, exchange = session._register_pending_request( + b"token", event, container + ) + try: + session._dispatch_coap( + build_coap(TYPE_NON, 0x45, 0x1234, b"token", [], b"ok") + ) + + assert event.is_set() + assert container["code"] == 0x45 + assert container["mid"] == 0x1234 + assert container["payload"] == b"ok" + finally: + session._unregister_pending_request( + b"token", request_mid, exchange + ) diff --git a/tests/test_dtls_session_reader_death.py b/tests/test_dtls_session_reader_death.py index 3a5dfc3..d341e30 100644 --- a/tests/test_dtls_session_reader_death.py +++ b/tests/test_dtls_session_reader_death.py @@ -135,6 +135,26 @@ def test_fatal_socket_error_exits_with_warning(caplog): assert "reader exiting" in warnings[0].getMessage() +def test_fatal_reader_exit_drains_mid_registry_and_wakes_waiters(): + sess = _make_session() + pending = [] + for token in (b"get", b"post"): + event = threading.Event() + container = {} + sess._register_pending_request(token, event, container) + pending.append((event, container)) + + _run_reader(sess, [OSError(errno.EBADF, "bad file descriptor")]) + + assert sess._pending == {} + assert sess._pending_mids == {} + assert all(event.is_set() for event, _container in pending) + assert all( + isinstance(container.get("err"), SessionClosedError) + for _event, container in pending + ) + + def test_request_fails_fast_after_reader_death(): sess = _make_session() _run_reader(sess, [OSError(errno.EBADF, "bad file descriptor")])