flwr 1.14.0__py3-none-any.whl → 1.15.1__py3-none-any.whl

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

@@ -18,6 +18,7 @@
 from logging import DEBUG, INFO
 
 import grpc
+from google.protobuf.json_format import MessageToDict
 
 from flwr.common.logger import log
 from flwr.common.typing import InvalidRunStatusException
@@ -30,10 +31,10 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     DeleteNodeResponse,
     PingRequest,
     PingResponse,
-    PullTaskInsRequest,
-    PullTaskInsResponse,
-    PushTaskResRequest,
-    PushTaskResResponse,
+    PullMessagesRequest,
+    PullMessagesResponse,
+    PushMessagesRequest,
+    PushMessagesResponse,
 )
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs_factory import FfsFactory
@@ -56,13 +57,13 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
     ) -> CreateNodeResponse:
         """."""
         log(INFO, "[Fleet.CreateNode] Request ping_interval=%s", request.ping_interval)
-        log(DEBUG, "[Fleet.CreateNode] Request: %s", request)
+        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", response)
+        log(DEBUG, "[Fleet.CreateNode] Response: %s", MessageToDict(response))
         return response
 
     def DeleteNode(
@@ -70,7 +71,7 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
     ) -> DeleteNodeResponse:
         """."""
         log(INFO, "[Fleet.DeleteNode] Delete node_id=%s", request.node.node_id)
-        log(DEBUG, "[Fleet.DeleteNode] Request: %s", request)
+        log(DEBUG, "[Fleet.DeleteNode] Request: %s", MessageToDict(request))
         return message_handler.delete_node(
             request=request,
             state=self.state_factory.state(),
@@ -78,38 +79,38 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
 
     def Ping(self, request: PingRequest, context: grpc.ServicerContext) -> PingResponse:
         """."""
-        log(DEBUG, "[Fleet.Ping] Request: %s", request)
+        log(DEBUG, "[Fleet.Ping] Request: %s", MessageToDict(request))
         return message_handler.ping(
             request=request,
             state=self.state_factory.state(),
         )
 
-    def PullTaskIns(
-        self, request: PullTaskInsRequest, context: grpc.ServicerContext
-    ) -> PullTaskInsResponse:
-        """Pull TaskIns."""
-        log(INFO, "[Fleet.PullTaskIns] node_id=%s", request.node.node_id)
-        log(DEBUG, "[Fleet.PullTaskIns] Request: %s", request)
-        return message_handler.pull_task_ins(
+    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", MessageToDict(request))
+        return message_handler.pull_messages(
             request=request,
             state=self.state_factory.state(),
         )
 
-    def PushTaskRes(
-        self, request: PushTaskResRequest, context: grpc.ServicerContext
-    ) -> PushTaskResResponse:
-        """Push TaskRes."""
-        if request.task_res_list:
+    def PushMessages(
+        self, request: PushMessagesRequest, context: grpc.ServicerContext
+    ) -> PushMessagesResponse:
+        """Push Messages."""
+        if request.messages_list:
             log(
                 INFO,
-                "[Fleet.PushTaskRes] Push results from node_id=%s",
-                request.task_res_list[0].task.producer.node_id,
+                "[Fleet.PushMessages] Push results from node_id=%s",
+                request.messages_list[0].metadata.src_node_id,
             )
         else:
-            log(INFO, "[Fleet.PushTaskRes] No task results to push")
+            log(INFO, "[Fleet.PushMessages] No task results to push")
 
         try:
-            res = message_handler.push_task_res(
+            res = message_handler.push_messages(
                 request=request,
                 state=self.state_factory.state(),
             )

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

@@ -15,93 +15,62 @@
 """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,
+    SYSTEM_TIME_TOLERANCE,
+    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"
+MIN_TIMESTAMP_DIFF = -SYSTEM_TIME_TOLERANCE
+MAX_TIMESTAMP_DIFF = TIMESTAMP_TOLERANCE + SYSTEM_TIME_TOLERANCE
 
-Request = Union[
-    CreateNodeRequest,
-    DeleteNodeRequest,
-    PullTaskInsRequest,
-    PushTaskResRequest,
-    GetRunRequest,
-    PingRequest,
-    GetFabRequest,
-]
-
-Response = Union[
-    CreateNodeResponse,
-    DeleteNodeResponse,
-    PullTaskInsResponse,
-    PushTaskResResponse,
-    GetRunResponse,
-    PingResponse,
-    GetFabResponse,
-]
 
+def _unary_unary_rpc_terminator(
+    message: str, code: Any = grpc.StatusCode.UNAUTHENTICATED
+) -> grpc.RpcMethodHandler:
+    def terminate(_request: GrpcMessage, context: grpc.ServicerContext) -> GrpcMessage:
+        context.abort(code, message)
+        raise RuntimeError("Should not reach this point")  # Make mypy happy
 
-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()
-
-    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.auto_auth = auto_auth
 
-        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)
-
-    def intercept_service(
+    def intercept_service(  # pylint: disable=too-many-return-statements
         self,
         continuation: Callable[[Any], Any],
         handler_call_details: grpc.HandlerCallDetails,
@@ -112,116 +81,89 @@ 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.
         """
+        # Filter out non-Fleet service calls
+        if not handler_call_details.method.startswith("/flwr.proto.Fleet/"):
+            return _unary_unary_rpc_terminator(
+                "This request should be sent to a different service.",
+                grpc.StatusCode.FAILED_PRECONDITION,
+            )
+
+        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 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
         # `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,10 +37,10 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     DeleteNodeResponse,
     PingRequest,
     PingResponse,
-    PullTaskInsRequest,
-    PullTaskInsResponse,
-    PushTaskResRequest,
-    PushTaskResResponse,
+    PullMessagesRequest,
+    PullMessagesResponse,
+    PushMessagesRequest,
+    PushMessagesResponse,
     Reconnect,
 )
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
@@ -42,7 +49,7 @@ from flwr.proto.run_pb2 import (  # pylint: disable=E0611
     GetRunResponse,
     Run,
 )
-from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
+from flwr.proto.task_pb2 import TaskIns  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs import Ffs
 from flwr.server.superlink.linkstate import LinkState
 from flwr.server.superlink.utils import check_abort
@@ -55,13 +62,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
@@ -78,27 +85,33 @@ def ping(
     return PingResponse(success=res)
 
 
-def pull_task_ins(request: PullTaskInsRequest, state: LinkState) -> PullTaskInsResponse:
-    """Pull TaskIns handler."""
+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: 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)
 
-    # Build response
-    response = PullTaskInsResponse(
-        task_ins_list=task_ins_list,
-    )
-    return response
+    # 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
-    task_res: TaskRes = request.task_res_list[0]
-    # pylint: enable=no-member
+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(
@@ -113,12 +126,12 @@ def push_task_res(request: PushTaskResRequest, state: LinkState) -> PushTaskResR
     task_res.task.pushed_at = time.time()
 
     # Store TaskRes in State
-    task_id: Optional[UUID] = state.store_task_res(task_res=task_res)
+    message_id: Optional[UUID] = state.store_task_res(task_res=task_res)
 
     # Build response
-    response = PushTaskResResponse(
+    response = PushMessagesResponse(
         reconnect=Reconnect(reconnect=5),
-        results={str(task_id): 0},
+        results={str(message_id): 0},
     )
     return response
 

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,10 +31,10 @@ from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     DeleteNodeResponse,
     PingRequest,
     PingResponse,
-    PullTaskInsRequest,
-    PullTaskInsResponse,
-    PushTaskResRequest,
-    PushTaskResResponse,
+    PullMessagesRequest,
+    PullMessagesResponse,
+    PushMessagesRequest,
+    PushMessagesResponse,
 )
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs import Ffs
@@ -51,7 +50,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)
@@ -107,25 +106,24 @@ async def delete_node(request: DeleteNodeRequest) -> DeleteNodeResponse:
     return message_handler.delete_node(request=request, state=state)
 
 
-@rest_request_response(PullTaskInsRequest)
-async def pull_task_ins(request: PullTaskInsRequest) -> PullTaskInsResponse:
-    """Pull TaskIns."""
+@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_task_ins(request=request, state=state)
+    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:
-    """Push TaskRes."""
+@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_task_res(request=request, state=state)
+    return message_handler.push_messages(request=request, state=state)
 
 
 @rest_request_response(PingRequest)
@@ -164,8 +162,8 @@ async def get_fab(request: GetFabRequest) -> GetFabResponse:
 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/push-task-res", push_task_res, methods=["POST"]),
+    Route("/api/v0/fleet/pull-messages", pull_message, 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: