flwr 1.22.0__py3-none-any.whl → 1.24.0__py3-none-any.whl

flwr/server/superlink/fleet/grpc_rere/fleet_servicer.py

@@ -15,25 +15,31 @@
 """Fleet API gRPC request-response servicer."""
 
 
-from logging import DEBUG, INFO
+import threading
+from logging import DEBUG, ERROR, INFO
 
 import grpc
 from google.protobuf.json_format import MessageToDict
 
+from flwr.common.constant import PUBLIC_KEY_ALREADY_IN_USE_MESSAGE
 from flwr.common.inflatable import UnexpectedObjectContentError
 from flwr.common.logger import log
 from flwr.common.typing import InvalidRunStatusException
 from flwr.proto import fleet_pb2_grpc  # pylint: disable=E0611
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
-    CreateNodeRequest,
-    CreateNodeResponse,
-    DeleteNodeRequest,
-    DeleteNodeResponse,
+    ActivateNodeRequest,
+    ActivateNodeResponse,
+    DeactivateNodeRequest,
+    DeactivateNodeResponse,
     PullMessagesRequest,
     PullMessagesResponse,
     PushMessagesRequest,
     PushMessagesResponse,
+    RegisterNodeFleetRequest,
+    RegisterNodeFleetResponse,
+    UnregisterNodeFleetRequest,
+    UnregisterNodeFleetResponse,
 )
 from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
     SendNodeHeartbeatRequest,
@@ -63,49 +69,137 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
         state_factory: LinkStateFactory,
         ffs_factory: FfsFactory,
         objectstore_factory: ObjectStoreFactory,
+        enable_supernode_auth: bool,
     ) -> None:
         self.state_factory = state_factory
         self.ffs_factory = ffs_factory
         self.objectstore_factory = objectstore_factory
+        self.enable_supernode_auth = enable_supernode_auth
+        self.lock = threading.Lock()
 
-    def CreateNode(
-        self, request: CreateNodeRequest, context: grpc.ServicerContext
-    ) -> CreateNodeResponse:
-        """."""
-        log(
-            INFO,
-            "[Fleet.CreateNode] Request heartbeat_interval=%s",
-            request.heartbeat_interval,
-        )
-        log(DEBUG, "[Fleet.CreateNode] Request: %s", MessageToDict(request))
-        response = message_handler.create_node(
-            request=request,
-            state=self.state_factory.state(),
-        )
-        log(INFO, "[Fleet.CreateNode] Created node_id=%s", response.node.node_id)
-        log(DEBUG, "[Fleet.CreateNode] Response: %s", MessageToDict(response))
-        return response
+    def RegisterNode(
+        self, request: RegisterNodeFleetRequest, context: grpc.ServicerContext
+    ) -> RegisterNodeFleetResponse:
+        """Register a node."""
+        # Prevent registration when SuperNode authentication is enabled
+        if self.enable_supernode_auth:
+            log(ERROR, "SuperNode registration is disabled through Fleet API.")
+            context.abort(
+                grpc.StatusCode.FAILED_PRECONDITION,
+                "SuperNode authentication is enabled. "
+                "All SuperNodes must be registered via the CLI.",
+            )
 
-    def DeleteNode(
-        self, request: DeleteNodeRequest, context: grpc.ServicerContext
-    ) -> DeleteNodeResponse:
-        """."""
-        log(INFO, "[Fleet.DeleteNode] Delete node_id=%s", request.node.node_id)
-        log(DEBUG, "[Fleet.DeleteNode] Request: %s", MessageToDict(request))
-        return message_handler.delete_node(
-            request=request,
-            state=self.state_factory.state(),
-        )
+        try:
+            response = message_handler.register_node(
+                request=request,
+                state=self.state_factory.state(),
+            )
+            log(DEBUG, "[Fleet.RegisterNode] Registered node_id=%s", response.node_id)
+            return response
+        except ValueError:
+            # Public key already in use
+            # This should NEVER happen due to the public keys should be automatically
+            # generated and unique for each SuperNode instance.
+            log(
+                ERROR,
+                "[Fleet.RegisterNode] Registration failed: %s",
+                PUBLIC_KEY_ALREADY_IN_USE_MESSAGE,
+            )
+            context.abort(
+                grpc.StatusCode.FAILED_PRECONDITION, PUBLIC_KEY_ALREADY_IN_USE_MESSAGE
+            )
+
+        raise RuntimeError  # Make mypy happy
+
+    def ActivateNode(
+        self, request: ActivateNodeRequest, context: grpc.ServicerContext
+    ) -> ActivateNodeResponse:
+        """Activate a node."""
+        try:
+            response = message_handler.activate_node(
+                request=request,
+                state=self.state_factory.state(),
+            )
+            log(INFO, "[Fleet.ActivateNode] Activated node_id=%s", response.node_id)
+            return response
+        except message_handler.InvalidHeartbeatIntervalError:
+            # Heartbeat interval is invalid
+            log(ERROR, "[Fleet.ActivateNode] Invalid heartbeat interval")
+            context.abort(
+                grpc.StatusCode.INVALID_ARGUMENT, "Invalid heartbeat interval"
+            )
+        except ValueError as e:
+            log(ERROR, "[Fleet.ActivateNode] Activation failed: %s", str(e))
+            context.abort(grpc.StatusCode.FAILED_PRECONDITION, str(e))
+
+        raise RuntimeError  # Make mypy happy
+
+    def DeactivateNode(
+        self, request: DeactivateNodeRequest, context: grpc.ServicerContext
+    ) -> DeactivateNodeResponse:
+        """Deactivate a node."""
+        try:
+            response = message_handler.deactivate_node(
+                request=request,
+                state=self.state_factory.state(),
+            )
+            log(INFO, "[Fleet.DeactivateNode] Deactivated node_id=%s", request.node_id)
+            return response
+        except ValueError as e:
+            log(ERROR, "[Fleet.DeactivateNode] Deactivation failed: %s", str(e))
+            context.abort(grpc.StatusCode.FAILED_PRECONDITION, str(e))
+
+        raise RuntimeError  # Make mypy happy
+
+    def UnregisterNode(
+        self, request: UnregisterNodeFleetRequest, context: grpc.ServicerContext
+    ) -> UnregisterNodeFleetResponse:
+        """Unregister a node."""
+        # Prevent unregistration when SuperNode authentication is enabled
+        if self.enable_supernode_auth:
+            log(ERROR, "SuperNode unregistration is disabled through Fleet API.")
+            context.abort(
+                grpc.StatusCode.FAILED_PRECONDITION,
+                "SuperNode authentication is enabled. "
+                "All SuperNodes must be unregistered via the CLI.",
+            )
+
+        try:
+            response = message_handler.unregister_node(
+                request=request,
+                state=self.state_factory.state(),
+            )
+            log(
+                DEBUG, "[Fleet.UnregisterNode] Unregistered node_id=%s", request.node_id
+            )
+            return response
+        except ValueError as e:
+            log(
+                ERROR,
+                "[Fleet.UnregisterNode] Unregistration failed: %s",
+                str(e),
+            )
+            context.abort(grpc.StatusCode.FAILED_PRECONDITION, str(e))
+            raise RuntimeError from None  # Make mypy happy
 
     def SendNodeHeartbeat(
         self, request: SendNodeHeartbeatRequest, context: grpc.ServicerContext
     ) -> SendNodeHeartbeatResponse:
         """."""
         log(DEBUG, "[Fleet.SendNodeHeartbeat] Request: %s", MessageToDict(request))
-        return message_handler.send_node_heartbeat(
-            request=request,
-            state=self.state_factory.state(),
-        )
+        try:
+            return message_handler.send_node_heartbeat(
+                request=request,
+                state=self.state_factory.state(),
+            )
+        except message_handler.InvalidHeartbeatIntervalError:
+            # Heartbeat interval is invalid
+            log(ERROR, "[Fleet.SendNodeHeartbeat] Invalid heartbeat interval")
+            context.abort(
+                grpc.StatusCode.INVALID_ARGUMENT, "Invalid heartbeat interval"
+            )
+        raise RuntimeError  # Make mypy happy
 
     def PullMessages(
         self, request: PullMessagesRequest, context: grpc.ServicerContext
@@ -155,8 +249,8 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
                 state=self.state_factory.state(),
                 store=self.objectstore_factory.store(),
             )
-        except InvalidRunStatusException as e:
-            abort_grpc_context(e.message, context)
+        except (InvalidRunStatusException, ValueError) as e:
+            abort_grpc_context(str(e), context)
 
         return res
 
@@ -172,8 +266,8 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
                 state=self.state_factory.state(),
                 store=self.objectstore_factory.store(),
             )
-        except InvalidRunStatusException as e:
-            abort_grpc_context(e.message, context)
+        except (InvalidRunStatusException, ValueError) as e:
+            abort_grpc_context(str(e), context)
 
         return res
 
@@ -183,7 +277,7 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
         """Push an object to the ObjectStore."""
         log(
             DEBUG,
-            "[ServerAppIoServicer.PushObject] Push Object with object_id=%s",
+            "[Fleet.PushObject] Push Object with object_id=%s",
             request.object_id,
         )
 
@@ -208,7 +302,7 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
         """Pull an object from the ObjectStore."""
         log(
             DEBUG,
-            "[ServerAppIoServicer.PullObject] Pull Object with object_id=%s",
+            "[Fleet.PullObject] Pull Object with object_id=%s",
             request.object_id,
         )
 

flwr/server/superlink/fleet/grpc_rere/{server_interceptor.py → node_auth_server_interceptor.py}

@@ -16,7 +16,8 @@
 
 
 import datetime
-from typing import Any, Callable, Optional, cast
+from collections.abc import Callable
+from typing import Any, cast
 
 import grpc
 from google.protobuf.message import Message as GrpcMessage
@@ -29,15 +30,12 @@ from flwr.common.constant import (
     TIMESTAMP_HEADER,
     TIMESTAMP_TOLERANCE,
 )
-from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
-    bytes_to_public_key,
-    verify_signature,
-)
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
-    CreateNodeRequest,
-    CreateNodeResponse,
+    ActivateNodeRequest,
+    RegisterNodeFleetRequest,
 )
 from flwr.server.superlink.linkstate import LinkStateFactory
+from flwr.supercore.primitives.asymmetric import bytes_to_public_key, verify_signature
 
 MIN_TIMESTAMP_DIFF = -SYSTEM_TIME_TOLERANCE
 MAX_TIMESTAMP_DIFF = TIMESTAMP_TOLERANCE + SYSTEM_TIME_TOLERANCE
@@ -53,22 +51,17 @@ def _unary_unary_rpc_terminator(
     return grpc.unary_unary_rpc_method_handler(terminate)
 
 
-class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
+class NodeAuthServerInterceptor(grpc.ServerInterceptor):  # type: ignore
     """Server interceptor for node authentication.
 
     Parameters
     ----------
     state_factory : LinkStateFactory
         A factory for creating new instances of LinkState.
-    auto_auth : bool (default: False)
-        If True, nodes are authenticated without requiring their public keys to be
-        pre-stored in the LinkState. If False, only nodes with pre-stored public keys
-        can be authenticated.
     """
 
-    def __init__(self, state_factory: LinkStateFactory, auto_auth: bool = False):
+    def __init__(self, state_factory: LinkStateFactory):
         self.state_factory = state_factory
-        self.auto_auth = auto_auth
 
     def intercept_service(  # pylint: disable=too-many-return-statements
         self,
@@ -85,7 +78,6 @@ class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
         if not handler_call_details.method.startswith("/flwr.proto.Fleet/"):
             return continuation(handler_call_details)
 
-        state = self.state_factory.state()
         metadata_dict = dict(handler_call_details.invocation_metadata)
 
         # Retrieve info from the metadata
@@ -96,11 +88,6 @@ class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
         except KeyError:
             return _unary_unary_rpc_terminator("Missing authentication metadata")
 
-        if not self.auto_auth:
-            # Abort the RPC call if the node public key is not found
-            if node_pk_bytes not in state.get_node_public_keys():
-                return _unary_unary_rpc_terminator("Public key not recognized")
-
         # Verify the signature
         node_pk = bytes_to_public_key(node_pk_bytes)
         if not verify_signature(node_pk, timestamp_iso.encode("ascii"), signature):
@@ -113,50 +100,40 @@ class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
         if not MIN_TIMESTAMP_DIFF < time_diff.total_seconds() < MAX_TIMESTAMP_DIFF:
             return _unary_unary_rpc_terminator("Invalid timestamp")
 
-        # Continue the RPC call
-        expected_node_id = state.get_node_id(node_pk_bytes)
-        if not handler_call_details.method.endswith("CreateNode"):
-            # All calls, except for `CreateNode`, must provide a public key that is
-            # already mapped to a `node_id` (in `LinkState`)
-            if expected_node_id is None:
-                return _unary_unary_rpc_terminator("Invalid node ID")
-        # One of the method handlers in
+        # Continue the RPC call: One of the method handlers in
         # `flwr.server.superlink.fleet.grpc_rere.fleet_server.FleetServicer`
         method_handler: grpc.RpcMethodHandler = continuation(handler_call_details)
-        return self._wrap_method_handler(
-            method_handler, expected_node_id, node_pk_bytes
-        )
+        return self._wrap_method_handler(method_handler, node_pk_bytes)
 
     def _wrap_method_handler(
         self,
         method_handler: grpc.RpcMethodHandler,
-        expected_node_id: Optional[int],
-        node_public_key: bytes,
+        expected_public_key: bytes,
     ) -> grpc.RpcMethodHandler:
         def _generic_method_handler(
             request: GrpcMessage,
             context: grpc.ServicerContext,
         ) -> GrpcMessage:
-            # Verify the node ID
-            if not isinstance(request, CreateNodeRequest):
-                try:
-                    if request.node.node_id != expected_node_id:  # type: ignore
-                        raise ValueError
-                except (AttributeError, ValueError):
-                    context.abort(grpc.StatusCode.UNAUTHENTICATED, "Invalid node ID")
+            # Note: This function runs in a different thread
+            # than the `intercept_service` function.
+
+            # Retrieve the public key
+            if isinstance(request, (RegisterNodeFleetRequest | ActivateNodeRequest)):
+                actual_public_key = request.public_key
+            else:
+                if hasattr(request, "node"):
+                    node_id = request.node.node_id
+                else:
+                    node_id = request.node_id  # type: ignore[attr-defined]
+                actual_public_key = self.state_factory.state().get_node_public_key(
+                    node_id
+                )
+
+            # Verify the public key
+            if actual_public_key != expected_public_key:
+                context.abort(grpc.StatusCode.UNAUTHENTICATED, "Invalid node ID")
 
             response: GrpcMessage = method_handler.unary_unary(request, context)
-
-            # Set the public key after a successful CreateNode request
-            if isinstance(response, CreateNodeResponse):
-                state = self.state_factory.state()
-                try:
-                    state.set_node_public_key(response.node.node_id, node_public_key)
-                except ValueError as e:
-                    # Remove newly created node if setting the public key fails
-                    state.delete_node(response.node.node_id)
-                    context.abort(grpc.StatusCode.UNAUTHENTICATED, str(e))
-
             return response
 
         return grpc.unary_unary_rpc_method_handler(

flwr/server/superlink/fleet/message_handler/message_handler.py

@@ -15,29 +15,38 @@
 """Fleet API message handlers."""
 
 from logging import ERROR
-from typing import Optional
 
 from flwr.common import Message, log
-from flwr.common.constant import Status
+from flwr.common.constant import (
+    HEARTBEAT_MAX_INTERVAL,
+    HEARTBEAT_MIN_INTERVAL,
+    NOOP_ACCOUNT_NAME,
+    NOOP_FLWR_AID,
+    Status,
+)
 from flwr.common.inflatable import UnexpectedObjectContentError
 from flwr.common.serde import (
     fab_to_proto,
     message_from_proto,
     message_to_proto,
-    user_config_to_proto,
+    run_to_proto,
 )
-from flwr.common.typing import Fab, InvalidRunStatusException
+from flwr.common.typing import Fab, InvalidRunStatusException, Run
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
-    CreateNodeRequest,
-    CreateNodeResponse,
-    DeleteNodeRequest,
-    DeleteNodeResponse,
+    ActivateNodeRequest,
+    ActivateNodeResponse,
+    DeactivateNodeRequest,
+    DeactivateNodeResponse,
     PullMessagesRequest,
     PullMessagesResponse,
     PushMessagesRequest,
     PushMessagesResponse,
     Reconnect,
+    RegisterNodeFleetRequest,
+    RegisterNodeFleetResponse,
+    UnregisterNodeFleetRequest,
+    UnregisterNodeFleetResponse,
 )
 from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
     SendNodeHeartbeatRequest,
@@ -51,38 +60,59 @@ from flwr.proto.message_pb2 import (  # pylint: disable=E0611
     PushObjectRequest,
     PushObjectResponse,
 )
-from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
-from flwr.proto.run_pb2 import (  # pylint: disable=E0611
-    GetRunRequest,
-    GetRunResponse,
-    Run,
-)
+from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 from flwr.server.superlink.linkstate import LinkState
 from flwr.server.superlink.utils import check_abort
 from flwr.supercore.ffs import Ffs
 from flwr.supercore.object_store import NoObjectInStoreError, ObjectStore
-from flwr.supercore.object_store.utils import store_mapping_and_register_objects
 
 
-def create_node(
-    request: CreateNodeRequest,  # pylint: disable=unused-argument
+class InvalidHeartbeatIntervalError(Exception):
+    """Invalid heartbeat interval exception."""
+
+
+def register_node(
+    request: RegisterNodeFleetRequest,
     state: LinkState,
-) -> CreateNodeResponse:
-    """."""
-    # Create node
-    node_id = state.create_node(heartbeat_interval=request.heartbeat_interval)
-    return CreateNodeResponse(node=Node(node_id=node_id))
+) -> RegisterNodeFleetResponse:
+    """Register a node (Fleet API only)."""
+    node_id = state.create_node(NOOP_FLWR_AID, NOOP_ACCOUNT_NAME, request.public_key, 0)
+    return RegisterNodeFleetResponse(node_id=node_id)
 
 
-def delete_node(request: DeleteNodeRequest, state: LinkState) -> DeleteNodeResponse:
-    """."""
-    # Validate node_id
-    if request.node.node_id == 0:  # i.e. unset `node_id`
-        return DeleteNodeResponse()
+def activate_node(
+    request: ActivateNodeRequest,
+    state: LinkState,
+) -> ActivateNodeResponse:
+    """Activate a node."""
+    node_id = state.get_node_id_by_public_key(request.public_key)
+    if node_id is None:
+        raise ValueError("No SuperNode found with the given public key.")
+    _validate_heartbeat_interval(request.heartbeat_interval)
+    if not state.activate_node(node_id, request.heartbeat_interval):
+        raise ValueError(f"SuperNode with node ID {node_id} could not be activated.")
+    return ActivateNodeResponse(node_id=node_id)
+
+
+def deactivate_node(
+    request: DeactivateNodeRequest,
+    state: LinkState,
+) -> DeactivateNodeResponse:
+    """Deactivate a node."""
+    if not state.deactivate_node(request.node_id):
+        raise ValueError(
+            f"SuperNode with node ID {request.node_id} could not be deactivated."
+        )
+    return DeactivateNodeResponse()
 
-    # Update state
-    state.delete_node(node_id=request.node.node_id)
-    return DeleteNodeResponse()
+
+def unregister_node(
+    request: UnregisterNodeFleetRequest,
+    state: LinkState,
+) -> UnregisterNodeFleetResponse:
+    """Unregister a node (Fleet API only)."""
+    state.delete_node(NOOP_FLWR_AID, request.node_id)
+    return UnregisterNodeFleetResponse()
 
 
 def send_node_heartbeat(
@@ -90,6 +120,7 @@ def send_node_heartbeat(
     state: LinkState,  # pylint: disable=unused-argument
 ) -> SendNodeHeartbeatResponse:
     """."""
+    _validate_heartbeat_interval(request.heartbeat_interval)
     res = state.acknowledge_node_heartbeat(
         request.node.node_id, request.heartbeat_interval
     )
@@ -137,10 +168,11 @@ def push_messages(
     """Push Messages handler."""
     # Convert Message from proto
     msg = message_from_proto(message_proto=request.messages_list[0])
+    run_id = msg.metadata.run_id
 
     # Abort if the run is not running
     abort_msg = check_abort(
-        msg.metadata.run_id,
+        run_id,
         [Status.PENDING, Status.STARTING, Status.FINISHED],
         state,
         store,
@@ -148,11 +180,12 @@ def push_messages(
     if abort_msg:
         raise InvalidRunStatusException(abort_msg)
 
-    # Store Message in State
-    message_id: Optional[str] = state.store_message_res(message=msg)
-
     # Store Message object to descendants mapping and preregister objects
-    objects_to_push = store_mapping_and_register_objects(store, request=request)
+    objects_to_push: set[str] = set()
+    for object_tree in request.message_object_trees:
+        objects_to_push |= set(store.preregister(run_id, object_tree))
+    # Store Message in State
+    message_id: str | None = state.store_message_res(message=msg)
 
     # Build response
     response = PushMessagesResponse(
@@ -167,10 +200,8 @@ def get_run(
     request: GetRunRequest, state: LinkState, store: ObjectStore
 ) -> GetRunResponse:
     """Get run information."""
-    run = state.get_run(request.run_id)
-
-    if run is None:
-        return GetRunResponse()
+    # Validate that the requesting SuperNode is part of the federation
+    run = _validate_node_in_federation(state, request.node.node_id, request.run_id)
 
     # Abort if the run is not running
     abort_msg = check_abort(
@@ -182,21 +213,16 @@ def get_run(
     if abort_msg:
         raise InvalidRunStatusException(abort_msg)
 
-    return GetRunResponse(
-        run=Run(
-            run_id=run.run_id,
-            fab_id=run.fab_id,
-            fab_version=run.fab_version,
-            override_config=user_config_to_proto(run.override_config),
-            fab_hash=run.fab_hash,
-        )
-    )
+    return GetRunResponse(run=run_to_proto(run))
 
 
 def get_fab(
     request: GetFabRequest, ffs: Ffs, state: LinkState, store: ObjectStore
 ) -> GetFabResponse:
     """Get FAB."""
+    # Validate that the requesting SuperNode is part of the federation
+    _validate_node_in_federation(state, request.node.node_id, request.run_id)
+
     # Abort if the run is not running
     abort_msg = check_abort(
         request.run_id,
@@ -208,7 +234,7 @@ def get_fab(
         raise InvalidRunStatusException(abort_msg)
 
     if result := ffs.get(request.hash_str):
-        fab = Fab(request.hash_str, result[0])
+        fab = Fab(request.hash_str, result[0], result[1])
         return GetFabResponse(fab=fab_to_proto(fab))
 
     raise ValueError(f"Found no FAB with hash: {request.hash_str}")
@@ -284,3 +310,28 @@ def confirm_message_received(
     store.delete(request.message_object_id)
 
     return ConfirmMessageReceivedResponse()
+
+
+def _validate_heartbeat_interval(interval: float) -> None:
+    """Raise if heartbeat interval is out of bounds."""
+    if not HEARTBEAT_MIN_INTERVAL <= interval <= HEARTBEAT_MAX_INTERVAL:
+        raise InvalidHeartbeatIntervalError(
+            f"Heartbeat interval {interval} is out of bounds "
+            f"[{HEARTBEAT_MIN_INTERVAL}, {HEARTBEAT_MAX_INTERVAL}]."
+        )
+
+
+def _validate_node_in_federation(
+    state: LinkState,
+    node_id: int,
+    run_id: int,
+) -> Run:
+    """Raise if the requesting SuperNode is not part of the federation the run belongs
+    to."""
+    run = state.get_run(run_id)
+    if not run:
+        raise ValueError(f"Run ID not found: {run_id}")
+
+    if not state.federation_manager.has_node(node_id, run.federation):
+        raise ValueError(f"SuperNode is not part of the federation '{run.federation}'.")
+    return run