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

flwr/supernode/start_client_internal.py

@@ -15,23 +15,26 @@
 """Main loop for Flower SuperNode."""
 
 
+import hashlib
+import json
 import os
 import subprocess
 import time
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
 from contextlib import contextmanager
 from functools import partial
-from logging import INFO
+from logging import ERROR, INFO, WARN
 from pathlib import Path
-from typing import Callable, Optional, Union, cast
+from typing import cast
 
 import grpc
-from cryptography.hazmat.primitives.asymmetric import ec
+from cryptography.hazmat.primitives.asymmetric import ec, ed25519
+from cryptography.hazmat.primitives.serialization.ssh import load_ssh_public_key
 from grpc import RpcError
 
 from flwr.client.grpc_adapter_client.connection import grpc_adapter
 from flwr.client.grpc_rere_client.connection import grpc_request_response
-from flwr.common import GRPC_MAX_MESSAGE_LENGTH, Context, Message, RecordDict
+from flwr.common import GRPC_MAX_MESSAGE_LENGTH, Context, Error, Message, RecordDict
 from flwr.common.address import parse_address
 from flwr.common.config import get_flwr_dir, get_fused_config_from_fab
 from flwr.common.constant import (
@@ -41,11 +44,17 @@ from flwr.common.constant import (
     TRANSPORT_TYPE_GRPC_RERE,
     TRANSPORT_TYPE_REST,
     TRANSPORT_TYPES,
+    ErrorCode,
     ExecPluginType,
 )
 from flwr.common.exit import ExitCode, flwr_exit, register_signal_handlers
 from flwr.common.grpc import generic_create_grpc_server
-from flwr.common.inflatable import iterate_object_tree
+from flwr.common.inflatable import (
+    get_all_nested_objects,
+    get_object_tree,
+    iterate_object_tree,
+    no_object_id_recompute,
+)
 from flwr.common.inflatable_utils import (
     pull_objects,
     push_object_contents_from_iterable,
@@ -60,11 +69,18 @@ from flwr.proto.message_pb2 import ObjectTree  # pylint: disable=E0611
 from flwr.supercore.ffs import Ffs, FfsFactory
 from flwr.supercore.grpc_health import run_health_server_grpc_no_tls
 from flwr.supercore.object_store import ObjectStore, ObjectStoreFactory
+from flwr.supercore.primitives.asymmetric_ed25519 import (
+    create_message_to_sign,
+    decode_base64url,
+    verify_signature,
+)
 from flwr.supernode.nodestate import NodeState, NodeStateFactory
 from flwr.supernode.servicer.clientappio import ClientAppIoServicer
 
 DEFAULT_FFS_DIR = get_flwr_dir() / "supernode" / "ffs"
 
+FAB_VERIFICATION_ERROR = Error(ErrorCode.INVALID_FAB, "The FAB could not be verified.")
+
 
 # pylint: disable=import-outside-toplevel
 # pylint: disable=too-many-branches
@@ -75,18 +91,19 @@ def start_client_internal(
     *,
     server_address: str,
     node_config: UserConfig,
-    root_certificates: Optional[Union[bytes, str]] = None,
-    insecure: Optional[bool] = None,
+    root_certificates: bytes | str | None = None,
+    insecure: bool | None = None,
     transport: str,
-    authentication_keys: Optional[
-        tuple[ec.EllipticCurvePrivateKey, ec.EllipticCurvePublicKey]
-    ] = None,
-    max_retries: Optional[int] = None,
-    max_wait_time: Optional[float] = None,
-    flwr_path: Optional[Path] = None,
+    authentication_keys: (
+        tuple[ec.EllipticCurvePrivateKey, ec.EllipticCurvePublicKey] | None
+    ) = None,
+    max_retries: int | None = None,
+    max_wait_time: float | None = None,
+    flwr_path: Path | None = None,
     isolation: str = ISOLATION_MODE_SUBPROCESS,
     clientappio_api_address: str = CLIENTAPPIO_API_DEFAULT_SERVER_ADDRESS,
-    health_server_address: Optional[str] = None,
+    health_server_address: str | None = None,
+    trusted_entities: dict[str, str] | None = None,
 ) -> None:
     """Start a Flower client node which connects to a Flower server.
 
@@ -138,6 +155,10 @@ def start_client_internal(
     health_server_address : Optional[str] (default: None)
         The address of the health server. If `None` is provided, the health server will
         NOT be started.
+    trusted_entities : Optional[dict[str, str]] (default: None)
+        A dictionary mapping public key IDs to public keys.
+        Only apps verified by at least one of these
+        entities can run on a supernode.
     """
     if insecure is None:
         insecure = root_certificates is None
@@ -156,9 +177,9 @@ def start_client_internal(
         )
 
     # Initialize factories
-    state_factory = NodeStateFactory()
-    ffs_factory = FfsFactory(get_flwr_dir(flwr_path) / "supernode" / "ffs")  # type: ignore
     object_store_factory = ObjectStoreFactory()
+    state_factory = NodeStateFactory(objectstore_factory=object_store_factory)
+    ffs_factory = FfsFactory(get_flwr_dir(flwr_path) / "supernode" / "ffs")  # type: ignore
 
     # Launch ClientAppIo API server
     grpc_servers = []
@@ -218,6 +239,7 @@ def start_client_internal(
         ) = conn
         # Store node_id in state
         state.set_node_id(node_id)
+        log(INFO, "SuperNode ID: %s", node_id)
 
         # pylint: disable=too-many-nested-blocks
         while True:
@@ -233,6 +255,7 @@ def start_client_internal(
                 get_fab=get_fab,
                 pull_object=pull_object,
                 confirm_message_received=confirm_message_received,
+                trusted_entities=trusted_entities,
             )
 
             # No message has been pulled therefore we can skip the push stage.
@@ -249,17 +272,34 @@ def start_client_internal(
             )
 
 
+def _insert_message(msg: Message, state: NodeState, store: ObjectStore) -> None:
+    """Insert a message into the NodeState and ObjectStore."""
+    with no_object_id_recompute():
+        # Store message in state
+        msg.metadata.__dict__["_message_id"] = msg.object_id  # Set message_id
+        state.store_message(msg)
+
+        # Preregister objects in ObjectStore
+        store.preregister(msg.metadata.run_id, get_object_tree(msg))
+
+        # Store all objects in ObjectStore
+        all_objects = get_all_nested_objects(msg)
+        for obj_id, obj in all_objects.items():
+            store.put(obj_id, obj.deflate())
+
+
 def _pull_and_store_message(  # pylint: disable=too-many-positional-arguments
     state: NodeState,
     ffs: Ffs,
     object_store: ObjectStore,
     node_config: UserConfig,
-    receive: Callable[[], Optional[tuple[Message, ObjectTree]]],
+    receive: Callable[[], tuple[Message, ObjectTree] | None],
     get_run: Callable[[int], Run],
     get_fab: Callable[[str, int], Fab],
     pull_object: Callable[[int, str], bytes],
     confirm_message_received: Callable[[int, str], None],
-) -> Optional[int]:
+    trusted_entities: dict[str, str] | None,
+) -> int | None:
     """Pull a message from the SuperLink and store it in the state.
 
     This function current returns None if no message is received,
@@ -267,6 +307,7 @@ def _pull_and_store_message(  # pylint: disable=too-many-positional-arguments
     This behavior will change in the future to return None after
     completing transition to the `NodeState`-based SuperNode.
     """
+    # pylint: disable=too-many-nested-blocks
     message = None
     try:
         # Pull message
@@ -287,7 +328,7 @@ def _pull_and_store_message(  # pylint: disable=too-many-positional-arguments
             log(INFO, "[RUN %s]", message.metadata.run_id)
         log(
             INFO,
-            "Received: %s message %s",
+            "Receiving: %s message (ID: %s)",
             message.metadata.message_type,
             message.metadata.message_id,
         )
@@ -299,11 +340,32 @@ def _pull_and_store_message(  # pylint: disable=too-many-positional-arguments
         if (run_info := state.get_run(run_id)) is None:
             # Pull run info from SuperLink
             run_info = get_run(run_id)
-            state.store_run(run_info)
 
             # Pull and store the FAB
             fab = get_fab(run_info.fab_hash, run_id)
-            ffs.put(fab.content, {})
+
+            # Verify the received FAB
+            # FAB must be signed if trust entities provided
+            if trusted_entities:
+                if not fab.verifications.get("valid_license", ""):
+                    log(
+                        WARN,
+                        "App verification is not supported by the connected SuperLink.",
+                    )
+                else:
+                    fab_verified = _verify_fab(fab, trusted_entities)
+                    if not fab_verified:
+                        # Insert an error message in the state
+                        # when FAB verification fails
+                        log(
+                            ERROR,
+                            "FAB verification failed: the provided trusted entities "
+                            "could not verify the FAB. An error reply "
+                            "has been generated.",
+                        )
+                        reply = Message(FAB_VERIFICATION_ERROR, reply_to=message)
+                        _insert_message(reply, state, object_store)
+                        return run_id
 
             # Initialize the context
             run_cfg = get_fused_config_from_fab(fab.content, run_info)
@@ -314,7 +376,11 @@ def _pull_and_store_message(  # pylint: disable=too-many-positional-arguments
                 state=RecordDict(),
                 run_config=run_cfg,
             )
+
+            # Store in the state
             state.store_context(run_ctx)
+            state.store_run(run_info)
+            ffs.put(fab.content, fab.verifications)
 
         # Preregister the object tree of the message
         obj_ids_to_pull = object_store.preregister(run_id, object_tree)
@@ -322,16 +388,27 @@ def _pull_and_store_message(  # pylint: disable=too-many-positional-arguments
         # Store the message in the state (note this message has no content)
         state.store_message(message)
 
-        # Pull and store objects of the message in the ObjectStore
-        obj_contents = pull_objects(
-            obj_ids_to_pull,
-            pull_object_fn=lambda obj_id: pull_object(run_id, obj_id),
-        )
-        for obj_id in list(obj_contents.keys()):
-            object_store.put(obj_id, obj_contents.pop(obj_id))
+        try:
+            # Pull and store objects of the message in the ObjectStore
+            obj_contents = pull_objects(
+                obj_ids_to_pull,
+                pull_object_fn=lambda obj_id: pull_object(run_id, obj_id),
+            )
+            for obj_id in list(obj_contents.keys()):
+                object_store.put(obj_id, obj_contents.pop(obj_id))
 
-        # Confirm that the message was received
-        confirm_message_received(run_id, message.metadata.message_id)
+            # Confirm that the message was received
+            confirm_message_received(run_id, message.metadata.message_id)
+            log(INFO, "Received successfully")
+        except Exception as err:  # pylint: disable=broad-except
+            log(
+                ERROR,
+                "Failed to receive message %s: %s",
+                message.metadata.message_id,
+                err,
+            )
+            state.delete_messages(message_ids=[message.metadata.message_id])
+            object_store.delete(message.metadata.message_id)
 
     except RunNotRunningException:
         if message is None:
@@ -378,8 +455,9 @@ def _push_messages(
             log(INFO, "[RUN %s]", message.metadata.run_id)
         log(
             INFO,
-            "Sending: %s message",
+            "Sending: %s message (ID: %s)",
             message.metadata.message_type,
+            message.metadata.message_id,
         )
 
         # Get the object tree for the message
@@ -424,6 +502,13 @@ def _push_messages(
                 message.metadata.run_id,
                 message.metadata.message_id,
             )
+        except Exception as err:  # pylint: disable=broad-except
+            log(
+                ERROR,
+                "Failed to send message %s: %s",
+                message.metadata.message_id,
+                err,
+            )
         finally:
             # Delete the message from the state
             state.delete_messages(
@@ -444,16 +529,16 @@ def _init_connection(  # pylint: disable=too-many-positional-arguments
     transport: str,
     server_address: str,
     insecure: bool,
-    root_certificates: Optional[Union[bytes, str]] = None,
-    authentication_keys: Optional[
-        tuple[ec.EllipticCurvePrivateKey, ec.EllipticCurvePublicKey]
-    ] = None,
-    max_retries: Optional[int] = None,
-    max_wait_time: Optional[float] = None,
+    root_certificates: bytes | str | None = None,
+    authentication_keys: (
+        tuple[ec.EllipticCurvePrivateKey, ec.EllipticCurvePublicKey] | None
+    ) = None,
+    max_retries: int | None = None,
+    max_wait_time: float | None = None,
 ) -> Iterator[
     tuple[
         int,
-        Callable[[], Optional[tuple[Message, ObjectTree]]],
+        Callable[[], tuple[Message, ObjectTree] | None],
         Callable[[Message, ObjectTree], set[str]],
         Callable[[int], Run],
         Callable[[str, int], Fab],
@@ -513,8 +598,8 @@ def _init_connection(  # pylint: disable=too-many-positional-arguments
 
 
 def _make_fleet_connection_retry_invoker(
-    max_retries: Optional[int] = None,
-    max_wait_time: Optional[float] = None,
+    max_retries: int | None = None,
+    max_wait_time: float | None = None,
     connection_error_type: type[Exception] = RpcError,
 ) -> RetryInvoker:
     """Create a retry invoker for fleet connection."""
@@ -533,7 +618,7 @@ def run_clientappio_api_grpc(
     state_factory: NodeStateFactory,
     ffs_factory: FfsFactory,
     objectstore_factory: ObjectStoreFactory,
-    certificates: Optional[tuple[bytes, bytes, bytes]],
+    certificates: tuple[bytes, bytes, bytes] | None,
 ) -> grpc.Server:
     """Run ClientAppIo API gRPC server."""
     clientappio_servicer: grpc.Server = ClientAppIoServicer(
@@ -554,3 +639,34 @@ def run_clientappio_api_grpc(
     log(INFO, "Flower Deployment Runtime: Starting ClientAppIo API on %s", address)
     clientappio_grpc_server.start()
     return clientappio_grpc_server
+
+
+def _verify_fab(fab: Fab, trusted_entities: dict[str, str]) -> bool:
+    """Verify a FAB using its verification data and the provided trusted entities.
+
+    The FAB is considered verified if at least one trusted entity matches the
+    information contained in its verification records.
+    """
+    verifications = fab.verifications
+    verif_full = {
+        k: json.loads(v) for k, v in verifications.items() if k != "valid_license"
+    }
+    fab_verified = False
+    for public_key_id, verif in verif_full.items():
+        if public_key_id in trusted_entities:
+            verifier_public_key = load_ssh_public_key(
+                trusted_entities[public_key_id].encode("utf-8")
+            )
+            message_to_verify = create_message_to_sign(
+                hashlib.sha256(fab.content).digest(),
+                verif["signed_at"],
+            )
+            assert isinstance(verifier_public_key, ed25519.Ed25519PublicKey)
+            if verify_signature(
+                verifier_public_key,
+                message_to_verify,
+                decode_base64url(verif["signature"]),
+            ):
+                fab_verified = True
+                break
+    return fab_verified

{flwr-1.23.0.dist-info → flwr-1.24.0.dist-info}/METADATA

@@ -1,12 +1,12 @@
 Metadata-Version: 2.3
 Name: flwr
-Version: 1.23.0
+Version: 1.24.0
 Summary: Flower: A Friendly Federated AI Framework
 License: Apache-2.0
 Keywords: Artificial Intelligence,Federated AI,Federated Analytics,Federated Evaluation,Federated Learning,Flower,Machine Learning
 Author: The Flower Authors
 Author-email: hello@flower.ai
-Requires-Python: >=3.9.2,<4.0.0
+Requires-Python: >=3.10,<4.0
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Science/Research
@@ -20,7 +20,6 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3 :: Only
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
@@ -33,21 +32,22 @@ Provides-Extra: rest
 Provides-Extra: simulation
 Requires-Dist: click (<8.2.0)
 Requires-Dist: cryptography (>=44.0.1,<45.0.0)
-Requires-Dist: grpcio (>=1.62.3,<2.0.0,!=1.65.0)
-Requires-Dist: grpcio-health-checking (>=1.62.3,<2.0.0)
+Requires-Dist: grpcio (>=1.70.0,<2.0.0)
+Requires-Dist: grpcio-health-checking (>=1.70.0,<2.0.0)
 Requires-Dist: iterators (>=0.0.2,<0.0.3)
 Requires-Dist: numpy (>=1.26.0,<3.0.0)
 Requires-Dist: pathspec (>=0.12.1,<0.13.0)
-Requires-Dist: protobuf (>=4.21.6,<5.0.0)
+Requires-Dist: protobuf (>=5.28.0,<7.0.0)
 Requires-Dist: pycryptodome (>=3.18.0,<4.0.0)
 Requires-Dist: pyyaml (>=6.0.2,<7.0.0)
-Requires-Dist: ray (==2.31.0) ; (python_version >= "3.9" and python_version < "3.13") and (extra == "simulation")
+Requires-Dist: ray (==2.51.1) ; (python_version >= "3.10" and python_version < "3.13") and (extra == "simulation")
+Requires-Dist: ray (==2.51.1) ; (sys_platform != "win32" and python_version == "3.13") and (extra == "simulation")
 Requires-Dist: requests (>=2.31.0,<3.0.0)
 Requires-Dist: rich (>=13.5.0,<14.0.0)
 Requires-Dist: starlette (>=0.45.2,<0.46.0) ; extra == "rest"
 Requires-Dist: tomli (>=2.0.1,<3.0.0)
 Requires-Dist: tomli-w (>=1.0.0,<2.0.0)
-Requires-Dist: typer (>=0.12.5,<0.13.0)
+Requires-Dist: typer (>=0.12.5,<0.21.0)
 Requires-Dist: uvicorn[standard] (>=0.34.0,<0.35.0) ; extra == "rest"
 Project-URL: Documentation, https://flower.ai
 Project-URL: Homepage, https://flower.ai