From bbf120214bf9a2d3e2b9aab02812eaa5183b57b8 Mon Sep 17 00:00:00 2001 From: YuanTingHsieh Date: Thu, 2 Jul 2026 09:35:08 -0700 Subject: [PATCH 1/4] Add Client API Cell protocol vocabulary and session-token auth primitives TE-1 (Wave 0, interface freeze #1) of the Client API Execution Modes program (see docs/design/client_api_execution_modes_plan.md on PR #4853). New pure-library package nvflare/client/cell consumed by the trainer engine (TE-3), external_process (EP-3), and attach (AT-2) tracks so the HELLO handshake is implemented once: - defs.py: the Client API control-protocol vocabulary - channel, PROTOCOL_VERSION, Topic (HELLO/HELLO_CHALLENGE/HELLO_PROOF/ HELLO_ACCEPTED, TASK_READY/TASK_ACCEPTED/TASK_FAILED, RESULT_READY/ RESULT_ACCEPTED/RESULT_REJECTED, LOG, HEARTBEAT, ABORT, SHUTDOWN, BYE, ERROR) and MsgKey payload keys. Topic values are "client_api."-prefixed to avoid colliding with ipc/defs.py values. - auth.py: session-token primitives per the design's attach auth contract - TokenScope (frozen, binds token to job/site/attach id/ target FQCN/trainer FQCN/rank policy/protocol version), generate_session_token/generate_nonce (secrets), token_digest (sha256; anything persisted stores only the digest), compute/verify_hello_proof (HMAC-SHA256 keyed by the token over a domain-tagged, length-prefixed serialization of nonce + full scope; constant-time compare), and SessionTokenManager (single-use nonces consumed on any attempt, attach_timeout expiry with injectable clock, single-session enforcement, invalidate; thread-safe; no cellnet/file-I/O deps). 37 unit tests: proof round-trip, tamper/replay/consumed-nonce/scope- mismatch/expiry/single-session/invalidation matrices, serialization non-ambiguity, and vocabulary uniqueness incl. non-collision with ipc. Co-Authored-By: Claude Fable 5 --- nvflare/client/cell/__init__.py | 13 + nvflare/client/cell/auth.py | 387 ++++++++++++++++++++ nvflare/client/cell/defs.py | 101 ++++++ tests/unit_test/client/cell/__init__.py | 13 + tests/unit_test/client/cell/auth_test.py | 432 +++++++++++++++++++++++ tests/unit_test/client/cell/defs_test.py | 121 +++++++ 6 files changed, 1067 insertions(+) create mode 100644 nvflare/client/cell/__init__.py create mode 100644 nvflare/client/cell/auth.py create mode 100644 nvflare/client/cell/defs.py create mode 100644 tests/unit_test/client/cell/__init__.py create mode 100644 tests/unit_test/client/cell/auth_test.py create mode 100644 tests/unit_test/client/cell/defs_test.py diff --git a/nvflare/client/cell/__init__.py b/nvflare/client/cell/__init__.py new file mode 100644 index 0000000000..4fc25d0d3c --- /dev/null +++ b/nvflare/client/cell/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/nvflare/client/cell/auth.py b/nvflare/client/cell/auth.py new file mode 100644 index 0000000000..9d8dbbb62c --- /dev/null +++ b/nvflare/client/cell/auth.py @@ -0,0 +1,387 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Session-token auth primitives for the Client API control protocol. + +Implements the attach auth contract of the Client API Execution Modes design +(docs/design/client_api_execution_modes.md, "Appendix B: Attach Topology and Auth"): + +- The executor generates a high-entropy session token per attach id / launch. The raw token + is held only in memory for the session lifetime; anything persisted stores only its digest. +- Token proof is challenge-response, not bearer presentation: HELLO -> HELLO_CHALLENGE (nonce) + -> HELLO_PROOF (HMAC keyed by the token over the nonce and the full token scope). The raw + token never crosses the wire, so observing attach traffic does not permit replay. +- The token is scoped to (job id, site name, attach id, target FQCN, trainer FQCN, rank + policy, protocol version), is single-session by default, and expires if the trainer does + not attach before attach_timeout. +- Two proof paths share the same HMAC primitive: + * The interactive HELLO_CHALLENGE / HELLO_PROOF path goes through SessionTokenManager, + which issues a self-generated single-use nonce (issue_nonce) and verifies the proof + against it (verify_proof). + * The one-round variant (Appendix B) folds the proof into HELLO, computed over an + executor nonce delivered in the bootstrap config combined with a trainer-supplied + nonce. This path does NOT go through SessionTokenManager.verify_proof (there is no + self-issued nonce); it uses the module-level verify_hello_proof directly over a + combined nonce (see combine_nonces). + +This module is part of interface freeze #1. It is a pure library: no Cell/cellnet imports, +no file I/O, no logging side effects. +""" + +import dataclasses +import hashlib +import hmac +import secrets +import threading +import time +from collections import OrderedDict +from typing import Callable, Optional + +from .defs import PROTOCOL_VERSION + +# Domain-separation tag for the HELLO proof HMAC, versioned with the proof format. +_PROOF_DOMAIN_TAG = "nvflare.client_api.hello_proof.v1" + +# Default sizes (bytes of entropy) for generated secrets. +DEFAULT_TOKEN_BYTES = 32 +DEFAULT_NONCE_BYTES = 16 + +# Minimum accepted length (hex chars) for a caller-supplied token: 16 bytes of entropy. +MIN_TOKEN_HEX_CHARS = 32 + + +@dataclasses.dataclass(frozen=True) +class TokenScope: + """The scope a session token is bound to. + + A token is valid only for this exact scope; the executor accepts an attach only when the + proof, the presented scope, and its own expected scope all match. + + Attributes: + job_id: id of the job the token was issued for. + site_name: name of the site (FL client) the token was issued for. + attach_id: attach/session id the token was issued for. + target_fqcn: FQCN of the CJ/job cell the trainer must talk to. + trainer_fqcn: FQCN the trainer cell must bind (routing name, not identity). + rank_policy: allowed rank policy for the session (e.g. which rank may attach). + protocol_version: Client API control protocol version. + """ + + job_id: str + site_name: str + attach_id: str + target_fqcn: str + trainer_fqcn: str + rank_policy: str + protocol_version: int = PROTOCOL_VERSION + + +def generate_session_token(num_bytes: int = DEFAULT_TOKEN_BYTES) -> str: + """Generate a new high-entropy session token. + + Args: + num_bytes: bytes of entropy in the token. + + Returns: + The token as a hex string. The raw token must be held only in memory; anything + persisted must store only its digest (see token_digest). + """ + return secrets.token_hex(num_bytes) + + +def generate_nonce(num_bytes: int = DEFAULT_NONCE_BYTES) -> str: + """Generate a new single-use challenge nonce as a hex string.""" + return secrets.token_hex(num_bytes) + + +def token_digest(token: str) -> str: + """Compute the persistable digest (SHA-256 hex) of a session token.""" + return hashlib.sha256(token.encode("utf-8")).hexdigest() + + +def _canonical_proof_message(nonce: str, scope: TokenScope) -> bytes: + """Serialize (nonce + all scope fields) canonically, unambiguously, and type-unambiguously. + + Each item is netstring-encoded (length-prefixed), so no combination of field values can + produce the same byte stream as a different combination (no concatenation ambiguity). + Each scope field additionally carries an explicit type tag (its Python type name) and is + rendered with repr(), so values that differ only by type -- e.g. protocol_version=1 (int) + vs protocol_version="1" (str) -- serialize differently and cannot be substituted for one + another (no cross-type proof forgery). Field order is fixed: domain tag, nonce, then + TokenScope fields in declaration order. + """ + items = [_PROOF_DOMAIN_TAG, nonce] + for f in dataclasses.fields(scope): + value = getattr(scope, f.name) + items.append(f"{type(value).__name__}:{value!r}") + + buf = bytearray() + for item in items: + encoded = item.encode("utf-8") + buf += f"{len(encoded)}:".encode("utf-8") + buf += encoded + buf += b"," + return bytes(buf) + + +def compute_hello_proof(token: str, nonce: str, scope: TokenScope) -> str: + """Compute the HELLO_PROOF value: HMAC-SHA256 keyed by the token over (nonce, scope). + + The proof covers the challenge nonce and the full token scope (attach id, job id, site + name, trainer FQCN, target FQCN, rank policy, protocol version), so a proof observed on + the wire is useless for any other nonce, scope, or protocol version. + + Args: + token: the raw session token (the HMAC key; never sent on the wire). + nonce: the challenge nonce from HELLO_CHALLENGE. + scope: the token scope the proof is computed over. + + Returns: + The proof as a hex string. + """ + return hmac.new(token.encode("utf-8"), _canonical_proof_message(nonce, scope), hashlib.sha256).hexdigest() + + +def verify_hello_proof(token: str, nonce: str, scope: TokenScope, proof: str) -> bool: + """Verify a HELLO_PROOF value in constant time. + + Args: + token: the raw session token. + nonce: the challenge nonce the proof must cover. + scope: the token scope the proof must cover. + proof: the presented proof (hex string). + + Returns: + True if the proof is valid for (token, nonce, scope). Never raises for a + malformed/attacker-supplied proof: a non-str, empty, or non-ASCII proof returns False. + """ + if not isinstance(proof, str) or not proof: + return False + try: + expected = compute_hello_proof(token, nonce, scope) + # Encode both operands to bytes before compare_digest: hmac.compare_digest raises + # TypeError on a non-ASCII str, so a non-ASCII (attacker-supplied) proof would raise + # instead of failing. Comparing bytes makes any malformed proof simply mismatch. + return hmac.compare_digest(expected.encode("utf-8"), proof.encode("utf-8")) + except (TypeError, ValueError): + return False + + +def combine_nonces(executor_nonce: str, trainer_nonce: str) -> str: + """Combine an executor nonce with a trainer nonce for the one-round HELLO proof variant. + + In the one-round variant (Appendix B), the executor delivers a nonce in the bootstrap + config and the trainer contributes its own nonce; the proof is computed over both so + neither side alone fixes the challenge. The combination is length-prefixed so the pair is + unambiguous regardless of nonce contents. Both sides must combine identically, then feed + the result as the ``nonce`` argument to compute_hello_proof / verify_hello_proof -- this + path does not use SessionTokenManager (there is no self-issued nonce). + """ + if not isinstance(executor_nonce, str) or not isinstance(trainer_nonce, str): + raise TypeError("executor_nonce and trainer_nonce must both be str") + return f"{len(executor_nonce)}:{executor_nonce}:{trainer_nonce}" + + +class SessionTokenManager: + """Executor-side manager for one session token and its challenge-response verification. + + Holds the raw token in memory for the session lifetime (plus its digest for anything + persisted), issues single-use challenge nonces, and verifies HELLO_PROOF messages while + enforcing the auth contract: + + - a nonce is consumed by its first verification attempt (success or failure), so a + captured proof cannot be replayed and a consumed nonce cannot be reused; + - the presented scope must equal the expected scope exactly; + - the token expires if no session is established within attach_timeout of creation; + - the token is single-session: while a session is active, further attach attempts are + rejected; + - after invalidate(), all verification fails permanently (reconnect requires a fresh + token through a new manager); + - issue_nonce() refuses (raises RuntimeError) once the manager is invalidated, expired, or + a session is active, and the set of outstanding pending nonces is bounded by + MAX_PENDING_NONCES (oldest evicted first) so a HELLO flood cannot grow memory without + bound. + + Binding the accepted session to the peer's FQCN (recording/enforcing the bound peer) is + intentionally out of scope here and handled separately (see AT-2). + + The clock is injectable for testability; it must be a monotonic seconds counter. + Thread-safe. + """ + + # Upper bound on outstanding (issued-but-unconsumed) nonces; oldest are evicted past this. + MAX_PENDING_NONCES = 1024 + + def __init__( + self, + scope: TokenScope, + attach_timeout: Optional[float] = None, + token: Optional[str] = None, + clock: Callable[[], float] = time.monotonic, + ): + """Create a manager for one session token. + + Args: + scope: the scope the token is bound to. + attach_timeout: seconds from creation within which the trainer must attach; + None means no attach expiry. + token: the raw token to manage; a new high-entropy token is generated if None. + clock: monotonic seconds source (injectable for tests). + """ + if attach_timeout is not None and attach_timeout <= 0: + raise ValueError(f"attach_timeout must be positive but got {attach_timeout}") + if not isinstance(scope, TokenScope): + raise TypeError(f"scope must be TokenScope but got {type(scope)}") + if token is not None: + if not isinstance(token, str): + raise TypeError(f"token must be str but got {type(token)}") + if len(token) < MIN_TOKEN_HEX_CHARS: + raise ValueError( + f"token must have at least {MIN_TOKEN_HEX_CHARS} chars " + f"(>= 16 bytes of entropy) but got {len(token)}" + ) + + self._scope = scope + self._token = token if token else generate_session_token() + self._digest = token_digest(self._token) + self._attach_timeout = attach_timeout + self._clock = clock + self._created_at = clock() + # Insertion-ordered map nonce -> issued_at, so the oldest pending nonce is evictable + # in O(1) when the cap is exceeded, and each nonce carries an issuance time for TTL. + self._pending_nonces = OrderedDict() + self._session_active = False + self._invalidated = False + self._lock = threading.Lock() + + @property + def token(self) -> str: + """The raw session token (in-memory only; never persist or send on the wire).""" + return self._token + + @property + def digest(self) -> str: + """The persistable SHA-256 hex digest of the token.""" + return self._digest + + @property + def scope(self) -> TokenScope: + """The scope this token is bound to.""" + return self._scope + + @property + def session_active(self) -> bool: + """Whether a session is currently bound to this token.""" + with self._lock: + return self._session_active + + @property + def invalidated(self) -> bool: + """Whether this token has been explicitly invalidated.""" + with self._lock: + return self._invalidated + + def is_expired(self) -> bool: + """Whether the attach window has closed without a session being established.""" + with self._lock: + return self._is_expired() + + def _is_expired(self) -> bool: + if self._attach_timeout is None or self._session_active: + return False + return (self._clock() - self._created_at) > self._attach_timeout + + def issue_nonce(self) -> str: + """Issue a fresh single-use challenge nonce for HELLO_CHALLENGE. + + Refuses to issue once the manager is no longer in the pre-attach state, so a stale or + flooded caller cannot accumulate nonces after the session has been settled: + + Raises: + RuntimeError: if the token has been invalidated, the attach window has expired, or + a session is already active. + + The number of outstanding (issued-but-unconsumed) nonces is capped at + MAX_PENDING_NONCES; issuing past the cap evicts the oldest pending nonce so a HELLO + flood cannot grow memory without bound. + """ + with self._lock: + if self._invalidated: + raise RuntimeError("cannot issue nonce: session token has been invalidated") + if self._session_active: + raise RuntimeError("cannot issue nonce: a session is already active") + if self._is_expired(): + raise RuntimeError("cannot issue nonce: attach window has expired") + nonce = generate_nonce() + self._pending_nonces[nonce] = self._clock() + # bound memory: evict oldest outstanding nonces once the cap is exceeded + while len(self._pending_nonces) > self.MAX_PENDING_NONCES: + self._pending_nonces.popitem(last=False) + return nonce + + def verify_proof(self, nonce: str, scope: TokenScope, proof: str) -> bool: + """Verify a HELLO_PROOF attach attempt; on success the session becomes active. + + The nonce is consumed by this call whether or not verification succeeds. The proof is + verified against the manager's EXPECTED scope (self._scope), not the caller-supplied + scope: a wrong presented scope therefore yields a proof mismatch in constant time + (hmac.compare_digest) rather than leaking, via a short-circuiting tuple compare, which + scope field mismatched -- and the accepted proof is bound to the scope the executor + expects. (Recording/enforcing the bound peer FQCN is out of scope here; see AT-2.) + + Args: + nonce: the challenge nonce the proof claims to answer; must have been issued by + issue_nonce() and not yet consumed. + scope: the scope presented by the attaching trainer. + proof: the presented proof (hex string). + + Returns: + True only if the nonce is valid and unconsumed, the token is not invalidated or + expired, no session is already active, the proof verifies against the expected + scope, and (redundantly) the presented scope equals the expected scope. + """ + if not isinstance(nonce, str): + return False + with self._lock: + # single-use: consumed by any verification attempt (success or failure) + issued = self._pending_nonces.pop(nonce, None) is not None + + if self._invalidated: + return False + if not issued: + return False + # An old nonce cannot be redeemed after the attach window: _is_expired() covers + # it (a nonce's issue time is >= the manager's creation time under a monotonic + # clock, so the manager-level window always trips first) -- no separate per-nonce + # TTL is needed while no session is active. + if self._is_expired(): + return False + if self._session_active: + # single-session: reject further attach attempts while a session is active + return False + # constant-time proof check bound to the EXPECTED scope + if not verify_hello_proof(self._token, nonce, self._scope, proof): + return False + # redundant defense-in-depth guard (the proof is already bound to self._scope) + if scope != self._scope: + return False + + self._session_active = True + return True + + def invalidate(self) -> None: + """Invalidate the token and session permanently; all future verification fails.""" + with self._lock: + self._invalidated = True + self._session_active = False + self._pending_nonces.clear() diff --git a/nvflare/client/cell/defs.py b/nvflare/client/cell/defs.py new file mode 100644 index 0000000000..c4ca97884f --- /dev/null +++ b/nvflare/client/cell/defs.py @@ -0,0 +1,101 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Protocol vocabulary for the Client API control protocol over Cell. + +This module is interface freeze #1 of the Client API Execution Modes design +(docs/design/client_api_execution_modes.md, "Control Protocol"). It is consumed by the +trainer-side Cell engine, the external_process backend, and the attach backend. It is a +pure vocabulary module: constants only, no Cell/cellnet imports, no I/O. +""" + +# Cell channel used for all Client API control protocol messages. +# Clearly namespaced so it cannot collide with the legacy FlareAgent channel +# ("flare_agent" in nvflare/client/ipc/defs.py) or other Cell channels. +# This is a frozen wire constant; its exact value is part of the cross-track contract. +CHANNEL = "client_api" + +# The Client API control protocol version carried in HELLO. +# V1 supports exactly one protocol version; the field exists so later versions +# can define a compatibility window. +PROTOCOL_VERSION = 1 + + +class Topic: + """Topics of the Client API control protocol messages over the Cell CHANNEL. + + Reply-type messages (HELLO_CHALLENGE, HELLO_ACCEPTED, TASK_ACCEPTED, RESULT_ACCEPTED, + RESULT_REJECTED) may be modeled as Cell request replies at runtime rather than separate + sends; the constants still name them so state machines, logs, and tests share one + vocabulary. + + Values are prefixed with "client_api." so they never collide with legacy topic values + (e.g. "hello"/"heartbeat"/"abort"/"bye" in nvflare/client/ipc/defs.py). + """ + + # Session setup (all out-of-process modes) + HELLO = "client_api.hello" + HELLO_CHALLENGE = "client_api.hello_challenge" + HELLO_PROOF = "client_api.hello_proof" + HELLO_ACCEPTED = "client_api.hello_accepted" + + # Per task (every round) + TASK_READY = "client_api.task_ready" + TASK_ACCEPTED = "client_api.task_accepted" + TASK_FAILED = "client_api.task_failed" + RESULT_READY = "client_api.result_ready" + RESULT_ACCEPTED = "client_api.result_accepted" + RESULT_REJECTED = "client_api.result_rejected" + + # Throughout the session + LOG = "client_api.log" + HEARTBEAT = "client_api.heartbeat" + + # Teardown / failure + ABORT = "client_api.abort" + SHUTDOWN = "client_api.shutdown" + BYE = "client_api.bye" + ERROR = "client_api.error" + + +class MsgKey: + """Keys for Client API control protocol message payloads. + + These string values are the frozen wire contract shared across tracks; renaming a value is + a protocol break, not a refactor. + """ + + SESSION_ID = "session_id" + ATTACH_ID = "attach_id" + JOB_ID = "job_id" + SITE_NAME = "site_name" + TRAINER_FQCN = "trainer_fqcn" + TARGET_FQCN = "target_fqcn" + RANK = "rank" + # Rank policy the session is scoped to (a TokenScope / HELLO_PROOF-covered field); + # distinct from RANK, which is the concrete rank a trainer reports in HELLO. + RANK_POLICY = "rank_policy" + PROTOCOL_VERSION = "protocol_version" + NONCE = "nonce" + PROOF = "proof" + REASON = "reason" + TASK_ID = "task_id" + # TASK_READY carries the task name alongside the task id. + TASK_NAME = "task_name" + # TASK_READY carries the FLModel reference and its params (lazy refs, not materialized bytes). + MODEL = "model" + PARAMS = "params" + RESULT_ID = "result_id" + TRANSFER_ID = "transfer_id" + # RESULT_READY carries the payload manifest describing the result envelope. + MANIFEST = "manifest" diff --git a/tests/unit_test/client/cell/__init__.py b/tests/unit_test/client/cell/__init__.py new file mode 100644 index 0000000000..4fc25d0d3c --- /dev/null +++ b/tests/unit_test/client/cell/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/unit_test/client/cell/auth_test.py b/tests/unit_test/client/cell/auth_test.py new file mode 100644 index 0000000000..c5ced5040e --- /dev/null +++ b/tests/unit_test/client/cell/auth_test.py @@ -0,0 +1,432 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import dataclasses +import hashlib + +import pytest + +from nvflare.client.cell.auth import ( + MIN_TOKEN_HEX_CHARS, + SessionTokenManager, + TokenScope, + combine_nonces, + compute_hello_proof, + generate_nonce, + generate_session_token, + token_digest, + verify_hello_proof, +) +from nvflare.client.cell.defs import PROTOCOL_VERSION + + +class _FakeClock: + """Injectable monotonic clock for expiry tests (no sleeps).""" + + def __init__(self, start: float = 1000.0): + self.now = start + + def advance(self, seconds: float): + self.now += seconds + + def __call__(self) -> float: + return self.now + + +def _make_scope(**overrides) -> TokenScope: + kwargs = dict( + job_id="job-1234", + site_name="site-1", + attach_id="attach-abcd", + target_fqcn="site-1.job-1234", + trainer_fqcn="site-1.-client_api_attach-abcd", + rank_policy="0", + protocol_version=PROTOCOL_VERSION, + ) + kwargs.update(overrides) + return TokenScope(**kwargs) + + +class TestTokenAndDigest: + def test_digest_is_stable_sha256_of_token(self): + token = generate_session_token() + expected = hashlib.sha256(token.encode("utf-8")).hexdigest() + assert token_digest(token) == expected + assert token_digest(token) == token_digest(token) + + def test_manager_digest_matches_token_digest(self): + mgr = SessionTokenManager(scope=_make_scope()) + assert mgr.digest == token_digest(mgr.token) + + def test_token_entropy_and_uniqueness(self): + tokens = {generate_session_token() for _ in range(100)} + assert len(tokens) == 100 + # default 32 bytes of entropy -> 64 hex chars + assert all(len(t) == 64 for t in tokens) + + def test_nonce_uniqueness(self): + nonces = {generate_nonce() for _ in range(100)} + assert len(nonces) == 100 + + +class TestHelloProof: + def test_proof_round_trip(self): + token = generate_session_token() + nonce = generate_nonce() + scope = _make_scope() + proof = compute_hello_proof(token, nonce, scope) + assert verify_hello_proof(token, nonce, scope, proof) + + def test_tampered_nonce_fails(self): + token = generate_session_token() + nonce = generate_nonce() + scope = _make_scope() + proof = compute_hello_proof(token, nonce, scope) + assert not verify_hello_proof(token, generate_nonce(), scope, proof) + + @pytest.mark.parametrize( + "field,value", + [ + ("job_id", "job-9999"), + ("site_name", "site-2"), + ("attach_id", "attach-zzzz"), + ("target_fqcn", "site-2.job-9999"), + ("trainer_fqcn", "site-1.-client_api_other"), + ("rank_policy", "1"), + ], + ) + def test_tampered_scope_field_fails(self, field, value): + token = generate_session_token() + nonce = generate_nonce() + scope = _make_scope() + proof = compute_hello_proof(token, nonce, scope) + tampered = dataclasses.replace(scope, **{field: value}) + assert not verify_hello_proof(token, nonce, tampered, proof) + + def test_wrong_protocol_version_fails(self): + token = generate_session_token() + nonce = generate_nonce() + scope = _make_scope() + proof = compute_hello_proof(token, nonce, scope) + skewed = dataclasses.replace(scope, protocol_version=PROTOCOL_VERSION + 1) + assert not verify_hello_proof(token, nonce, skewed, proof) + + def test_wrong_token_fails(self): + nonce = generate_nonce() + scope = _make_scope() + proof = compute_hello_proof(generate_session_token(), nonce, scope) + assert not verify_hello_proof(generate_session_token(), nonce, scope, proof) + + def test_empty_or_non_str_proof_fails(self): + token = generate_session_token() + nonce = generate_nonce() + scope = _make_scope() + assert not verify_hello_proof(token, nonce, scope, "") + assert not verify_hello_proof(token, nonce, scope, None) + + def test_serialization_is_unambiguous_across_field_boundaries(self): + # shifting bytes between adjacent fields must not produce the same proof + token = generate_session_token() + nonce = generate_nonce() + scope_a = _make_scope(job_id="ab", site_name="c") + scope_b = _make_scope(job_id="a", site_name="bc") + assert compute_hello_proof(token, nonce, scope_a) != compute_hello_proof(token, nonce, scope_b) + + def test_non_ascii_proof_returns_false_without_raising(self): + # hmac.compare_digest raises TypeError on a non-ASCII str; verify must return False. + token = generate_session_token() + nonce = generate_nonce() + scope = _make_scope() + assert verify_hello_proof(token, nonce, scope, "deadbeefÿ") is False + assert verify_hello_proof(token, nonce, scope, "中文") is False + + def test_cross_typed_scope_produces_different_proof_and_fails(self): + # protocol_version=1 (int) and protocol_version="1" (str) must not be interchangeable. + token = generate_session_token() + nonce = generate_nonce() + scope_int = _make_scope(protocol_version=1) + scope_str = _make_scope(protocol_version="1") + proof_int = compute_hello_proof(token, nonce, scope_int) + proof_str = compute_hello_proof(token, nonce, scope_str) + assert proof_int != proof_str + # a proof computed over the int scope must not verify for the str scope, and vice versa + assert not verify_hello_proof(token, nonce, scope_str, proof_int) + assert not verify_hello_proof(token, nonce, scope_int, proof_str) + + +class TestOneRoundProof: + """The one-round variant folds the proof into HELLO over a combined executor+trainer nonce + and uses the module-level functions directly (no SessionTokenManager / issue_nonce).""" + + def test_one_round_round_trip(self): + token = generate_session_token() + scope = _make_scope() + executor_nonce = generate_nonce() # delivered in bootstrap config + trainer_nonce = generate_nonce() # contributed by the trainer + combined = combine_nonces(executor_nonce, trainer_nonce) + proof = compute_hello_proof(token, combined, scope) + assert verify_hello_proof(token, combined, scope, proof) + + def test_one_round_tampered_trainer_nonce_fails(self): + token = generate_session_token() + scope = _make_scope() + executor_nonce = generate_nonce() + trainer_nonce = generate_nonce() + proof = compute_hello_proof(token, combine_nonces(executor_nonce, trainer_nonce), scope) + forged = combine_nonces(executor_nonce, generate_nonce()) + assert not verify_hello_proof(token, forged, scope, proof) + + def test_one_round_tampered_executor_nonce_fails(self): + token = generate_session_token() + scope = _make_scope() + executor_nonce = generate_nonce() + trainer_nonce = generate_nonce() + proof = compute_hello_proof(token, combine_nonces(executor_nonce, trainer_nonce), scope) + forged = combine_nonces(generate_nonce(), trainer_nonce) + assert not verify_hello_proof(token, forged, scope, proof) + + def test_combine_nonces_is_unambiguous(self): + # the split point between the two nonces must not be shiftable + assert combine_nonces("ab", "c") != combine_nonces("a", "bc") + + +class TestSessionTokenManager: + def test_successful_attach(self): + scope = _make_scope() + mgr = SessionTokenManager(scope=scope) + nonce = mgr.issue_nonce() + proof = compute_hello_proof(mgr.token, nonce, scope) + assert mgr.verify_proof(nonce, scope, proof) + assert mgr.session_active + + def test_consumed_nonce_rejected(self): + scope = _make_scope() + mgr = SessionTokenManager(scope=scope) + nonce = mgr.issue_nonce() + proof = compute_hello_proof(mgr.token, nonce, scope) + assert mgr.verify_proof(nonce, scope, proof) + # the same nonce (already consumed) must be rejected even with a valid proof + assert not mgr.verify_proof(nonce, scope, proof) + + def test_failed_attempt_also_consumes_nonce(self): + scope = _make_scope() + mgr = SessionTokenManager(scope=scope) + nonce = mgr.issue_nonce() + assert not mgr.verify_proof(nonce, scope, "bad-proof") + # nonce was consumed by the failed attempt; a now-valid proof must still fail + proof = compute_hello_proof(mgr.token, nonce, scope) + assert not mgr.verify_proof(nonce, scope, proof) + + def test_replayed_proof_against_fresh_nonce_fails(self): + scope = _make_scope() + mgr = SessionTokenManager(scope=scope) + nonce1 = mgr.issue_nonce() + observed_proof = compute_hello_proof(mgr.token, nonce1, scope) + # attacker observed the proof for nonce1 but is challenged with a fresh nonce + nonce2 = mgr.issue_nonce() + assert not mgr.verify_proof(nonce2, scope, observed_proof) + + def test_unissued_nonce_rejected(self): + scope = _make_scope() + mgr = SessionTokenManager(scope=scope) + nonce = generate_nonce() # not issued by the manager + proof = compute_hello_proof(mgr.token, nonce, scope) + assert not mgr.verify_proof(nonce, scope, proof) + + def test_scope_mismatch_rejected(self): + scope = _make_scope() + mgr = SessionTokenManager(scope=scope) + nonce = mgr.issue_nonce() + wrong_scope = _make_scope(attach_id="attach-other") + # proof is internally consistent with the wrong scope, but scope != expected scope + proof = compute_hello_proof(mgr.token, nonce, wrong_scope) + assert not mgr.verify_proof(nonce, wrong_scope, proof) + assert not mgr.session_active + + def test_wrong_protocol_version_rejected(self): + scope = _make_scope() + mgr = SessionTokenManager(scope=scope) + nonce = mgr.issue_nonce() + skewed = dataclasses.replace(scope, protocol_version=PROTOCOL_VERSION + 1) + proof = compute_hello_proof(mgr.token, nonce, skewed) + assert not mgr.verify_proof(nonce, skewed, proof) + + def test_expiry_after_attach_timeout(self): + clock = _FakeClock() + scope = _make_scope() + mgr = SessionTokenManager(scope=scope, attach_timeout=30.0, clock=clock) + nonce = mgr.issue_nonce() + proof = compute_hello_proof(mgr.token, nonce, scope) + clock.advance(30.1) + assert mgr.is_expired() + assert not mgr.verify_proof(nonce, scope, proof) + # issuing a fresh nonce after expiry is refused + with pytest.raises(RuntimeError): + mgr.issue_nonce() + + def test_attach_within_timeout_succeeds(self): + clock = _FakeClock() + scope = _make_scope() + mgr = SessionTokenManager(scope=scope, attach_timeout=30.0, clock=clock) + nonce = mgr.issue_nonce() + proof = compute_hello_proof(mgr.token, nonce, scope) + clock.advance(29.9) + assert not mgr.is_expired() + assert mgr.verify_proof(nonce, scope, proof) + + def test_active_session_does_not_expire(self): + clock = _FakeClock() + scope = _make_scope() + mgr = SessionTokenManager(scope=scope, attach_timeout=30.0, clock=clock) + nonce = mgr.issue_nonce() + proof = compute_hello_proof(mgr.token, nonce, scope) + assert mgr.verify_proof(nonce, scope, proof) + clock.advance(3600) + # after attach, the session is governed by heartbeat/job lifetime, not attach_timeout + assert not mgr.is_expired() + + def test_single_session_enforced(self): + scope = _make_scope() + mgr = SessionTokenManager(scope=scope) + # both nonces are issued before the session becomes active (issuing while a session is + # active is refused; see test_issue_nonce_refused_while_session_active) + nonce1 = mgr.issue_nonce() + nonce2 = mgr.issue_nonce() + proof1 = compute_hello_proof(mgr.token, nonce1, scope) + proof2 = compute_hello_proof(mgr.token, nonce2, scope) + assert mgr.verify_proof(nonce1, scope, proof1) + # a second, otherwise fully valid attach attempt must be rejected + assert not mgr.verify_proof(nonce2, scope, proof2) + assert mgr.session_active + # and no further nonce can be issued while the session is active + with pytest.raises(RuntimeError): + mgr.issue_nonce() + + def test_invalidate_kills_future_verification(self): + scope = _make_scope() + mgr = SessionTokenManager(scope=scope) + # a nonce issued before invalidate must not verify afterwards + nonce = mgr.issue_nonce() + proof = compute_hello_proof(mgr.token, nonce, scope) + mgr.invalidate() + assert mgr.invalidated + assert not mgr.session_active + assert not mgr.verify_proof(nonce, scope, proof) + # and issuing a new nonce after invalidate is refused + with pytest.raises(RuntimeError): + mgr.issue_nonce() + + def test_invalidate_after_attach_ends_session(self): + scope = _make_scope() + mgr = SessionTokenManager(scope=scope) + nonce = mgr.issue_nonce() + proof = compute_hello_proof(mgr.token, nonce, scope) + assert mgr.verify_proof(nonce, scope, proof) + mgr.invalidate() + assert not mgr.session_active + # after invalidate no fresh nonce can be issued and verification stays dead + with pytest.raises(RuntimeError): + mgr.issue_nonce() + + def test_provided_token_is_used(self): + token = generate_session_token() + mgr = SessionTokenManager(scope=_make_scope(), token=token) + assert mgr.token == token + assert mgr.digest == token_digest(token) + + def test_invalid_attach_timeout_rejected(self): + with pytest.raises(ValueError): + SessionTokenManager(scope=_make_scope(), attach_timeout=0) + with pytest.raises(ValueError): + SessionTokenManager(scope=_make_scope(), attach_timeout=-1.0) + + def test_invalid_scope_type_rejected(self): + with pytest.raises(TypeError): + SessionTokenManager(scope={"job_id": "job-1234"}) + + def test_non_ascii_proof_consumes_nonce_and_fails_without_raising(self): + scope = _make_scope() + mgr = SessionTokenManager(scope=scope) + nonce = mgr.issue_nonce() + # a non-ASCII proof must fail (not raise) and still consume the nonce + assert mgr.verify_proof(nonce, scope, "deadbeefÿ") is False + assert not mgr.session_active + # the nonce was consumed; a now-valid proof over the same nonce must still fail + proof = compute_hello_proof(mgr.token, nonce, scope) + assert not mgr.verify_proof(nonce, scope, proof) + + def test_issue_nonce_refused_while_session_active(self): + scope = _make_scope() + mgr = SessionTokenManager(scope=scope) + nonce = mgr.issue_nonce() + assert mgr.verify_proof(nonce, scope, compute_hello_proof(mgr.token, nonce, scope)) + assert mgr.session_active + with pytest.raises(RuntimeError): + mgr.issue_nonce() + + def test_issue_nonce_refused_after_invalidate(self): + mgr = SessionTokenManager(scope=_make_scope()) + mgr.invalidate() + with pytest.raises(RuntimeError): + mgr.issue_nonce() + + def test_issue_nonce_refused_after_expiry(self): + clock = _FakeClock() + mgr = SessionTokenManager(scope=_make_scope(), attach_timeout=30.0, clock=clock) + clock.advance(30.1) + with pytest.raises(RuntimeError): + mgr.issue_nonce() + + def test_pending_nonces_are_bounded_oldest_evicted(self): + scope = _make_scope() + mgr = SessionTokenManager(scope=scope) + cap = SessionTokenManager.MAX_PENDING_NONCES + first = mgr.issue_nonce() + # overflow the cap; the first (oldest) nonce must be evicted + for _ in range(cap + 50): + mgr.issue_nonce() + # the pending set never exceeds the cap (a HELLO flood cannot grow memory unboundedly) + assert len(mgr._pending_nonces) == cap + # the evicted oldest nonce no longer verifies even with a correct proof + assert not mgr.verify_proof(first, scope, compute_hello_proof(mgr.token, first, scope)) + + def test_hello_flood_keeps_memory_bounded(self): + mgr = SessionTokenManager(scope=_make_scope()) + cap = SessionTokenManager.MAX_PENDING_NONCES + for _ in range(cap * 3): + mgr.issue_nonce() + assert len(mgr._pending_nonces) <= cap + + def test_nonce_cannot_be_redeemed_after_attach_window(self): + clock = _FakeClock() + scope = _make_scope() + mgr = SessionTokenManager(scope=scope, attach_timeout=30.0, clock=clock) + nonce = mgr.issue_nonce() + proof = compute_hello_proof(mgr.token, nonce, scope) + # advance past the attach window; the manager is expired so the nonce cannot be redeemed + clock.advance(30.1) + assert not mgr.verify_proof(nonce, scope, proof) + + def test_short_token_rejected(self): + with pytest.raises(ValueError): + SessionTokenManager(scope=_make_scope(), token="a" * (MIN_TOKEN_HEX_CHARS - 1)) + + def test_minimum_length_token_accepted(self): + token = "a" * MIN_TOKEN_HEX_CHARS + mgr = SessionTokenManager(scope=_make_scope(), token=token) + assert mgr.token == token + + def test_non_str_token_rejected(self): + with pytest.raises(TypeError): + SessionTokenManager(scope=_make_scope(), token=12345678901234567890123456789012) diff --git a/tests/unit_test/client/cell/defs_test.py b/tests/unit_test/client/cell/defs_test.py new file mode 100644 index 0000000000..3eba49f251 --- /dev/null +++ b/tests/unit_test/client/cell/defs_test.py @@ -0,0 +1,121 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from nvflare.client.cell.defs import CHANNEL, PROTOCOL_VERSION, MsgKey, Topic +from nvflare.client.ipc import defs as ipc_defs + + +def _public_str_values(clazz) -> dict: + return {name: value for name, value in vars(clazz).items() if not name.startswith("_") and isinstance(value, str)} + + +class TestDefs: + def test_protocol_version(self): + assert PROTOCOL_VERSION == 1 + + def test_expected_topics_present(self): + expected = { + "HELLO", + "HELLO_CHALLENGE", + "HELLO_PROOF", + "HELLO_ACCEPTED", + "TASK_READY", + "TASK_ACCEPTED", + "TASK_FAILED", + "RESULT_READY", + "RESULT_ACCEPTED", + "RESULT_REJECTED", + "LOG", + "HEARTBEAT", + "ABORT", + "SHUTDOWN", + "BYE", + "ERROR", + } + assert expected == set(_public_str_values(Topic).keys()) + + def test_topic_values_unique(self): + values = list(_public_str_values(Topic).values()) + assert len(values) == len(set(values)) + + def test_msg_key_values_unique(self): + values = list(_public_str_values(MsgKey).values()) + assert len(values) == len(set(values)) + + def test_no_collision_with_legacy_ipc_defs(self): + # the Client API control protocol must not collide with the legacy FlareAgent + # channel/topic values in nvflare/client/ipc/defs.py + assert CHANNEL != ipc_defs.CHANNEL + legacy_topics = { + value for name, value in vars(ipc_defs).items() if name.startswith("TOPIC_") and isinstance(value, str) + } + assert not legacy_topics.intersection(_public_str_values(Topic).values()) + + def test_channel_wire_value(self): + # frozen cross-track wire constant: renaming this value is a protocol break + assert CHANNEL == "client_api" + + def test_topic_wire_values(self): + # exact frozen wire strings for every topic; a value rename must fail CI + expected = { + "HELLO": "client_api.hello", + "HELLO_CHALLENGE": "client_api.hello_challenge", + "HELLO_PROOF": "client_api.hello_proof", + "HELLO_ACCEPTED": "client_api.hello_accepted", + "TASK_READY": "client_api.task_ready", + "TASK_ACCEPTED": "client_api.task_accepted", + "TASK_FAILED": "client_api.task_failed", + "RESULT_READY": "client_api.result_ready", + "RESULT_ACCEPTED": "client_api.result_accepted", + "RESULT_REJECTED": "client_api.result_rejected", + "LOG": "client_api.log", + "HEARTBEAT": "client_api.heartbeat", + "ABORT": "client_api.abort", + "SHUTDOWN": "client_api.shutdown", + "BYE": "client_api.bye", + "ERROR": "client_api.error", + } + assert _public_str_values(Topic) == expected + # every topic is namespaced under the channel prefix + assert all(value.startswith("client_api.") for value in expected.values()) + + def test_msg_key_wire_values(self): + # exact frozen wire strings for every message key; a value rename must fail CI + expected = { + "SESSION_ID": "session_id", + "ATTACH_ID": "attach_id", + "JOB_ID": "job_id", + "SITE_NAME": "site_name", + "TRAINER_FQCN": "trainer_fqcn", + "TARGET_FQCN": "target_fqcn", + "RANK": "rank", + "RANK_POLICY": "rank_policy", + "PROTOCOL_VERSION": "protocol_version", + "NONCE": "nonce", + "PROOF": "proof", + "REASON": "reason", + "TASK_ID": "task_id", + "TASK_NAME": "task_name", + "MODEL": "model", + "PARAMS": "params", + "RESULT_ID": "result_id", + "TRANSFER_ID": "transfer_id", + "MANIFEST": "manifest", + } + assert _public_str_values(MsgKey) == expected + + def test_msg_key_names_present(self): + # the protocol/Appendix B wire keys must exist by name + for name in ("TASK_NAME", "MANIFEST", "RANK_POLICY", "MODEL", "PARAMS", "NONCE", "PROOF"): + assert hasattr(MsgKey, name) From a5d60ea18c38b4189bd382536131eda5434d1fe0 Mon Sep 17 00:00:00 2001 From: YuanTingHsieh Date: Thu, 2 Jul 2026 11:54:43 -0700 Subject: [PATCH 2/4] TE-1: trim to stateless proof toolkit; defer SessionTokenManager to attach MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Scope refinement (PR #4856): external_process — the primary mode, which shells out to torchrun/Deepspeed/Horovod/mpirun and talks only to rank 0 over a localhost connection NVFlare itself created — needs only a lightweight launch-token proof. The stateful executor-side SessionTokenManager (single-use nonce issuance, attach-window expiry, single-session enforcement, invalidation) is a challenge-response concern specific to attach mode, where NVFlare does not own the trainer process and the token is delivered out-of-band by an untrusted starter. Move SessionTokenManager to the attach backend (AT-2) so this freeze ships only what P0 (external_process) needs: the defs vocabulary and the stateless primitives (TokenScope, generate_session_token/generate_nonce, token_digest, compute_hello_proof/verify_hello_proof, combine_nonces). 30 tests remain (the manager's tests move with it). Co-Authored-By: Claude Fable 5 --- nvflare/client/cell/auth.py | 241 ++-------------------- tests/unit_test/client/cell/auth_test.py | 250 ----------------------- 2 files changed, 22 insertions(+), 469 deletions(-) diff --git a/nvflare/client/cell/auth.py b/nvflare/client/cell/auth.py index 9d8dbbb62c..ca5e7cd857 100644 --- a/nvflare/client/cell/auth.py +++ b/nvflare/client/cell/auth.py @@ -13,26 +13,28 @@ # limitations under the License. """Session-token auth primitives for the Client API control protocol. -Implements the attach auth contract of the Client API Execution Modes design -(docs/design/client_api_execution_modes.md, "Appendix B: Attach Topology and Auth"): - -- The executor generates a high-entropy session token per attach id / launch. The raw token - is held only in memory for the session lifetime; anything persisted stores only its digest. -- Token proof is challenge-response, not bearer presentation: HELLO -> HELLO_CHALLENGE (nonce) - -> HELLO_PROOF (HMAC keyed by the token over the nonce and the full token scope). The raw - token never crosses the wire, so observing attach traffic does not permit replay. -- The token is scoped to (job id, site name, attach id, target FQCN, trainer FQCN, rank - policy, protocol version), is single-session by default, and expires if the trainer does - not attach before attach_timeout. -- Two proof paths share the same HMAC primitive: - * The interactive HELLO_CHALLENGE / HELLO_PROOF path goes through SessionTokenManager, - which issues a self-generated single-use nonce (issue_nonce) and verifies the proof - against it (verify_proof). - * The one-round variant (Appendix B) folds the proof into HELLO, computed over an - executor nonce delivered in the bootstrap config combined with a trainer-supplied - nonce. This path does NOT go through SessionTokenManager.verify_proof (there is no - self-issued nonce); it uses the module-level verify_hello_proof directly over a - combined nonce (see combine_nonces). +Implements the token/proof primitives of the Client API Execution Modes design +(docs/design/client_api_execution_modes.md, "Control Protocol" HELLO handshake and +"Appendix B: Attach Topology and Auth"): + +- Token/nonce generation and the persistable digest: the raw token is held only in memory; + anything persisted stores only its digest (token_digest). +- The HELLO proof is challenge-response, not bearer presentation: the proof is an HMAC keyed + by the token over a challenge nonce and the full token scope, so the raw token never crosses + the wire and an observed proof is useless for any other nonce/scope (compute_hello_proof / + verify_hello_proof). +- TokenScope binds a token to (job id, site name, attach id, target FQCN, trainer FQCN, rank + policy, protocol version). +- One-round variant (Appendix B, for local/confidential channels such as an external_process + trainer on localhost): the proof is computed over an executor nonce combined with a + trainer-supplied nonce (combine_nonces) and verified with the module-level verify_hello_proof. + +Scope note: this module is the shared, stateless proof toolkit consumed by external_process +(EP-3) and attach (AT-2). The stateful, executor-side session manager -- single-use nonce +issuance, attach-window expiry, single-session enforcement, invalidation -- is an attach-mode +requirement (NVFlare does not own the trainer process there) and lands with the attach backend +(AT-2), not here. external_process launches the trainer itself on localhost and needs only the +lightweight launch-token proof these functions provide. This module is part of interface freeze #1. It is a pure library: no Cell/cellnet imports, no file I/O, no logging side effects. @@ -42,10 +44,6 @@ import hashlib import hmac import secrets -import threading -import time -from collections import OrderedDict -from typing import Callable, Optional from .defs import PROTOCOL_VERSION @@ -190,198 +188,3 @@ def combine_nonces(executor_nonce: str, trainer_nonce: str) -> str: if not isinstance(executor_nonce, str) or not isinstance(trainer_nonce, str): raise TypeError("executor_nonce and trainer_nonce must both be str") return f"{len(executor_nonce)}:{executor_nonce}:{trainer_nonce}" - - -class SessionTokenManager: - """Executor-side manager for one session token and its challenge-response verification. - - Holds the raw token in memory for the session lifetime (plus its digest for anything - persisted), issues single-use challenge nonces, and verifies HELLO_PROOF messages while - enforcing the auth contract: - - - a nonce is consumed by its first verification attempt (success or failure), so a - captured proof cannot be replayed and a consumed nonce cannot be reused; - - the presented scope must equal the expected scope exactly; - - the token expires if no session is established within attach_timeout of creation; - - the token is single-session: while a session is active, further attach attempts are - rejected; - - after invalidate(), all verification fails permanently (reconnect requires a fresh - token through a new manager); - - issue_nonce() refuses (raises RuntimeError) once the manager is invalidated, expired, or - a session is active, and the set of outstanding pending nonces is bounded by - MAX_PENDING_NONCES (oldest evicted first) so a HELLO flood cannot grow memory without - bound. - - Binding the accepted session to the peer's FQCN (recording/enforcing the bound peer) is - intentionally out of scope here and handled separately (see AT-2). - - The clock is injectable for testability; it must be a monotonic seconds counter. - Thread-safe. - """ - - # Upper bound on outstanding (issued-but-unconsumed) nonces; oldest are evicted past this. - MAX_PENDING_NONCES = 1024 - - def __init__( - self, - scope: TokenScope, - attach_timeout: Optional[float] = None, - token: Optional[str] = None, - clock: Callable[[], float] = time.monotonic, - ): - """Create a manager for one session token. - - Args: - scope: the scope the token is bound to. - attach_timeout: seconds from creation within which the trainer must attach; - None means no attach expiry. - token: the raw token to manage; a new high-entropy token is generated if None. - clock: monotonic seconds source (injectable for tests). - """ - if attach_timeout is not None and attach_timeout <= 0: - raise ValueError(f"attach_timeout must be positive but got {attach_timeout}") - if not isinstance(scope, TokenScope): - raise TypeError(f"scope must be TokenScope but got {type(scope)}") - if token is not None: - if not isinstance(token, str): - raise TypeError(f"token must be str but got {type(token)}") - if len(token) < MIN_TOKEN_HEX_CHARS: - raise ValueError( - f"token must have at least {MIN_TOKEN_HEX_CHARS} chars " - f"(>= 16 bytes of entropy) but got {len(token)}" - ) - - self._scope = scope - self._token = token if token else generate_session_token() - self._digest = token_digest(self._token) - self._attach_timeout = attach_timeout - self._clock = clock - self._created_at = clock() - # Insertion-ordered map nonce -> issued_at, so the oldest pending nonce is evictable - # in O(1) when the cap is exceeded, and each nonce carries an issuance time for TTL. - self._pending_nonces = OrderedDict() - self._session_active = False - self._invalidated = False - self._lock = threading.Lock() - - @property - def token(self) -> str: - """The raw session token (in-memory only; never persist or send on the wire).""" - return self._token - - @property - def digest(self) -> str: - """The persistable SHA-256 hex digest of the token.""" - return self._digest - - @property - def scope(self) -> TokenScope: - """The scope this token is bound to.""" - return self._scope - - @property - def session_active(self) -> bool: - """Whether a session is currently bound to this token.""" - with self._lock: - return self._session_active - - @property - def invalidated(self) -> bool: - """Whether this token has been explicitly invalidated.""" - with self._lock: - return self._invalidated - - def is_expired(self) -> bool: - """Whether the attach window has closed without a session being established.""" - with self._lock: - return self._is_expired() - - def _is_expired(self) -> bool: - if self._attach_timeout is None or self._session_active: - return False - return (self._clock() - self._created_at) > self._attach_timeout - - def issue_nonce(self) -> str: - """Issue a fresh single-use challenge nonce for HELLO_CHALLENGE. - - Refuses to issue once the manager is no longer in the pre-attach state, so a stale or - flooded caller cannot accumulate nonces after the session has been settled: - - Raises: - RuntimeError: if the token has been invalidated, the attach window has expired, or - a session is already active. - - The number of outstanding (issued-but-unconsumed) nonces is capped at - MAX_PENDING_NONCES; issuing past the cap evicts the oldest pending nonce so a HELLO - flood cannot grow memory without bound. - """ - with self._lock: - if self._invalidated: - raise RuntimeError("cannot issue nonce: session token has been invalidated") - if self._session_active: - raise RuntimeError("cannot issue nonce: a session is already active") - if self._is_expired(): - raise RuntimeError("cannot issue nonce: attach window has expired") - nonce = generate_nonce() - self._pending_nonces[nonce] = self._clock() - # bound memory: evict oldest outstanding nonces once the cap is exceeded - while len(self._pending_nonces) > self.MAX_PENDING_NONCES: - self._pending_nonces.popitem(last=False) - return nonce - - def verify_proof(self, nonce: str, scope: TokenScope, proof: str) -> bool: - """Verify a HELLO_PROOF attach attempt; on success the session becomes active. - - The nonce is consumed by this call whether or not verification succeeds. The proof is - verified against the manager's EXPECTED scope (self._scope), not the caller-supplied - scope: a wrong presented scope therefore yields a proof mismatch in constant time - (hmac.compare_digest) rather than leaking, via a short-circuiting tuple compare, which - scope field mismatched -- and the accepted proof is bound to the scope the executor - expects. (Recording/enforcing the bound peer FQCN is out of scope here; see AT-2.) - - Args: - nonce: the challenge nonce the proof claims to answer; must have been issued by - issue_nonce() and not yet consumed. - scope: the scope presented by the attaching trainer. - proof: the presented proof (hex string). - - Returns: - True only if the nonce is valid and unconsumed, the token is not invalidated or - expired, no session is already active, the proof verifies against the expected - scope, and (redundantly) the presented scope equals the expected scope. - """ - if not isinstance(nonce, str): - return False - with self._lock: - # single-use: consumed by any verification attempt (success or failure) - issued = self._pending_nonces.pop(nonce, None) is not None - - if self._invalidated: - return False - if not issued: - return False - # An old nonce cannot be redeemed after the attach window: _is_expired() covers - # it (a nonce's issue time is >= the manager's creation time under a monotonic - # clock, so the manager-level window always trips first) -- no separate per-nonce - # TTL is needed while no session is active. - if self._is_expired(): - return False - if self._session_active: - # single-session: reject further attach attempts while a session is active - return False - # constant-time proof check bound to the EXPECTED scope - if not verify_hello_proof(self._token, nonce, self._scope, proof): - return False - # redundant defense-in-depth guard (the proof is already bound to self._scope) - if scope != self._scope: - return False - - self._session_active = True - return True - - def invalidate(self) -> None: - """Invalidate the token and session permanently; all future verification fails.""" - with self._lock: - self._invalidated = True - self._session_active = False - self._pending_nonces.clear() diff --git a/tests/unit_test/client/cell/auth_test.py b/tests/unit_test/client/cell/auth_test.py index c5ced5040e..e7fb5cb2c5 100644 --- a/tests/unit_test/client/cell/auth_test.py +++ b/tests/unit_test/client/cell/auth_test.py @@ -18,8 +18,6 @@ import pytest from nvflare.client.cell.auth import ( - MIN_TOKEN_HEX_CHARS, - SessionTokenManager, TokenScope, combine_nonces, compute_hello_proof, @@ -31,19 +29,6 @@ from nvflare.client.cell.defs import PROTOCOL_VERSION -class _FakeClock: - """Injectable monotonic clock for expiry tests (no sleeps).""" - - def __init__(self, start: float = 1000.0): - self.now = start - - def advance(self, seconds: float): - self.now += seconds - - def __call__(self) -> float: - return self.now - - def _make_scope(**overrides) -> TokenScope: kwargs = dict( job_id="job-1234", @@ -65,10 +50,6 @@ def test_digest_is_stable_sha256_of_token(self): assert token_digest(token) == expected assert token_digest(token) == token_digest(token) - def test_manager_digest_matches_token_digest(self): - mgr = SessionTokenManager(scope=_make_scope()) - assert mgr.digest == token_digest(mgr.token) - def test_token_entropy_and_uniqueness(self): tokens = {generate_session_token() for _ in range(100)} assert len(tokens) == 100 @@ -199,234 +180,3 @@ def test_one_round_tampered_executor_nonce_fails(self): def test_combine_nonces_is_unambiguous(self): # the split point between the two nonces must not be shiftable assert combine_nonces("ab", "c") != combine_nonces("a", "bc") - - -class TestSessionTokenManager: - def test_successful_attach(self): - scope = _make_scope() - mgr = SessionTokenManager(scope=scope) - nonce = mgr.issue_nonce() - proof = compute_hello_proof(mgr.token, nonce, scope) - assert mgr.verify_proof(nonce, scope, proof) - assert mgr.session_active - - def test_consumed_nonce_rejected(self): - scope = _make_scope() - mgr = SessionTokenManager(scope=scope) - nonce = mgr.issue_nonce() - proof = compute_hello_proof(mgr.token, nonce, scope) - assert mgr.verify_proof(nonce, scope, proof) - # the same nonce (already consumed) must be rejected even with a valid proof - assert not mgr.verify_proof(nonce, scope, proof) - - def test_failed_attempt_also_consumes_nonce(self): - scope = _make_scope() - mgr = SessionTokenManager(scope=scope) - nonce = mgr.issue_nonce() - assert not mgr.verify_proof(nonce, scope, "bad-proof") - # nonce was consumed by the failed attempt; a now-valid proof must still fail - proof = compute_hello_proof(mgr.token, nonce, scope) - assert not mgr.verify_proof(nonce, scope, proof) - - def test_replayed_proof_against_fresh_nonce_fails(self): - scope = _make_scope() - mgr = SessionTokenManager(scope=scope) - nonce1 = mgr.issue_nonce() - observed_proof = compute_hello_proof(mgr.token, nonce1, scope) - # attacker observed the proof for nonce1 but is challenged with a fresh nonce - nonce2 = mgr.issue_nonce() - assert not mgr.verify_proof(nonce2, scope, observed_proof) - - def test_unissued_nonce_rejected(self): - scope = _make_scope() - mgr = SessionTokenManager(scope=scope) - nonce = generate_nonce() # not issued by the manager - proof = compute_hello_proof(mgr.token, nonce, scope) - assert not mgr.verify_proof(nonce, scope, proof) - - def test_scope_mismatch_rejected(self): - scope = _make_scope() - mgr = SessionTokenManager(scope=scope) - nonce = mgr.issue_nonce() - wrong_scope = _make_scope(attach_id="attach-other") - # proof is internally consistent with the wrong scope, but scope != expected scope - proof = compute_hello_proof(mgr.token, nonce, wrong_scope) - assert not mgr.verify_proof(nonce, wrong_scope, proof) - assert not mgr.session_active - - def test_wrong_protocol_version_rejected(self): - scope = _make_scope() - mgr = SessionTokenManager(scope=scope) - nonce = mgr.issue_nonce() - skewed = dataclasses.replace(scope, protocol_version=PROTOCOL_VERSION + 1) - proof = compute_hello_proof(mgr.token, nonce, skewed) - assert not mgr.verify_proof(nonce, skewed, proof) - - def test_expiry_after_attach_timeout(self): - clock = _FakeClock() - scope = _make_scope() - mgr = SessionTokenManager(scope=scope, attach_timeout=30.0, clock=clock) - nonce = mgr.issue_nonce() - proof = compute_hello_proof(mgr.token, nonce, scope) - clock.advance(30.1) - assert mgr.is_expired() - assert not mgr.verify_proof(nonce, scope, proof) - # issuing a fresh nonce after expiry is refused - with pytest.raises(RuntimeError): - mgr.issue_nonce() - - def test_attach_within_timeout_succeeds(self): - clock = _FakeClock() - scope = _make_scope() - mgr = SessionTokenManager(scope=scope, attach_timeout=30.0, clock=clock) - nonce = mgr.issue_nonce() - proof = compute_hello_proof(mgr.token, nonce, scope) - clock.advance(29.9) - assert not mgr.is_expired() - assert mgr.verify_proof(nonce, scope, proof) - - def test_active_session_does_not_expire(self): - clock = _FakeClock() - scope = _make_scope() - mgr = SessionTokenManager(scope=scope, attach_timeout=30.0, clock=clock) - nonce = mgr.issue_nonce() - proof = compute_hello_proof(mgr.token, nonce, scope) - assert mgr.verify_proof(nonce, scope, proof) - clock.advance(3600) - # after attach, the session is governed by heartbeat/job lifetime, not attach_timeout - assert not mgr.is_expired() - - def test_single_session_enforced(self): - scope = _make_scope() - mgr = SessionTokenManager(scope=scope) - # both nonces are issued before the session becomes active (issuing while a session is - # active is refused; see test_issue_nonce_refused_while_session_active) - nonce1 = mgr.issue_nonce() - nonce2 = mgr.issue_nonce() - proof1 = compute_hello_proof(mgr.token, nonce1, scope) - proof2 = compute_hello_proof(mgr.token, nonce2, scope) - assert mgr.verify_proof(nonce1, scope, proof1) - # a second, otherwise fully valid attach attempt must be rejected - assert not mgr.verify_proof(nonce2, scope, proof2) - assert mgr.session_active - # and no further nonce can be issued while the session is active - with pytest.raises(RuntimeError): - mgr.issue_nonce() - - def test_invalidate_kills_future_verification(self): - scope = _make_scope() - mgr = SessionTokenManager(scope=scope) - # a nonce issued before invalidate must not verify afterwards - nonce = mgr.issue_nonce() - proof = compute_hello_proof(mgr.token, nonce, scope) - mgr.invalidate() - assert mgr.invalidated - assert not mgr.session_active - assert not mgr.verify_proof(nonce, scope, proof) - # and issuing a new nonce after invalidate is refused - with pytest.raises(RuntimeError): - mgr.issue_nonce() - - def test_invalidate_after_attach_ends_session(self): - scope = _make_scope() - mgr = SessionTokenManager(scope=scope) - nonce = mgr.issue_nonce() - proof = compute_hello_proof(mgr.token, nonce, scope) - assert mgr.verify_proof(nonce, scope, proof) - mgr.invalidate() - assert not mgr.session_active - # after invalidate no fresh nonce can be issued and verification stays dead - with pytest.raises(RuntimeError): - mgr.issue_nonce() - - def test_provided_token_is_used(self): - token = generate_session_token() - mgr = SessionTokenManager(scope=_make_scope(), token=token) - assert mgr.token == token - assert mgr.digest == token_digest(token) - - def test_invalid_attach_timeout_rejected(self): - with pytest.raises(ValueError): - SessionTokenManager(scope=_make_scope(), attach_timeout=0) - with pytest.raises(ValueError): - SessionTokenManager(scope=_make_scope(), attach_timeout=-1.0) - - def test_invalid_scope_type_rejected(self): - with pytest.raises(TypeError): - SessionTokenManager(scope={"job_id": "job-1234"}) - - def test_non_ascii_proof_consumes_nonce_and_fails_without_raising(self): - scope = _make_scope() - mgr = SessionTokenManager(scope=scope) - nonce = mgr.issue_nonce() - # a non-ASCII proof must fail (not raise) and still consume the nonce - assert mgr.verify_proof(nonce, scope, "deadbeefÿ") is False - assert not mgr.session_active - # the nonce was consumed; a now-valid proof over the same nonce must still fail - proof = compute_hello_proof(mgr.token, nonce, scope) - assert not mgr.verify_proof(nonce, scope, proof) - - def test_issue_nonce_refused_while_session_active(self): - scope = _make_scope() - mgr = SessionTokenManager(scope=scope) - nonce = mgr.issue_nonce() - assert mgr.verify_proof(nonce, scope, compute_hello_proof(mgr.token, nonce, scope)) - assert mgr.session_active - with pytest.raises(RuntimeError): - mgr.issue_nonce() - - def test_issue_nonce_refused_after_invalidate(self): - mgr = SessionTokenManager(scope=_make_scope()) - mgr.invalidate() - with pytest.raises(RuntimeError): - mgr.issue_nonce() - - def test_issue_nonce_refused_after_expiry(self): - clock = _FakeClock() - mgr = SessionTokenManager(scope=_make_scope(), attach_timeout=30.0, clock=clock) - clock.advance(30.1) - with pytest.raises(RuntimeError): - mgr.issue_nonce() - - def test_pending_nonces_are_bounded_oldest_evicted(self): - scope = _make_scope() - mgr = SessionTokenManager(scope=scope) - cap = SessionTokenManager.MAX_PENDING_NONCES - first = mgr.issue_nonce() - # overflow the cap; the first (oldest) nonce must be evicted - for _ in range(cap + 50): - mgr.issue_nonce() - # the pending set never exceeds the cap (a HELLO flood cannot grow memory unboundedly) - assert len(mgr._pending_nonces) == cap - # the evicted oldest nonce no longer verifies even with a correct proof - assert not mgr.verify_proof(first, scope, compute_hello_proof(mgr.token, first, scope)) - - def test_hello_flood_keeps_memory_bounded(self): - mgr = SessionTokenManager(scope=_make_scope()) - cap = SessionTokenManager.MAX_PENDING_NONCES - for _ in range(cap * 3): - mgr.issue_nonce() - assert len(mgr._pending_nonces) <= cap - - def test_nonce_cannot_be_redeemed_after_attach_window(self): - clock = _FakeClock() - scope = _make_scope() - mgr = SessionTokenManager(scope=scope, attach_timeout=30.0, clock=clock) - nonce = mgr.issue_nonce() - proof = compute_hello_proof(mgr.token, nonce, scope) - # advance past the attach window; the manager is expired so the nonce cannot be redeemed - clock.advance(30.1) - assert not mgr.verify_proof(nonce, scope, proof) - - def test_short_token_rejected(self): - with pytest.raises(ValueError): - SessionTokenManager(scope=_make_scope(), token="a" * (MIN_TOKEN_HEX_CHARS - 1)) - - def test_minimum_length_token_accepted(self): - token = "a" * MIN_TOKEN_HEX_CHARS - mgr = SessionTokenManager(scope=_make_scope(), token=token) - assert mgr.token == token - - def test_non_str_token_rejected(self): - with pytest.raises(TypeError): - SessionTokenManager(scope=_make_scope(), token=12345678901234567890123456789012) From bdb520abf9dcb5b3a8a34792a89b966137f4616a Mon Sep 17 00:00:00 2001 From: YuanTingHsieh Date: Thu, 2 Jul 2026 14:20:28 -0700 Subject: [PATCH 3/4] =?UTF-8?q?TE-1:=20address=20review=20=E2=80=94=20proo?= =?UTF-8?q?f=20length=20guard,=20entropy=20floors,=20HELLO=5FREJECTED?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR #4856 review fixes: - verify_hello_proof rejects any proof whose length is not the fixed SHA-256 hex length (PROOF_HEX_LEN=64) up front, so an arbitrarily long attacker proof cannot force HMAC/compare work; also guards AttributeError so a non-str token/nonce (e.g. a None token from a config miss on the one-round path) returns False instead of raising. - generate_session_token / generate_nonce reject below a 128-bit entropy floor (_MIN_ENTROPY_BYTES=16) so a public primitive can't mint an empty or weak token/nonce (num_bytes=0). - add Topic.HELLO_REJECTED so a handshake auth refusal (bad proof, wrong scope, consumed/expired nonce, single-session) is a distinct semantic signal from the generic ERROR (protocol/transport fault) for TE-3/AT-2. Tests cover the entropy floors, wrong/huge proof lengths, non-str token/nonce, a 64-char non-ASCII proof reaching compare_digest, and the HELLO_REJECTED wire value. Co-Authored-By: Claude Fable 5 --- nvflare/client/cell/auth.py | 32 ++++++++++++++++---- nvflare/client/cell/defs.py | 4 +++ tests/unit_test/client/cell/auth_test.py | 38 ++++++++++++++++++++++-- tests/unit_test/client/cell/defs_test.py | 2 ++ 4 files changed, 69 insertions(+), 7 deletions(-) diff --git a/nvflare/client/cell/auth.py b/nvflare/client/cell/auth.py index ca5e7cd857..5b78f347b3 100644 --- a/nvflare/client/cell/auth.py +++ b/nvflare/client/cell/auth.py @@ -56,6 +56,12 @@ # Minimum accepted length (hex chars) for a caller-supplied token: 16 bytes of entropy. MIN_TOKEN_HEX_CHARS = 32 +# Minimum bytes of entropy accepted by the token/nonce generators (128-bit floor). +_MIN_ENTROPY_BYTES = 16 +# The HELLO proof is a SHA-256 HMAC hex digest, so its length is fixed. A presented proof of +# any other length cannot match and is rejected up front (avoids HMAC/compare work on an +# arbitrarily long attacker-supplied string). +PROOF_HEX_LEN = hashlib.sha256().digest_size * 2 # 64 @dataclasses.dataclass(frozen=True) @@ -93,12 +99,23 @@ def generate_session_token(num_bytes: int = DEFAULT_TOKEN_BYTES) -> str: Returns: The token as a hex string. The raw token must be held only in memory; anything persisted must store only its digest (see token_digest). + + Raises: + ValueError: if num_bytes is below the 128-bit entropy floor (_MIN_ENTROPY_BYTES). """ + if num_bytes < _MIN_ENTROPY_BYTES: + raise ValueError(f"num_bytes must be >= {_MIN_ENTROPY_BYTES} (128-bit floor) but got {num_bytes}") return secrets.token_hex(num_bytes) def generate_nonce(num_bytes: int = DEFAULT_NONCE_BYTES) -> str: - """Generate a new single-use challenge nonce as a hex string.""" + """Generate a new single-use challenge nonce as a hex string. + + Raises: + ValueError: if num_bytes is below the 128-bit entropy floor (_MIN_ENTROPY_BYTES). + """ + if num_bytes < _MIN_ENTROPY_BYTES: + raise ValueError(f"num_bytes must be >= {_MIN_ENTROPY_BYTES} (128-bit floor) but got {num_bytes}") return secrets.token_hex(num_bytes) @@ -160,18 +177,23 @@ def verify_hello_proof(token: str, nonce: str, scope: TokenScope, proof: str) -> proof: the presented proof (hex string). Returns: - True if the proof is valid for (token, nonce, scope). Never raises for a - malformed/attacker-supplied proof: a non-str, empty, or non-ASCII proof returns False. + True if the proof is valid for (token, nonce, scope). Never raises for + malformed/attacker-supplied input: a non-str, wrong-length, or non-ASCII proof returns + False, and a non-str token/nonce (e.g. a None token from a config miss on the one-round + path) returns False rather than raising. """ - if not isinstance(proof, str) or not proof: + # A valid proof is always a SHA-256 hex digest of fixed length; reject any other length up + # front so an arbitrarily long attacker-supplied proof cannot force HMAC/compare work. + if not isinstance(proof, str) or len(proof) != PROOF_HEX_LEN: return False try: expected = compute_hello_proof(token, nonce, scope) # Encode both operands to bytes before compare_digest: hmac.compare_digest raises # TypeError on a non-ASCII str, so a non-ASCII (attacker-supplied) proof would raise # instead of failing. Comparing bytes makes any malformed proof simply mismatch. + # AttributeError covers a non-str token/nonce (e.g. .encode on None) on the one-round path. return hmac.compare_digest(expected.encode("utf-8"), proof.encode("utf-8")) - except (TypeError, ValueError): + except (TypeError, ValueError, AttributeError): return False diff --git a/nvflare/client/cell/defs.py b/nvflare/client/cell/defs.py index c4ca97884f..57ae47d64f 100644 --- a/nvflare/client/cell/defs.py +++ b/nvflare/client/cell/defs.py @@ -48,6 +48,10 @@ class Topic: HELLO_CHALLENGE = "client_api.hello_challenge" HELLO_PROOF = "client_api.hello_proof" HELLO_ACCEPTED = "client_api.hello_accepted" + # Distinct from ERROR (a protocol/transport error): HELLO_REJECTED is a clean, semantic + # auth/handshake refusal (bad proof, wrong scope, consumed/expired nonce, single-session), + # so TE-3/AT-2 state machines can tell a recoverable auth failure from a protocol fault. + HELLO_REJECTED = "client_api.hello_rejected" # Per task (every round) TASK_READY = "client_api.task_ready" diff --git a/tests/unit_test/client/cell/auth_test.py b/tests/unit_test/client/cell/auth_test.py index e7fb5cb2c5..1509f0b161 100644 --- a/tests/unit_test/client/cell/auth_test.py +++ b/tests/unit_test/client/cell/auth_test.py @@ -60,6 +60,17 @@ def test_nonce_uniqueness(self): nonces = {generate_nonce() for _ in range(100)} assert len(nonces) == 100 + def test_generators_reject_below_entropy_floor(self): + # public primitives must not mint an empty/weak token or nonce + for bad in (0, -1, 8, 15): + with pytest.raises(ValueError): + generate_session_token(bad) + with pytest.raises(ValueError): + generate_nonce(bad) + # the floor itself is accepted + assert len(generate_session_token(16)) == 32 + assert len(generate_nonce(16)) == 32 + class TestHelloProof: def test_proof_round_trip(self): @@ -116,6 +127,28 @@ def test_empty_or_non_str_proof_fails(self): assert not verify_hello_proof(token, nonce, scope, "") assert not verify_hello_proof(token, nonce, scope, None) + def test_wrong_length_proof_rejected(self): + # a valid proof is a fixed-length SHA-256 hex digest; any other length is rejected + # up front (no HMAC/compare work on an arbitrarily long attacker string) + from nvflare.client.cell.auth import PROOF_HEX_LEN + + token = generate_session_token() + nonce = generate_nonce() + scope = _make_scope() + assert PROOF_HEX_LEN == 64 + assert not verify_hello_proof(token, nonce, scope, "ab") # too short + assert not verify_hello_proof(token, nonce, scope, "a" * (PROOF_HEX_LEN - 1)) + assert not verify_hello_proof(token, nonce, scope, "a" * (PROOF_HEX_LEN + 1)) + assert not verify_hello_proof(token, nonce, scope, "a" * 100_000) # huge + + def test_non_str_token_or_nonce_returns_false_without_raising(self): + # the one-round path calls verify_hello_proof directly; a None token/nonce (e.g. a + # config miss) must return False, not raise AttributeError + scope = _make_scope() + proof = "a" * 64 + assert verify_hello_proof(None, generate_nonce(), scope, proof) is False + assert verify_hello_proof(generate_session_token(), None, scope, proof) is False + def test_serialization_is_unambiguous_across_field_boundaries(self): # shifting bytes between adjacent fields must not produce the same proof token = generate_session_token() @@ -126,11 +159,12 @@ def test_serialization_is_unambiguous_across_field_boundaries(self): def test_non_ascii_proof_returns_false_without_raising(self): # hmac.compare_digest raises TypeError on a non-ASCII str; verify must return False. + # Use a 64-char proof so it passes the length check and actually reaches compare_digest. token = generate_session_token() nonce = generate_nonce() scope = _make_scope() - assert verify_hello_proof(token, nonce, scope, "deadbeefÿ") is False - assert verify_hello_proof(token, nonce, scope, "中文") is False + assert verify_hello_proof(token, nonce, scope, "a" * 63 + "ÿ") is False + assert verify_hello_proof(token, nonce, scope, "中" * 32) is False def test_cross_typed_scope_produces_different_proof_and_fails(self): # protocol_version=1 (int) and protocol_version="1" (str) must not be interchangeable. diff --git a/tests/unit_test/client/cell/defs_test.py b/tests/unit_test/client/cell/defs_test.py index 3eba49f251..b5644a09f2 100644 --- a/tests/unit_test/client/cell/defs_test.py +++ b/tests/unit_test/client/cell/defs_test.py @@ -30,6 +30,7 @@ def test_expected_topics_present(self): "HELLO_CHALLENGE", "HELLO_PROOF", "HELLO_ACCEPTED", + "HELLO_REJECTED", "TASK_READY", "TASK_ACCEPTED", "TASK_FAILED", @@ -73,6 +74,7 @@ def test_topic_wire_values(self): "HELLO_CHALLENGE": "client_api.hello_challenge", "HELLO_PROOF": "client_api.hello_proof", "HELLO_ACCEPTED": "client_api.hello_accepted", + "HELLO_REJECTED": "client_api.hello_rejected", "TASK_READY": "client_api.task_ready", "TASK_ACCEPTED": "client_api.task_accepted", "TASK_FAILED": "client_api.task_failed", From 4bd2d09a9961384d67c0067427151bc31079c93c Mon Sep 17 00:00:00 2001 From: YuanTingHsieh Date: Thu, 2 Jul 2026 15:36:45 -0700 Subject: [PATCH 4/4] TE-1: reduce freeze to the protocol vocabulary; drop auth mechanism MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Trim auth.py (and its tests) out of the interface freeze. The token is three roles — rendezvous, anti-mixup, and authentication — and whether external_process needs a *secret* HMAC proof at all is a host-trust decision (single-tenant: OS isolation suffices; multi-tenant: needed) that belongs to EP-3, not to a frozen module shipped ahead of it. #4856 now freezes only what is undisputed and shared: the Cell control- protocol vocabulary (client/cell/defs.py — Topics, MsgKeys, CHANNEL, PROTOCOL_VERSION). The proof helpers (TokenScope, compute/verify_hello_proof, combine_nonces) move to EP-3 where the auth model is decided and first consumed; the generic generators (generate_session_token/nonce, token_digest) go to the fuel/sec consolidation (FLARE-3017); the stateful SessionTokenManager remains attach-only (AT-2). Plan updated to match. The removed code is preserved in this branch's history for EP-3 to draw on. Co-Authored-By: Claude Fable 5 --- nvflare/client/cell/auth.py | 212 ---------------------- tests/unit_test/client/cell/auth_test.py | 216 ----------------------- 2 files changed, 428 deletions(-) delete mode 100644 nvflare/client/cell/auth.py delete mode 100644 tests/unit_test/client/cell/auth_test.py diff --git a/nvflare/client/cell/auth.py b/nvflare/client/cell/auth.py deleted file mode 100644 index 5b78f347b3..0000000000 --- a/nvflare/client/cell/auth.py +++ /dev/null @@ -1,212 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Session-token auth primitives for the Client API control protocol. - -Implements the token/proof primitives of the Client API Execution Modes design -(docs/design/client_api_execution_modes.md, "Control Protocol" HELLO handshake and -"Appendix B: Attach Topology and Auth"): - -- Token/nonce generation and the persistable digest: the raw token is held only in memory; - anything persisted stores only its digest (token_digest). -- The HELLO proof is challenge-response, not bearer presentation: the proof is an HMAC keyed - by the token over a challenge nonce and the full token scope, so the raw token never crosses - the wire and an observed proof is useless for any other nonce/scope (compute_hello_proof / - verify_hello_proof). -- TokenScope binds a token to (job id, site name, attach id, target FQCN, trainer FQCN, rank - policy, protocol version). -- One-round variant (Appendix B, for local/confidential channels such as an external_process - trainer on localhost): the proof is computed over an executor nonce combined with a - trainer-supplied nonce (combine_nonces) and verified with the module-level verify_hello_proof. - -Scope note: this module is the shared, stateless proof toolkit consumed by external_process -(EP-3) and attach (AT-2). The stateful, executor-side session manager -- single-use nonce -issuance, attach-window expiry, single-session enforcement, invalidation -- is an attach-mode -requirement (NVFlare does not own the trainer process there) and lands with the attach backend -(AT-2), not here. external_process launches the trainer itself on localhost and needs only the -lightweight launch-token proof these functions provide. - -This module is part of interface freeze #1. It is a pure library: no Cell/cellnet imports, -no file I/O, no logging side effects. -""" - -import dataclasses -import hashlib -import hmac -import secrets - -from .defs import PROTOCOL_VERSION - -# Domain-separation tag for the HELLO proof HMAC, versioned with the proof format. -_PROOF_DOMAIN_TAG = "nvflare.client_api.hello_proof.v1" - -# Default sizes (bytes of entropy) for generated secrets. -DEFAULT_TOKEN_BYTES = 32 -DEFAULT_NONCE_BYTES = 16 - -# Minimum accepted length (hex chars) for a caller-supplied token: 16 bytes of entropy. -MIN_TOKEN_HEX_CHARS = 32 -# Minimum bytes of entropy accepted by the token/nonce generators (128-bit floor). -_MIN_ENTROPY_BYTES = 16 -# The HELLO proof is a SHA-256 HMAC hex digest, so its length is fixed. A presented proof of -# any other length cannot match and is rejected up front (avoids HMAC/compare work on an -# arbitrarily long attacker-supplied string). -PROOF_HEX_LEN = hashlib.sha256().digest_size * 2 # 64 - - -@dataclasses.dataclass(frozen=True) -class TokenScope: - """The scope a session token is bound to. - - A token is valid only for this exact scope; the executor accepts an attach only when the - proof, the presented scope, and its own expected scope all match. - - Attributes: - job_id: id of the job the token was issued for. - site_name: name of the site (FL client) the token was issued for. - attach_id: attach/session id the token was issued for. - target_fqcn: FQCN of the CJ/job cell the trainer must talk to. - trainer_fqcn: FQCN the trainer cell must bind (routing name, not identity). - rank_policy: allowed rank policy for the session (e.g. which rank may attach). - protocol_version: Client API control protocol version. - """ - - job_id: str - site_name: str - attach_id: str - target_fqcn: str - trainer_fqcn: str - rank_policy: str - protocol_version: int = PROTOCOL_VERSION - - -def generate_session_token(num_bytes: int = DEFAULT_TOKEN_BYTES) -> str: - """Generate a new high-entropy session token. - - Args: - num_bytes: bytes of entropy in the token. - - Returns: - The token as a hex string. The raw token must be held only in memory; anything - persisted must store only its digest (see token_digest). - - Raises: - ValueError: if num_bytes is below the 128-bit entropy floor (_MIN_ENTROPY_BYTES). - """ - if num_bytes < _MIN_ENTROPY_BYTES: - raise ValueError(f"num_bytes must be >= {_MIN_ENTROPY_BYTES} (128-bit floor) but got {num_bytes}") - return secrets.token_hex(num_bytes) - - -def generate_nonce(num_bytes: int = DEFAULT_NONCE_BYTES) -> str: - """Generate a new single-use challenge nonce as a hex string. - - Raises: - ValueError: if num_bytes is below the 128-bit entropy floor (_MIN_ENTROPY_BYTES). - """ - if num_bytes < _MIN_ENTROPY_BYTES: - raise ValueError(f"num_bytes must be >= {_MIN_ENTROPY_BYTES} (128-bit floor) but got {num_bytes}") - return secrets.token_hex(num_bytes) - - -def token_digest(token: str) -> str: - """Compute the persistable digest (SHA-256 hex) of a session token.""" - return hashlib.sha256(token.encode("utf-8")).hexdigest() - - -def _canonical_proof_message(nonce: str, scope: TokenScope) -> bytes: - """Serialize (nonce + all scope fields) canonically, unambiguously, and type-unambiguously. - - Each item is netstring-encoded (length-prefixed), so no combination of field values can - produce the same byte stream as a different combination (no concatenation ambiguity). - Each scope field additionally carries an explicit type tag (its Python type name) and is - rendered with repr(), so values that differ only by type -- e.g. protocol_version=1 (int) - vs protocol_version="1" (str) -- serialize differently and cannot be substituted for one - another (no cross-type proof forgery). Field order is fixed: domain tag, nonce, then - TokenScope fields in declaration order. - """ - items = [_PROOF_DOMAIN_TAG, nonce] - for f in dataclasses.fields(scope): - value = getattr(scope, f.name) - items.append(f"{type(value).__name__}:{value!r}") - - buf = bytearray() - for item in items: - encoded = item.encode("utf-8") - buf += f"{len(encoded)}:".encode("utf-8") - buf += encoded - buf += b"," - return bytes(buf) - - -def compute_hello_proof(token: str, nonce: str, scope: TokenScope) -> str: - """Compute the HELLO_PROOF value: HMAC-SHA256 keyed by the token over (nonce, scope). - - The proof covers the challenge nonce and the full token scope (attach id, job id, site - name, trainer FQCN, target FQCN, rank policy, protocol version), so a proof observed on - the wire is useless for any other nonce, scope, or protocol version. - - Args: - token: the raw session token (the HMAC key; never sent on the wire). - nonce: the challenge nonce from HELLO_CHALLENGE. - scope: the token scope the proof is computed over. - - Returns: - The proof as a hex string. - """ - return hmac.new(token.encode("utf-8"), _canonical_proof_message(nonce, scope), hashlib.sha256).hexdigest() - - -def verify_hello_proof(token: str, nonce: str, scope: TokenScope, proof: str) -> bool: - """Verify a HELLO_PROOF value in constant time. - - Args: - token: the raw session token. - nonce: the challenge nonce the proof must cover. - scope: the token scope the proof must cover. - proof: the presented proof (hex string). - - Returns: - True if the proof is valid for (token, nonce, scope). Never raises for - malformed/attacker-supplied input: a non-str, wrong-length, or non-ASCII proof returns - False, and a non-str token/nonce (e.g. a None token from a config miss on the one-round - path) returns False rather than raising. - """ - # A valid proof is always a SHA-256 hex digest of fixed length; reject any other length up - # front so an arbitrarily long attacker-supplied proof cannot force HMAC/compare work. - if not isinstance(proof, str) or len(proof) != PROOF_HEX_LEN: - return False - try: - expected = compute_hello_proof(token, nonce, scope) - # Encode both operands to bytes before compare_digest: hmac.compare_digest raises - # TypeError on a non-ASCII str, so a non-ASCII (attacker-supplied) proof would raise - # instead of failing. Comparing bytes makes any malformed proof simply mismatch. - # AttributeError covers a non-str token/nonce (e.g. .encode on None) on the one-round path. - return hmac.compare_digest(expected.encode("utf-8"), proof.encode("utf-8")) - except (TypeError, ValueError, AttributeError): - return False - - -def combine_nonces(executor_nonce: str, trainer_nonce: str) -> str: - """Combine an executor nonce with a trainer nonce for the one-round HELLO proof variant. - - In the one-round variant (Appendix B), the executor delivers a nonce in the bootstrap - config and the trainer contributes its own nonce; the proof is computed over both so - neither side alone fixes the challenge. The combination is length-prefixed so the pair is - unambiguous regardless of nonce contents. Both sides must combine identically, then feed - the result as the ``nonce`` argument to compute_hello_proof / verify_hello_proof -- this - path does not use SessionTokenManager (there is no self-issued nonce). - """ - if not isinstance(executor_nonce, str) or not isinstance(trainer_nonce, str): - raise TypeError("executor_nonce and trainer_nonce must both be str") - return f"{len(executor_nonce)}:{executor_nonce}:{trainer_nonce}" diff --git a/tests/unit_test/client/cell/auth_test.py b/tests/unit_test/client/cell/auth_test.py deleted file mode 100644 index 1509f0b161..0000000000 --- a/tests/unit_test/client/cell/auth_test.py +++ /dev/null @@ -1,216 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import dataclasses -import hashlib - -import pytest - -from nvflare.client.cell.auth import ( - TokenScope, - combine_nonces, - compute_hello_proof, - generate_nonce, - generate_session_token, - token_digest, - verify_hello_proof, -) -from nvflare.client.cell.defs import PROTOCOL_VERSION - - -def _make_scope(**overrides) -> TokenScope: - kwargs = dict( - job_id="job-1234", - site_name="site-1", - attach_id="attach-abcd", - target_fqcn="site-1.job-1234", - trainer_fqcn="site-1.-client_api_attach-abcd", - rank_policy="0", - protocol_version=PROTOCOL_VERSION, - ) - kwargs.update(overrides) - return TokenScope(**kwargs) - - -class TestTokenAndDigest: - def test_digest_is_stable_sha256_of_token(self): - token = generate_session_token() - expected = hashlib.sha256(token.encode("utf-8")).hexdigest() - assert token_digest(token) == expected - assert token_digest(token) == token_digest(token) - - def test_token_entropy_and_uniqueness(self): - tokens = {generate_session_token() for _ in range(100)} - assert len(tokens) == 100 - # default 32 bytes of entropy -> 64 hex chars - assert all(len(t) == 64 for t in tokens) - - def test_nonce_uniqueness(self): - nonces = {generate_nonce() for _ in range(100)} - assert len(nonces) == 100 - - def test_generators_reject_below_entropy_floor(self): - # public primitives must not mint an empty/weak token or nonce - for bad in (0, -1, 8, 15): - with pytest.raises(ValueError): - generate_session_token(bad) - with pytest.raises(ValueError): - generate_nonce(bad) - # the floor itself is accepted - assert len(generate_session_token(16)) == 32 - assert len(generate_nonce(16)) == 32 - - -class TestHelloProof: - def test_proof_round_trip(self): - token = generate_session_token() - nonce = generate_nonce() - scope = _make_scope() - proof = compute_hello_proof(token, nonce, scope) - assert verify_hello_proof(token, nonce, scope, proof) - - def test_tampered_nonce_fails(self): - token = generate_session_token() - nonce = generate_nonce() - scope = _make_scope() - proof = compute_hello_proof(token, nonce, scope) - assert not verify_hello_proof(token, generate_nonce(), scope, proof) - - @pytest.mark.parametrize( - "field,value", - [ - ("job_id", "job-9999"), - ("site_name", "site-2"), - ("attach_id", "attach-zzzz"), - ("target_fqcn", "site-2.job-9999"), - ("trainer_fqcn", "site-1.-client_api_other"), - ("rank_policy", "1"), - ], - ) - def test_tampered_scope_field_fails(self, field, value): - token = generate_session_token() - nonce = generate_nonce() - scope = _make_scope() - proof = compute_hello_proof(token, nonce, scope) - tampered = dataclasses.replace(scope, **{field: value}) - assert not verify_hello_proof(token, nonce, tampered, proof) - - def test_wrong_protocol_version_fails(self): - token = generate_session_token() - nonce = generate_nonce() - scope = _make_scope() - proof = compute_hello_proof(token, nonce, scope) - skewed = dataclasses.replace(scope, protocol_version=PROTOCOL_VERSION + 1) - assert not verify_hello_proof(token, nonce, skewed, proof) - - def test_wrong_token_fails(self): - nonce = generate_nonce() - scope = _make_scope() - proof = compute_hello_proof(generate_session_token(), nonce, scope) - assert not verify_hello_proof(generate_session_token(), nonce, scope, proof) - - def test_empty_or_non_str_proof_fails(self): - token = generate_session_token() - nonce = generate_nonce() - scope = _make_scope() - assert not verify_hello_proof(token, nonce, scope, "") - assert not verify_hello_proof(token, nonce, scope, None) - - def test_wrong_length_proof_rejected(self): - # a valid proof is a fixed-length SHA-256 hex digest; any other length is rejected - # up front (no HMAC/compare work on an arbitrarily long attacker string) - from nvflare.client.cell.auth import PROOF_HEX_LEN - - token = generate_session_token() - nonce = generate_nonce() - scope = _make_scope() - assert PROOF_HEX_LEN == 64 - assert not verify_hello_proof(token, nonce, scope, "ab") # too short - assert not verify_hello_proof(token, nonce, scope, "a" * (PROOF_HEX_LEN - 1)) - assert not verify_hello_proof(token, nonce, scope, "a" * (PROOF_HEX_LEN + 1)) - assert not verify_hello_proof(token, nonce, scope, "a" * 100_000) # huge - - def test_non_str_token_or_nonce_returns_false_without_raising(self): - # the one-round path calls verify_hello_proof directly; a None token/nonce (e.g. a - # config miss) must return False, not raise AttributeError - scope = _make_scope() - proof = "a" * 64 - assert verify_hello_proof(None, generate_nonce(), scope, proof) is False - assert verify_hello_proof(generate_session_token(), None, scope, proof) is False - - def test_serialization_is_unambiguous_across_field_boundaries(self): - # shifting bytes between adjacent fields must not produce the same proof - token = generate_session_token() - nonce = generate_nonce() - scope_a = _make_scope(job_id="ab", site_name="c") - scope_b = _make_scope(job_id="a", site_name="bc") - assert compute_hello_proof(token, nonce, scope_a) != compute_hello_proof(token, nonce, scope_b) - - def test_non_ascii_proof_returns_false_without_raising(self): - # hmac.compare_digest raises TypeError on a non-ASCII str; verify must return False. - # Use a 64-char proof so it passes the length check and actually reaches compare_digest. - token = generate_session_token() - nonce = generate_nonce() - scope = _make_scope() - assert verify_hello_proof(token, nonce, scope, "a" * 63 + "ÿ") is False - assert verify_hello_proof(token, nonce, scope, "中" * 32) is False - - def test_cross_typed_scope_produces_different_proof_and_fails(self): - # protocol_version=1 (int) and protocol_version="1" (str) must not be interchangeable. - token = generate_session_token() - nonce = generate_nonce() - scope_int = _make_scope(protocol_version=1) - scope_str = _make_scope(protocol_version="1") - proof_int = compute_hello_proof(token, nonce, scope_int) - proof_str = compute_hello_proof(token, nonce, scope_str) - assert proof_int != proof_str - # a proof computed over the int scope must not verify for the str scope, and vice versa - assert not verify_hello_proof(token, nonce, scope_str, proof_int) - assert not verify_hello_proof(token, nonce, scope_int, proof_str) - - -class TestOneRoundProof: - """The one-round variant folds the proof into HELLO over a combined executor+trainer nonce - and uses the module-level functions directly (no SessionTokenManager / issue_nonce).""" - - def test_one_round_round_trip(self): - token = generate_session_token() - scope = _make_scope() - executor_nonce = generate_nonce() # delivered in bootstrap config - trainer_nonce = generate_nonce() # contributed by the trainer - combined = combine_nonces(executor_nonce, trainer_nonce) - proof = compute_hello_proof(token, combined, scope) - assert verify_hello_proof(token, combined, scope, proof) - - def test_one_round_tampered_trainer_nonce_fails(self): - token = generate_session_token() - scope = _make_scope() - executor_nonce = generate_nonce() - trainer_nonce = generate_nonce() - proof = compute_hello_proof(token, combine_nonces(executor_nonce, trainer_nonce), scope) - forged = combine_nonces(executor_nonce, generate_nonce()) - assert not verify_hello_proof(token, forged, scope, proof) - - def test_one_round_tampered_executor_nonce_fails(self): - token = generate_session_token() - scope = _make_scope() - executor_nonce = generate_nonce() - trainer_nonce = generate_nonce() - proof = compute_hello_proof(token, combine_nonces(executor_nonce, trainer_nonce), scope) - forged = combine_nonces(generate_nonce(), trainer_nonce) - assert not verify_hello_proof(token, forged, scope, proof) - - def test_combine_nonces_is_unambiguous(self): - # the split point between the two nonces must not be shiftable - assert combine_nonces("ab", "c") != combine_nonces("a", "bc")