flwr-nightly 1.15.0.dev20250104__py3-none-any.whl → 1.15.0.dev20250123__py3-none-any.whl

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

@@ -30,8 +30,12 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     DeleteNodeResponse,
     PingRequest,
     PingResponse,
+    PullMessagesRequest,
+    PullMessagesResponse,
     PullTaskInsRequest,
     PullTaskInsResponse,
+    PushMessagesRequest,
+    PushMessagesResponse,
     PushTaskResRequest,
     PushTaskResResponse,
 )
@@ -95,6 +99,17 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
             state=self.state_factory.state(),
         )
 
+    def PullMessages(
+        self, request: PullMessagesRequest, context: grpc.ServicerContext
+    ) -> PullMessagesResponse:
+        """Pull Messages."""
+        log(INFO, "[Fleet.PullMessages] node_id=%s", request.node.node_id)
+        log(DEBUG, "[Fleet.PullMessages] Request: %s", request)
+        return message_handler.pull_messages(
+            request=request,
+            state=self.state_factory.state(),
+        )
+
     def PushTaskRes(
         self, request: PushTaskResRequest, context: grpc.ServicerContext
     ) -> PushTaskResResponse:
@@ -118,6 +133,29 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
 
         return res
 
+    def PushMessages(
+        self, request: PushMessagesRequest, context: grpc.ServicerContext
+    ) -> PushMessagesResponse:
+        """Push Messages."""
+        if request.messages_list:
+            log(
+                INFO,
+                "[Fleet.PushMessages] Push results from node_id=%s",
+                request.messages_list[0].metadata.src_node_id,
+            )
+        else:
+            log(INFO, "[Fleet.PushMessages] No task results to push")
+
+        try:
+            res = message_handler.push_messages(
+                request=request,
+                state=self.state_factory.state(),
+            )
+        except InvalidRunStatusException as e:
+            abort_grpc_context(e.message, context)
+
+        return res
+
     def GetRun(
         self, request: GetRunRequest, context: grpc.ServicerContext
     ) -> GetRunResponse:

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

@@ -15,91 +15,54 @@
 """Flower server interceptor."""
 
 
-import base64
-from collections.abc import Sequence
-from logging import INFO, WARNING
-from typing import Any, Callable, Optional, Union
+import datetime
+from typing import Any, Callable, Optional, cast
 
 import grpc
-from cryptography.hazmat.primitives.asymmetric import ec
-
-from flwr.common.logger import log
+from google.protobuf.message import Message as GrpcMessage
+
+from flwr.common import now
+from flwr.common.constant import (
+    PUBLIC_KEY_HEADER,
+    SIGNATURE_HEADER,
+    TIMESTAMP_HEADER,
+    TIMESTAMP_TOLERANCE,
+)
 from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
-    bytes_to_private_key,
     bytes_to_public_key,
-    generate_shared_key,
-    verify_hmac,
+    verify_signature,
 )
-from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     CreateNodeRequest,
     CreateNodeResponse,
-    DeleteNodeRequest,
-    DeleteNodeResponse,
-    PingRequest,
-    PingResponse,
-    PullTaskInsRequest,
-    PullTaskInsResponse,
-    PushTaskResRequest,
-    PushTaskResResponse,
 )
-from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
-from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 from flwr.server.superlink.linkstate import LinkStateFactory
 
-_PUBLIC_KEY_HEADER = "public-key"
-_AUTH_TOKEN_HEADER = "auth-token"
-
-Request = Union[
-    CreateNodeRequest,
-    DeleteNodeRequest,
-    PullTaskInsRequest,
-    PushTaskResRequest,
-    GetRunRequest,
-    PingRequest,
-    GetFabRequest,
-]
-
-Response = Union[
-    CreateNodeResponse,
-    DeleteNodeResponse,
-    PullTaskInsResponse,
-    PushTaskResResponse,
-    GetRunResponse,
-    PingResponse,
-    GetFabResponse,
-]
-
 
-def _get_value_from_tuples(
-    key_string: str, tuples: Sequence[tuple[str, Union[str, bytes]]]
-) -> bytes:
-    value = next((value for key, value in tuples if key == key_string), "")
-    if isinstance(value, str):
-        return value.encode()
+def _unary_unary_rpc_terminator(message: str) -> grpc.RpcMethodHandler:
+    def terminate(_request: GrpcMessage, context: grpc.ServicerContext) -> GrpcMessage:
+        context.abort(grpc.StatusCode.UNAUTHENTICATED, message)
+        raise RuntimeError("Should not reach this point")  # Make mypy happy
 
-    return value
+    return grpc.unary_unary_rpc_method_handler(terminate)
 
 
 class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
-    """Server interceptor for node authentication."""
-
-    def __init__(self, state_factory: LinkStateFactory):
+    """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):
         self.state_factory = state_factory
-        state = self.state_factory.state()
-
-        self.node_public_keys = state.get_node_public_keys()
-        if len(self.node_public_keys) == 0:
-            log(WARNING, "Authentication enabled, but no known public keys configured")
-
-        private_key = state.get_server_private_key()
-        public_key = state.get_server_public_key()
-
-        if private_key is None or public_key is None:
-            raise ValueError("Error loading authentication keys")
-
-        self.server_private_key = bytes_to_private_key(private_key)
-        self.encoded_server_public_key = base64.urlsafe_b64encode(public_key)
+        self.auto_auth = auto_auth
 
     def intercept_service(
         self,
@@ -112,116 +75,80 @@ class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
         metadata sent by the node. Continue RPC call if node is authenticated, else,
         terminate RPC call by setting context to abort.
         """
+        state = self.state_factory.state()
+        metadata_dict = dict(handler_call_details.invocation_metadata)
+
+        # Retrieve info from the metadata
+        try:
+            node_pk_bytes = cast(bytes, metadata_dict[PUBLIC_KEY_HEADER])
+            timestamp_iso = cast(str, metadata_dict[TIMESTAMP_HEADER])
+            signature = cast(bytes, metadata_dict[SIGNATURE_HEADER])
+        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):
+            return _unary_unary_rpc_terminator("Invalid signature")
+
+        # Verify the timestamp
+        current = now()
+        time_diff = current - datetime.datetime.fromisoformat(timestamp_iso)
+        # Abort the RPC call if the timestamp is too old or in the future
+        if not 0 < time_diff.total_seconds() < TIMESTAMP_TOLERANCE:
+            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"):
+            if expected_node_id is None:
+                return _unary_unary_rpc_terminator("Invalid node ID")
         # 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._generic_auth_unary_method_handler(method_handler)
+        return self._wrap_method_handler(
+            method_handler, expected_node_id, node_pk_bytes
+        )
 
-    def _generic_auth_unary_method_handler(
-        self, method_handler: grpc.RpcMethodHandler
+    def _wrap_method_handler(
+        self,
+        method_handler: grpc.RpcMethodHandler,
+        expected_node_id: Optional[int],
+        node_public_key: bytes,
     ) -> grpc.RpcMethodHandler:
         def _generic_method_handler(
-            request: Request,
+            request: GrpcMessage,
             context: grpc.ServicerContext,
-        ) -> Response:
-            node_public_key_bytes = base64.urlsafe_b64decode(
-                _get_value_from_tuples(
-                    _PUBLIC_KEY_HEADER, context.invocation_metadata()
-                )
-            )
-            if node_public_key_bytes not in self.node_public_keys:
-                context.abort(grpc.StatusCode.UNAUTHENTICATED, "Access denied")
-
-            if isinstance(request, CreateNodeRequest):
-                response = self._create_authenticated_node(
-                    node_public_key_bytes, request, context
-                )
-                log(
-                    INFO,
-                    "AuthenticateServerInterceptor: Created node_id=%s",
-                    response.node.node_id,
-                )
-                return response
-
-            # Verify hmac value
-            hmac_value = base64.urlsafe_b64decode(
-                _get_value_from_tuples(
-                    _AUTH_TOKEN_HEADER, context.invocation_metadata()
-                )
-            )
-            public_key = bytes_to_public_key(node_public_key_bytes)
-
-            if not self._verify_hmac(public_key, request, hmac_value):
-                context.abort(grpc.StatusCode.UNAUTHENTICATED, "Access denied")
-
-            # Verify node_id
-            node_id = self.state_factory.state().get_node_id(node_public_key_bytes)
-
-            if not self._verify_node_id(node_id, request):
-                context.abort(grpc.StatusCode.UNAUTHENTICATED, "Access denied")
-
-            return method_handler.unary_unary(request, context)  # type: ignore
+        ) -> 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")
+
+            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(
             _generic_method_handler,
             request_deserializer=method_handler.request_deserializer,
             response_serializer=method_handler.response_serializer,
         )
-
-    def _verify_node_id(
-        self,
-        node_id: Optional[int],
-        request: Union[
-            DeleteNodeRequest,
-            PullTaskInsRequest,
-            PushTaskResRequest,
-            GetRunRequest,
-            PingRequest,
-            GetFabRequest,
-        ],
-    ) -> bool:
-        if node_id is None:
-            return False
-        if isinstance(request, PushTaskResRequest):
-            if len(request.task_res_list) == 0:
-                return False
-            return request.task_res_list[0].task.producer.node_id == node_id
-        if isinstance(request, GetRunRequest):
-            return node_id in self.state_factory.state().get_nodes(request.run_id)
-        return request.node.node_id == node_id
-
-    def _verify_hmac(
-        self, public_key: ec.EllipticCurvePublicKey, request: Request, hmac_value: bytes
-    ) -> bool:
-        shared_secret = generate_shared_key(self.server_private_key, public_key)
-        message_bytes = request.SerializeToString(deterministic=True)
-        return verify_hmac(shared_secret, message_bytes, hmac_value)
-
-    def _create_authenticated_node(
-        self,
-        public_key_bytes: bytes,
-        request: CreateNodeRequest,
-        context: grpc.ServicerContext,
-    ) -> CreateNodeResponse:
-        context.send_initial_metadata(
-            (
-                (
-                    _PUBLIC_KEY_HEADER,
-                    self.encoded_server_public_key,
-                ),
-            )
-        )
-        state = self.state_factory.state()
-        node_id = state.get_node_id(public_key_bytes)
-
-        # Handle `CreateNode` here instead of calling the default method handler
-        # Return previously assigned `node_id` for the provided `public_key`
-        if node_id is not None:
-            state.acknowledge_ping(node_id, request.ping_interval)
-            return CreateNodeResponse(node=Node(node_id=node_id, anonymous=False))
-
-        # No `node_id` exists for the provided `public_key`
-        # Handle `CreateNode` here instead of calling the default method handler
-        # Note: the innermost `CreateNode` method will never be called
-        node_id = state.create_node(request.ping_interval, public_key_bytes)
-        return CreateNodeResponse(node=Node(node_id=node_id, anonymous=False))

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

@@ -20,7 +20,14 @@ from typing import Optional
 from uuid import UUID
 
 from flwr.common.constant import Status
-from flwr.common.serde import fab_to_proto, user_config_to_proto
+from flwr.common.serde import (
+    fab_to_proto,
+    message_from_proto,
+    message_from_taskins,
+    message_to_proto,
+    message_to_taskres,
+    user_config_to_proto,
+)
 from flwr.common.typing import Fab, InvalidRunStatusException
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
@@ -30,8 +37,12 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     DeleteNodeResponse,
     PingRequest,
     PingResponse,
+    PullMessagesRequest,
+    PullMessagesResponse,
     PullTaskInsRequest,
     PullTaskInsResponse,
+    PushMessagesRequest,
+    PushMessagesResponse,
     PushTaskResRequest,
     PushTaskResResponse,
     Reconnect,
@@ -55,13 +66,13 @@ def create_node(
     """."""
     # Create node
     node_id = state.create_node(ping_interval=request.ping_interval)
-    return CreateNodeResponse(node=Node(node_id=node_id, anonymous=False))
+    return CreateNodeResponse(node=Node(node_id=node_id))
 
 
 def delete_node(request: DeleteNodeRequest, state: LinkState) -> DeleteNodeResponse:
     """."""
     # Validate node_id
-    if request.node.anonymous or request.node.node_id == 0:
+    if request.node.node_id == 0:  # i.e. unset `node_id`
         return DeleteNodeResponse()
 
     # Update state
@@ -80,9 +91,8 @@ def ping(
 
 def pull_task_ins(request: PullTaskInsRequest, state: LinkState) -> PullTaskInsResponse:
     """Pull TaskIns handler."""
-    # Get node_id if client node is not anonymous
     node = request.node  # pylint: disable=no-member
-    node_id: Optional[int] = None if node.anonymous else node.node_id
+    node_id: int = node.node_id
 
     # Retrieve TaskIns from State
     task_ins_list: list[TaskIns] = state.get_task_ins(node_id=node_id, limit=1)
@@ -94,6 +104,26 @@ def pull_task_ins(request: PullTaskInsRequest, state: LinkState) -> PullTaskInsR
     return response
 
 
+def pull_messages(
+    request: PullMessagesRequest, state: LinkState
+) -> PullMessagesResponse:
+    """Pull Messages handler."""
+    # Get node_id if client node is not anonymous
+    node = request.node  # pylint: disable=no-member
+    node_id: int = node.node_id
+
+    # Retrieve TaskIns from State
+    task_ins_list: list[TaskIns] = state.get_task_ins(node_id=node_id, limit=1)
+
+    # Convert to Messages
+    msg_proto = []
+    for task_ins in task_ins_list:
+        msg = message_from_taskins(task_ins)
+        msg_proto.append(message_to_proto(msg))
+
+    return PullMessagesResponse(messages_list=msg_proto)
+
+
 def push_task_res(request: PushTaskResRequest, state: LinkState) -> PushTaskResResponse:
     """Push TaskRes handler."""
     # pylint: disable=no-member
@@ -123,6 +153,37 @@ def push_task_res(request: PushTaskResRequest, state: LinkState) -> PushTaskResR
     return response
 
 
+def push_messages(
+    request: PushMessagesRequest, state: LinkState
+) -> PushMessagesResponse:
+    """Push Messages handler."""
+    # Convert Message to TaskRes
+    msg = message_from_proto(message_proto=request.messages_list[0])
+    task_res = message_to_taskres(msg)
+
+    # Abort if the run is not running
+    abort_msg = check_abort(
+        task_res.run_id,
+        [Status.PENDING, Status.STARTING, Status.FINISHED],
+        state,
+    )
+    if abort_msg:
+        raise InvalidRunStatusException(abort_msg)
+
+    # Set pushed_at (timestamp in seconds)
+    task_res.task.pushed_at = time.time()
+
+    # Store TaskRes in State
+    message_id: Optional[UUID] = state.store_task_res(task_res=task_res)
+
+    # Build response
+    response = PushMessagesResponse(
+        reconnect=Reconnect(reconnect=5),
+        results={str(message_id): 0},
+    )
+    return response
+
+
 def get_run(request: GetRunRequest, state: LinkState) -> GetRunResponse:
     """Get run information."""
     run = state.get_run(request.run_id)

flwr/server/superlink/fleet/rest_rere/rest_api.py

@@ -17,13 +17,12 @@
 
 from __future__ import annotations
 
-import sys
 from collections.abc import Awaitable
 from typing import Callable, TypeVar, cast
 
 from google.protobuf.message import Message as GrpcMessage
 
-from flwr.common.constant import MISSING_EXTRA_REST
+from flwr.common.exit import ExitCode, flwr_exit
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     CreateNodeRequest,
@@ -32,8 +31,12 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     DeleteNodeResponse,
     PingRequest,
     PingResponse,
+    PullMessagesRequest,
+    PullMessagesResponse,
     PullTaskInsRequest,
     PullTaskInsResponse,
+    PushMessagesRequest,
+    PushMessagesResponse,
     PushTaskResRequest,
     PushTaskResResponse,
 )
@@ -51,7 +54,7 @@ try:
     from starlette.responses import Response
     from starlette.routing import Route
 except ModuleNotFoundError:
-    sys.exit(MISSING_EXTRA_REST)
+    flwr_exit(ExitCode.COMMON_MISSING_EXTRA_REST)
 
 
 GrpcRequest = TypeVar("GrpcRequest", bound=GrpcMessage)
@@ -117,6 +120,16 @@ async def pull_task_ins(request: PullTaskInsRequest) -> PullTaskInsResponse:
     return message_handler.pull_task_ins(request=request, state=state)
 
 
+@rest_request_response(PullMessagesRequest)
+async def pull_message(request: PullMessagesRequest) -> PullMessagesResponse:
+    """Pull PullMessages."""
+    # Get state from app
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+
+    # Handle message
+    return message_handler.pull_messages(request=request, state=state)
+
+
 # Check if token is needed here
 @rest_request_response(PushTaskResRequest)
 async def push_task_res(request: PushTaskResRequest) -> PushTaskResResponse:
@@ -128,6 +141,16 @@ async def push_task_res(request: PushTaskResRequest) -> PushTaskResResponse:
     return message_handler.push_task_res(request=request, state=state)
 
 
+@rest_request_response(PushMessagesRequest)
+async def push_message(request: PushMessagesRequest) -> PushMessagesResponse:
+    """Pull PushMessages."""
+    # Get state from app
+    state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+
+    # Handle message
+    return message_handler.push_messages(request=request, state=state)
+
+
 @rest_request_response(PingRequest)
 async def ping(request: PingRequest) -> PingResponse:
     """Ping."""
@@ -165,7 +188,9 @@ routes = [
     Route("/api/v0/fleet/create-node", create_node, methods=["POST"]),
     Route("/api/v0/fleet/delete-node", delete_node, methods=["POST"]),
     Route("/api/v0/fleet/pull-task-ins", pull_task_ins, methods=["POST"]),
+    Route("/api/v0/fleet/pull-messages", pull_message, methods=["POST"]),
     Route("/api/v0/fleet/push-task-res", push_task_res, methods=["POST"]),
+    Route("/api/v0/fleet/push-messages", push_message, methods=["POST"]),
     Route("/api/v0/fleet/ping", ping, methods=["POST"]),
     Route("/api/v0/fleet/get-run", get_run, methods=["POST"]),
     Route("/api/v0/fleet/get-fab", get_fab, methods=["POST"]),

flwr/server/superlink/fleet/vce/vce_api.py

@@ -182,8 +182,8 @@ def run_api(
     f_stop: threading.Event,
 ) -> None:
     """Run the VCE."""
-    taskins_queue: "Queue[TaskIns]" = Queue()
-    taskres_queue: "Queue[TaskRes]" = Queue()
+    taskins_queue: Queue[TaskIns] = Queue()
+    taskres_queue: Queue[TaskRes] = Queue()
 
     try: