flwr-nightly 1.23.0.dev20250930__py3-none-any.whl → 1.26.0.dev20260121__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,41 @@ 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,
+        received_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):
+            # Note: This function runs in a different thread
+            # than the `intercept_service` function.
+
+            # Skip registration and activation requests
+            if not isinstance(request, (RegisterNodeFleetRequest, ActivateNodeRequest)):
+                # Retrieve the node ID from the request
+                if hasattr(request, "node"):
+                    received_node_id = request.node.node_id
+                else:
+                    received_node_id = request.node_id  # type: ignore[attr-defined]
+
+                # Get the actual node ID based on the received public key
+                node_id = self.state_factory.state().get_node_id_by_public_key(
+                    received_public_key
+                )
+
+                # Verify that the received node ID matches the actual node ID
+                if received_node_id != node_id:
                     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(