flwr-nightly 1.9.0.dev20240420__py3-none-any.whl → 1.9.0.dev20240509__py3-none-any.whl

flwr/common/record/recordset.py

@@ -24,6 +24,7 @@ from .parametersrecord import ParametersRecord
 from .typeddict import TypedDict
 
 
+@dataclass
 class RecordSetData:
     """Inner data container for the RecordSet class."""
 
@@ -82,7 +83,6 @@ class RecordSetData:
             )
 
 
-@dataclass
 class RecordSet:
     """RecordSet stores groups of parameters, metrics and configs."""
 
@@ -97,22 +97,34 @@ class RecordSet:
             metrics_records=metrics_records,
             configs_records=configs_records,
         )
-        setattr(self, "_data", data)  # noqa
+        self.__dict__["_data"] = data
 
     @property
     def parameters_records(self) -> TypedDict[str, ParametersRecord]:
         """Dictionary holding ParametersRecord instances."""
-        data = cast(RecordSetData, getattr(self, "_data"))  # noqa
+        data = cast(RecordSetData, self.__dict__["_data"])
         return data.parameters_records
 
     @property
     def metrics_records(self) -> TypedDict[str, MetricsRecord]:
         """Dictionary holding MetricsRecord instances."""
-        data = cast(RecordSetData, getattr(self, "_data"))  # noqa
+        data = cast(RecordSetData, self.__dict__["_data"])
         return data.metrics_records
 
     @property
     def configs_records(self) -> TypedDict[str, ConfigsRecord]:
         """Dictionary holding ConfigsRecord instances."""
-        data = cast(RecordSetData, getattr(self, "_data"))  # noqa
+        data = cast(RecordSetData, self.__dict__["_data"])
         return data.configs_records
+
+    def __repr__(self) -> str:
+        """Return a string representation of this instance."""
+        flds = ("parameters_records", "metrics_records", "configs_records")
+        view = ", ".join([f"{fld}={getattr(self, fld)!r}" for fld in flds])
+        return f"{self.__class__.__qualname__}({view})"
+
+    def __eq__(self, other: object) -> bool:
+        """Compare two instances of the class."""
+        if not isinstance(other, self.__class__):
+            raise NotImplementedError
+        return self.__dict__ == other.__dict__

flwr/common/secure_aggregation/crypto/symmetric_encryption.py

@@ -18,8 +18,9 @@
 import base64
 from typing import Tuple, cast
 
+from cryptography.exceptions import InvalidSignature
 from cryptography.fernet import Fernet
-from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives import hashes, hmac, serialization
 from cryptography.hazmat.primitives.asymmetric import ec
 from cryptography.hazmat.primitives.kdf.hkdf import HKDF
 
@@ -98,3 +99,36 @@ def decrypt(key: bytes, ciphertext: bytes) -> bytes:
     # The input key must be url safe
     fernet = Fernet(key)
     return fernet.decrypt(ciphertext)
+
+
+def compute_hmac(key: bytes, message: bytes) -> bytes:
+    """Compute hmac of a message using key as hash."""
+    computed_hmac = hmac.HMAC(key, hashes.SHA256())
+    computed_hmac.update(message)
+    return computed_hmac.finalize()
+
+
+def verify_hmac(key: bytes, message: bytes, hmac_value: bytes) -> bool:
+    """Verify hmac of a message using key as hash."""
+    computed_hmac = hmac.HMAC(key, hashes.SHA256())
+    computed_hmac.update(message)
+    try:
+        computed_hmac.verify(hmac_value)
+        return True
+    except InvalidSignature:
+        return False
+
+
+def ssh_types_to_elliptic_curve(
+    private_key: serialization.SSHPrivateKeyTypes,
+    public_key: serialization.SSHPublicKeyTypes,
+) -> Tuple[ec.EllipticCurvePrivateKey, ec.EllipticCurvePublicKey]:
+    """Cast SSH key types to elliptic curve."""
+    if isinstance(private_key, ec.EllipticCurvePrivateKey) and isinstance(
+        public_key, ec.EllipticCurvePublicKey
+    ):
+        return (private_key, public_key)
+
+    raise TypeError(
+        "The provided key is not an EllipticCurvePrivateKey or EllipticCurvePublicKey"
+    )

flwr/server/__init__.py

@@ -24,7 +24,6 @@ from .app import start_server as start_server
 from .client_manager import ClientManager as ClientManager
 from .client_manager import SimpleClientManager as SimpleClientManager
 from .compat import LegacyContext as LegacyContext
-from .compat import start_driver as start_driver
 from .driver import Driver as Driver
 from .history import History as History
 from .run_serverapp import run_server_app as run_server_app
@@ -45,7 +44,6 @@ __all__ = [
     "ServerApp",
     "ServerConfig",
     "SimpleClientManager",
-    "start_driver",
     "start_server",
     "strategy",
     "workflow",

flwr/server/app.py

@@ -16,15 +16,21 @@
 
 import argparse
 import asyncio
+import csv
 import importlib.util
 import sys
 import threading
 from logging import ERROR, INFO, WARN
 from os.path import isfile
 from pathlib import Path
-from typing import List, Optional, Tuple
+from typing import List, Optional, Sequence, Set, Tuple
 
 import grpc
+from cryptography.hazmat.primitives.asymmetric import ec
+from cryptography.hazmat.primitives.serialization import (
+    load_ssh_private_key,
+    load_ssh_public_key,
+)
 
 from flwr.common import GRPC_MAX_MESSAGE_LENGTH, EventType, event
 from flwr.common.address import parse_address
@@ -35,7 +41,12 @@ from flwr.common.constant import (
     TRANSPORT_TYPE_VCE,
 )
 from flwr.common.exit_handlers import register_exit_handlers
-from flwr.common.logger import log
+from flwr.common.logger import log, warn_deprecated_feature
+from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
+    private_key_to_bytes,
+    public_key_to_bytes,
+    ssh_types_to_elliptic_curve,
+)
 from flwr.proto.fleet_pb2_grpc import (  # pylint: disable=E0611
     add_FleetServicer_to_server,
 )
@@ -51,6 +62,7 @@ from .superlink.fleet.grpc_bidi.grpc_server import (
     start_grpc_server,
 )
 from .superlink.fleet.grpc_rere.fleet_servicer import FleetServicer
+from .superlink.fleet.grpc_rere.server_interceptor import AuthenticateServerInterceptor
 from .superlink.fleet.vce import start_vce
 from .superlink.state import StateFactory
 
@@ -184,6 +196,9 @@ def start_server(  # pylint: disable=too-many-arguments,too-many-locals
 def run_driver_api() -> None:
     """Run Flower server (Driver API)."""
     log(INFO, "Starting Flower server (Driver API)")
+    # Running `flower-driver-api` is deprecated
+    warn_deprecated_feature("flower-driver-api")
+    log(WARN, "Use `flower-superlink` instead")
     event(EventType.RUN_DRIVER_API_ENTER)
     args = _parse_args_run_driver_api().parse_args()
 
@@ -221,6 +236,9 @@ def run_driver_api() -> None:
 def run_fleet_api() -> None:
     """Run Flower server (Fleet API)."""
     log(INFO, "Starting Flower server (Fleet API)")
+    # Running `flower-fleet-api` is deprecated
+    warn_deprecated_feature("flower-fleet-api")
+    log(WARN, "Use `flower-superlink` instead")
     event(EventType.RUN_FLEET_API_ENTER)
     args = _parse_args_run_fleet_api().parse_args()
 
@@ -354,10 +372,33 @@ def run_superlink() -> None:
             sys.exit(f"Fleet IP address ({address_arg}) cannot be parsed.")
         host, port, is_v6 = parsed_address
         address = f"[{host}]:{port}" if is_v6 else f"{host}:{port}"
+
+        maybe_keys = _try_setup_client_authentication(args, certificates)
+        interceptors: Optional[Sequence[grpc.ServerInterceptor]] = None
+        if maybe_keys is not None:
+            (
+                client_public_keys,
+                server_private_key,
+                server_public_key,
+            ) = maybe_keys
+            state = state_factory.state()
+            state.store_client_public_keys(client_public_keys)
+            state.store_server_private_public_key(
+                private_key_to_bytes(server_private_key),
+                public_key_to_bytes(server_public_key),
+            )
+            log(
+                INFO,
+                "Client authentication enabled with %d known public keys",
+                len(client_public_keys),
+            )
+            interceptors = [AuthenticateServerInterceptor(state)]
+
         fleet_server = _run_fleet_api_grpc_rere(
             address=address,
             state_factory=state_factory,
             certificates=certificates,
+            interceptors=interceptors,
         )
         grpc_servers.append(fleet_server)
     elif args.fleet_api_type == TRANSPORT_TYPE_VCE:
@@ -390,6 +431,70 @@ def run_superlink() -> None:
         driver_server.wait_for_termination(timeout=1)
 
 
+def _try_setup_client_authentication(
+    args: argparse.Namespace,
+    certificates: Optional[Tuple[bytes, bytes, bytes]],
+) -> Optional[Tuple[Set[bytes], ec.EllipticCurvePrivateKey, ec.EllipticCurvePublicKey]]:
+    if not args.require_client_authentication:
+        return None
+
+    if certificates is None:
+        sys.exit(
+            "Client authentication only works over secure connections. "
+            "Please provide certificate paths using '--certificates' when "
+            "enabling '--require-client-authentication'."
+        )
+
+    client_keys_file_path = Path(args.require_client_authentication[0])
+    if not client_keys_file_path.exists():
+        sys.exit(
+            "The provided path to the client public keys CSV file does not exist: "
+            f"{client_keys_file_path}. "
+            "Please provide the CSV file path containing known client public keys "
+            "to '--require-client-authentication'."
+        )
+
+    client_public_keys: Set[bytes] = set()
+    ssh_private_key = load_ssh_private_key(
+        Path(args.require_client_authentication[1]).read_bytes(),
+        None,
+    )
+    ssh_public_key = load_ssh_public_key(
+        Path(args.require_client_authentication[2]).read_bytes()
+    )
+
+    try:
+        server_private_key, server_public_key = ssh_types_to_elliptic_curve(
+            ssh_private_key, ssh_public_key
+        )
+    except TypeError:
+        sys.exit(
+            "The file paths provided could not be read as a private and public "
+            "key pair. Client authentication requires an elliptic curve public and "
+            "private key pair. Please provide the file paths containing elliptic "
+            "curve private and public keys to '--require-client-authentication'."
+        )
+
+    with open(client_keys_file_path, newline="", encoding="utf-8") as csvfile:
+        reader = csv.reader(csvfile)
+        for row in reader:
+            for element in row:
+                public_key = load_ssh_public_key(element.encode())
+                if isinstance(public_key, ec.EllipticCurvePublicKey):
+                    client_public_keys.add(public_key_to_bytes(public_key))
+                else:
+                    sys.exit(
+                        "Error: Unable to parse the public keys in the .csv "
+                        "file. Please ensure that the .csv file contains valid "
+                        "SSH public keys and try again."
+                    )
+        return (
+            client_public_keys,
+            server_private_key,
+            server_public_key,
+        )
+
+
 def _try_obtain_certificates(
     args: argparse.Namespace,
 ) -> Optional[Tuple[bytes, bytes, bytes]]:
@@ -417,6 +522,7 @@ def _run_fleet_api_grpc_rere(
     address: str,
     state_factory: StateFactory,
     certificates: Optional[Tuple[bytes, bytes, bytes]],
+    interceptors: Optional[Sequence[grpc.ServerInterceptor]] = None,
 ) -> grpc.Server:
     """Run Fleet API (gRPC, request-response)."""
     # Create Fleet API gRPC server
@@ -429,6 +535,7 @@ def _run_fleet_api_grpc_rere(
         server_address=address,
         max_message_length=GRPC_MAX_MESSAGE_LENGTH,
         certificates=certificates,
+        interceptors=interceptors,
     )
 
     log(INFO, "Flower ECE: Starting Fleet API (gRPC-rere) on %s", address)
@@ -606,6 +713,15 @@ def _add_args_common(parser: argparse.ArgumentParser) -> None:
         "Flower will just create a state in memory.",
         default=DATABASE,
     )
+    parser.add_argument(
+        "--require-client-authentication",
+        nargs=3,
+        metavar=("CLIENT_KEYS", "SERVER_PRIVATE_KEY", "SERVER_PUBLIC_KEY"),
+        type=str,
+        help="Provide three file paths: (1) a .csv file containing a list of "
+        "known client public keys for authentication, (2) the server's private "
+        "key file, and (3) the server's public key file.",
+    )
 
 
 def _add_args_driver_api(parser: argparse.ArgumentParser) -> None:

flwr/server/compat/app.py

@@ -15,14 +15,11 @@
 """Flower driver app."""
 
 
-import sys
 from logging import INFO
-from pathlib import Path
-from typing import Optional, Union
+from typing import Optional
 
 from flwr.common import EventType, event
-from flwr.common.address import parse_address
-from flwr.common.logger import log, warn_deprecated_feature
+from flwr.common.logger import log
 from flwr.server.client_manager import ClientManager
 from flwr.server.history import History
 from flwr.server.server import Server, init_defaults, run_fl
@@ -32,33 +29,21 @@ from flwr.server.strategy import Strategy
 from ..driver import Driver
 from .app_utils import start_update_client_manager_thread
 
-DEFAULT_SERVER_ADDRESS_DRIVER = "[::]:9091"
-
-ERROR_MESSAGE_DRIVER_NOT_CONNECTED = """
-[Driver] Error: Not connected.
-
-Call `connect()` on the `Driver` instance before calling any of the other `Driver`
-methods.
-"""
-
 
 def start_driver(  # pylint: disable=too-many-arguments, too-many-locals
     *,
-    server_address: str = DEFAULT_SERVER_ADDRESS_DRIVER,
+    driver: Driver,
     server: Optional[Server] = None,
     config: Optional[ServerConfig] = None,
     strategy: Optional[Strategy] = None,
     client_manager: Optional[ClientManager] = None,
-    root_certificates: Optional[Union[bytes, str]] = None,
-    driver: Optional[Driver] = None,
 ) -> History:
     """Start a Flower Driver API server.
 
     Parameters
     ----------
-    server_address : Optional[str]
-        The IPv4 or IPv6 address of the Driver API server.
-        Defaults to `"[::]:8080"`.
+    driver : Driver
+        The Driver object to use.
     server : Optional[flwr.server.Server] (default: None)
         A server implementation, either `flwr.server.Server` or a subclass
         thereof. If no instance is provided, then `start_driver` will create
@@ -74,50 +59,14 @@ def start_driver(  # pylint: disable=too-many-arguments, too-many-locals
         An implementation of the class `flwr.server.ClientManager`. If no
         implementation is provided, then `start_driver` will use
         `flwr.server.SimpleClientManager`.
-    root_certificates : Optional[Union[bytes, str]] (default: None)
-        The PEM-encoded root certificates as a byte string or a path string.
-        If provided, a secure connection using the certificates will be
-        established to an SSL-enabled Flower server.
-    driver : Optional[Driver] (default: None)
-        The Driver object to use.
 
     Returns
     -------
     hist : flwr.server.history.History
         Object containing training and evaluation metrics.
-
-    Examples
-    --------
-    Starting a driver that connects to an insecure server:
-
-    >>> start_driver()
-
-    Starting a driver that connects to an SSL-enabled server:
-
-    >>> start_driver(
-    >>>     root_certificates=Path("/crts/root.pem").read_bytes()
-    >>> )
     """
     event(EventType.START_DRIVER_ENTER)
 
-    if driver is None:
-        # Not passing a `Driver` object is deprecated
-        warn_deprecated_feature("start_driver")
-
-        # Parse IP address
-        parsed_address = parse_address(server_address)
-        if not parsed_address:
-            sys.exit(f"Server IP address ({server_address}) cannot be parsed.")
-        host, port, is_v6 = parsed_address
-        address = f"[{host}]:{port}" if is_v6 else f"{host}:{port}"
-
-        # Create the Driver
-        if isinstance(root_certificates, str):
-            root_certificates = Path(root_certificates).read_bytes()
-        driver = Driver(
-            driver_service_address=address, root_certificates=root_certificates
-        )
-
     # Initialize the Driver API server and config
     initialized_server, initialized_config = init_defaults(
         server=server,

flwr/server/compat/app_utils.py

@@ -89,7 +89,7 @@ def _update_client_manager(
         for node_id in new_nodes:
             client_proxy = DriverClientProxy(
                 node_id=node_id,
-                driver=driver.grpc_driver_helper,  # type: ignore
+                driver=driver,
                 anonymous=False,
                 run_id=driver.run_id,  # type: ignore
             )

flwr/server/compat/driver_client_proxy.py

@@ -16,16 +16,14 @@
 
 
 import time
-from typing import List, Optional
+from typing import Optional
 
 from flwr import common
-from flwr.common import DEFAULT_TTL, MessageType, MessageTypeLegacy, RecordSet
+from flwr.common import Message, MessageType, MessageTypeLegacy, RecordSet
 from flwr.common import recordset_compat as compat
-from flwr.common import serde
-from flwr.proto import driver_pb2, node_pb2, task_pb2  # pylint: disable=E0611
 from flwr.server.client_proxy import ClientProxy
 
-from ..driver.driver import GrpcDriverHelper
+from ..driver.driver import Driver
 
 SLEEP_TIME = 1
 
@@ -33,9 +31,7 @@ SLEEP_TIME = 1
 class DriverClientProxy(ClientProxy):
     """Flower client proxy which delegates work using the Driver API."""
 
-    def __init__(
-        self, node_id: int, driver: GrpcDriverHelper, anonymous: bool, run_id: int
-    ):
+    def __init__(self, node_id: int, driver: Driver, anonymous: bool, run_id: int):
         super().__init__(str(node_id))
         self.node_id = node_id
         self.driver = driver
@@ -116,80 +112,39 @@ class DriverClientProxy(ClientProxy):
         timeout: Optional[float],
         group_id: Optional[int],
     ) -> RecordSet:
-        task_ins = task_pb2.TaskIns(  # pylint: disable=E1101
-            task_id="",
-            group_id=str(group_id) if group_id is not None else "",
-            run_id=self.run_id,
-            task=task_pb2.Task(  # pylint: disable=E1101
-                producer=node_pb2.Node(  # pylint: disable=E1101
-                    node_id=0,
-                    anonymous=True,
-                ),
-                consumer=node_pb2.Node(  # pylint: disable=E1101
-                    node_id=self.node_id,
-                    anonymous=self.anonymous,
-                ),
-                task_type=task_type,
-                recordset=serde.recordset_to_proto(recordset),
-                ttl=DEFAULT_TTL,
-            ),
-        )
-
-        # This would normally be recorded upon common.Message creation
-        # but this compatibility stack doesn't create Messages,
-        # so we need to inject `created_at` manually (needed for
-        # taskins validation by server.utils.validator)
-        task_ins.task.created_at = time.time()
 
-        push_task_ins_req = driver_pb2.PushTaskInsRequest(  # pylint: disable=E1101
-            task_ins_list=[task_ins]
+        # Create message
+        message = self.driver.create_message(
+            content=recordset,
+            message_type=task_type,
+            dst_node_id=self.node_id,
+            group_id=str(group_id) if group_id else "",
+            ttl=timeout,
         )
 
-        # Send TaskIns to Driver API
-        push_task_ins_res = self.driver.push_task_ins(req=push_task_ins_req)
-
-        if len(push_task_ins_res.task_ids) != 1:
-            raise ValueError("Unexpected number of task_ids")
+        # Push message
+        message_ids = list(self.driver.push_messages(messages=[message]))
+        if len(message_ids) != 1:
+            raise ValueError("Unexpected number of message_ids")
 
-        task_id = push_task_ins_res.task_ids[0]
-        if task_id == "":
-            raise ValueError(f"Failed to schedule task for node {self.node_id}")
+        message_id = message_ids[0]
+        if message_id == "":
+            raise ValueError(f"Failed to send message to node {self.node_id}")
 
         if timeout:
             start_time = time.time()
 
         while True:
-            pull_task_res_req = driver_pb2.PullTaskResRequest(  # pylint: disable=E1101
-                node=node_pb2.Node(node_id=0, anonymous=True),  # pylint: disable=E1101
-                task_ids=[task_id],
-            )
-
-            # Ask Driver API for TaskRes
-            pull_task_res_res = self.driver.pull_task_res(req=pull_task_res_req)
-
-            task_res_list: List[task_pb2.TaskRes] = list(  # pylint: disable=E1101
-                pull_task_res_res.task_res_list
-            )
-            if len(task_res_list) == 1:
-                task_res = task_res_list[0]
-
-                # This will raise an Exception if task_res carries an `error`
-                validate_task_res(task_res=task_res)
-
-                return serde.recordset_from_proto(task_res.task.recordset)
+            messages = list(self.driver.pull_messages(message_ids))
+            if len(messages) == 1:
+                msg: Message = messages[0]
+                if msg.has_error():
+                    raise ValueError(
+                        f"Message contains an Error (reason: {msg.error.reason}). "
+                        "It originated during client-side execution of a message."
+                    )
+                return msg.content
 
             if timeout is not None and time.time() > start_time + timeout:
                 raise RuntimeError("Timeout reached")
             time.sleep(SLEEP_TIME)
-
-
-def validate_task_res(
-    task_res: task_pb2.TaskRes,  # pylint: disable=E1101
-) -> None:
-    """Validate if a TaskRes is empty or not."""
-    if not task_res.HasField("task"):
-        raise ValueError("Invalid TaskRes, field `task` missing")
-    if task_res.task.HasField("error"):
-        raise ValueError("Exception during client-side task execution")
-    if not task_res.task.HasField("recordset"):
-        raise ValueError("Invalid TaskRes, both `recordset` and `error` are missing")

flwr/server/driver/__init__.py

@@ -16,7 +16,10 @@
 
 
 from .driver import Driver
+from .grpc_driver import GrpcDriver, GrpcDriverHelper
 
 __all__ = [
     "Driver",
+    "GrpcDriver",
+    "GrpcDriverHelper",
 ]