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

flwr/common/message.py

@@ -18,14 +18,13 @@ from __future__ import annotations
 
 import time
 import warnings
-from dataclasses import dataclass
+from typing import Optional, cast
 
 from .record import RecordSet
 
 DEFAULT_TTL = 3600
 
 
-@dataclass
 class Metadata:  # pylint: disable=too-many-instance-attributes
     """A dataclass holding metadata associated with the current message.
 
@@ -55,17 +54,6 @@ class Metadata:  # pylint: disable=too-many-instance-attributes
         is more relevant when conducting simulations.
     """
 
-    _run_id: int
-    _message_id: str
-    _src_node_id: int
-    _dst_node_id: int
-    _reply_to_message: str
-    _group_id: str
-    _ttl: float
-    _message_type: str
-    _partition_id: int | None
-    _created_at: float  # Unix timestamp (in seconds) to be set upon message creation
-
     def __init__(  # pylint: disable=too-many-arguments
         self,
         run_id: int,
@@ -78,98 +66,111 @@ class Metadata:  # pylint: disable=too-many-instance-attributes
         message_type: str,
         partition_id: int | None = None,
     ) -> None:
-        self._run_id = run_id
-        self._message_id = message_id
-        self._src_node_id = src_node_id
-        self._dst_node_id = dst_node_id
-        self._reply_to_message = reply_to_message
-        self._group_id = group_id
-        self._ttl = ttl
-        self._message_type = message_type
-        self._partition_id = partition_id
+        var_dict = {
+            "_run_id": run_id,
+            "_message_id": message_id,
+            "_src_node_id": src_node_id,
+            "_dst_node_id": dst_node_id,
+            "_reply_to_message": reply_to_message,
+            "_group_id": group_id,
+            "_ttl": ttl,
+            "_message_type": message_type,
+            "_partition_id": partition_id,
+        }
+        self.__dict__.update(var_dict)
 
     @property
     def run_id(self) -> int:
         """An identifier for the current run."""
-        return self._run_id
+        return cast(int, self.__dict__["_run_id"])
 
     @property
     def message_id(self) -> str:
         """An identifier for the current message."""
-        return self._message_id
+        return cast(str, self.__dict__["_message_id"])
 
     @property
     def src_node_id(self) -> int:
         """An identifier for the node sending this message."""
-        return self._src_node_id
+        return cast(int, self.__dict__["_src_node_id"])
 
     @property
     def reply_to_message(self) -> str:
         """An identifier for the message this message replies to."""
-        return self._reply_to_message
+        return cast(str, self.__dict__["_reply_to_message"])
 
     @property
     def dst_node_id(self) -> int:
         """An identifier for the node receiving this message."""
-        return self._dst_node_id
+        return cast(int, self.__dict__["_dst_node_id"])
 
     @dst_node_id.setter
     def dst_node_id(self, value: int) -> None:
         """Set dst_node_id."""
-        self._dst_node_id = value
+        self.__dict__["_dst_node_id"] = value
 
     @property
     def group_id(self) -> str:
         """An identifier for grouping messages."""
-        return self._group_id
+        return cast(str, self.__dict__["_group_id"])
 
     @group_id.setter
     def group_id(self, value: str) -> None:
         """Set group_id."""
-        self._group_id = value
+        self.__dict__["_group_id"] = value
 
     @property
     def created_at(self) -> float:
         """Unix timestamp when the message was created."""
-        return self._created_at
+        return cast(float, self.__dict__["_created_at"])
 
     @created_at.setter
     def created_at(self, value: float) -> None:
-        """Set creation timestamp for this messages."""
-        self._created_at = value
+        """Set creation timestamp for this message."""
+        self.__dict__["_created_at"] = value
 
     @property
     def ttl(self) -> float:
         """Time-to-live for this message."""
-        return self._ttl
+        return cast(float, self.__dict__["_ttl"])
 
     @ttl.setter
     def ttl(self, value: float) -> None:
         """Set ttl."""
-        self._ttl = value
+        self.__dict__["_ttl"] = value
 
     @property
     def message_type(self) -> str:
         """A string that encodes the action to be executed on the receiving end."""
-        return self._message_type
+        return cast(str, self.__dict__["_message_type"])
 
     @message_type.setter
     def message_type(self, value: str) -> None:
         """Set message_type."""
-        self._message_type = value
+        self.__dict__["_message_type"] = value
 
     @property
     def partition_id(self) -> int | None:
         """An identifier telling which data partition a ClientApp should use."""
-        return self._partition_id
+        return cast(int, self.__dict__["_partition_id"])
 
     @partition_id.setter
     def partition_id(self, value: int) -> None:
-        """Set patition_id."""
-        self._partition_id = value
+        """Set partition_id."""
+        self.__dict__["_partition_id"] = value
+
+    def __repr__(self) -> str:
+        """Return a string representation of this instance."""
+        view = ", ".join([f"{k.lstrip('_')}={v!r}" for k, v in self.__dict__.items()])
+        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__
 
 
-@dataclass
 class Error:
     """A dataclass that stores information about an error that occurred.
 
@@ -181,25 +182,35 @@ class Error:
         A reason for why the error arose (e.g. an exception stack-trace)
     """
 
-    _code: int
-    _reason: str | None = None
-
     def __init__(self, code: int, reason: str | None = None) -> None:
-        self._code = code
-        self._reason = reason
+        var_dict = {
+            "_code": code,
+            "_reason": reason,
+        }
+        self.__dict__.update(var_dict)
 
     @property
     def code(self) -> int:
         """Error code."""
-        return self._code
+        return cast(int, self.__dict__["_code"])
 
     @property
     def reason(self) -> str | None:
         """Reason reported about the error."""
-        return self._reason
+        return cast(Optional[str], self.__dict__["_reason"])
+
+    def __repr__(self) -> str:
+        """Return a string representation of this instance."""
+        view = ", ".join([f"{k.lstrip('_')}={v!r}" for k, v in self.__dict__.items()])
+        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__
 
 
-@dataclass
 class Message:
     """State of your application from the viewpoint of the entity using it.
 
@@ -215,88 +226,70 @@ class Message:
         when processing another message.
     """
 
-    _metadata: Metadata
-    _content: RecordSet | None = None
-    _error: Error | None = None
-
     def __init__(
         self,
         metadata: Metadata,
         content: RecordSet | None = None,
         error: Error | None = None,
     ) -> None:
-        self._metadata = metadata
-
-        # Set message creation timestamp
-        self._metadata.created_at = time.time()
-
         if not (content is None) ^ (error is None):
             raise ValueError("Either `content` or `error` must be set, but not both.")
 
-        self._content = content
-        self._error = error
+        metadata.created_at = time.time()  # Set the message creation timestamp
+        var_dict = {
+            "_metadata": metadata,
+            "_content": content,
+            "_error": error,
+        }
+        self.__dict__.update(var_dict)
 
     @property
     def metadata(self) -> Metadata:
         """A dataclass including information about the message to be executed."""
-        return self._metadata
+        return cast(Metadata, self.__dict__["_metadata"])
 
     @property
     def content(self) -> RecordSet:
         """The content of this message."""
-        if self._content is None:
+        if self.__dict__["_content"] is None:
             raise ValueError(
                 "Message content is None. Use <message>.has_content() "
                 "to check if a message has content."
             )
-        return self._content
+        return cast(RecordSet, self.__dict__["_content"])
 
     @content.setter
     def content(self, value: RecordSet) -> None:
         """Set content."""
-        if self._error is None:
-            self._content = value
+        if self.__dict__["_error"] is None:
+            self.__dict__["_content"] = value
         else:
             raise ValueError("A message with an error set cannot have content.")
 
     @property
     def error(self) -> Error:
         """Error captured by this message."""
-        if self._error is None:
+        if self.__dict__["_error"] is None:
             raise ValueError(
                 "Message error is None. Use <message>.has_error() "
                 "to check first if a message carries an error."
             )
-        return self._error
+        return cast(Error, self.__dict__["_error"])
 
     @error.setter
     def error(self, value: Error) -> None:
         """Set error."""
         if self.has_content():
             raise ValueError("A message with content set cannot carry an error.")
-        self._error = value
+        self.__dict__["_error"] = value
 
     def has_content(self) -> bool:
         """Return True if message has content, else False."""
-        return self._content is not None
+        return self.__dict__["_content"] is not None
 
     def has_error(self) -> bool:
         """Return True if message has an error, else False."""
-        return self._error is not None
-
-    def _create_reply_metadata(self, ttl: float) -> Metadata:
-        """Construct metadata for a reply message."""
-        return Metadata(
-            run_id=self.metadata.run_id,
-            message_id="",
-            src_node_id=self.metadata.dst_node_id,
-            dst_node_id=self.metadata.src_node_id,
-            reply_to_message=self.metadata.message_id,
-            group_id=self.metadata.group_id,
-            ttl=ttl,
-            message_type=self.metadata.message_type,
-            partition_id=self.metadata.partition_id,
-        )
+        return self.__dict__["_error"] is not None
 
     def create_error_reply(self, error: Error, ttl: float | None = None) -> Message:
         """Construct a reply message indicating an error happened.
@@ -323,7 +316,7 @@ class Message:
         # message creation)
         ttl_ = DEFAULT_TTL if ttl is None else ttl
         # Create reply with error
-        message = Message(metadata=self._create_reply_metadata(ttl_), error=error)
+        message = Message(metadata=_create_reply_metadata(self, ttl_), error=error)
 
         if ttl is None:
             # Set TTL equal to the remaining time for the received message to expire
@@ -369,7 +362,7 @@ class Message:
         ttl_ = DEFAULT_TTL if ttl is None else ttl
 
         message = Message(
-            metadata=self._create_reply_metadata(ttl_),
+            metadata=_create_reply_metadata(self, ttl_),
             content=content,
         )
 
@@ -381,3 +374,29 @@ class Message:
             message.metadata.ttl = ttl
 
         return message
+
+    def __repr__(self) -> str:
+        """Return a string representation of this instance."""
+        view = ", ".join(
+            [
+                f"{k.lstrip('_')}={v!r}"
+                for k, v in self.__dict__.items()
+                if v is not None
+            ]
+        )
+        return f"{self.__class__.__qualname__}({view})"
+
+
+def _create_reply_metadata(msg: Message, ttl: float) -> Metadata:
+    """Construct metadata for a reply message."""
+    return Metadata(
+        run_id=msg.metadata.run_id,
+        message_id="",
+        src_node_id=msg.metadata.dst_node_id,
+        dst_node_id=msg.metadata.src_node_id,
+        reply_to_message=msg.metadata.message_id,
+        group_id=msg.metadata.group_id,
+        ttl=ttl,
+        message_type=msg.metadata.message_type,
+        partition_id=msg.metadata.partition_id,
+    )

flwr/common/record/parametersrecord.py

@@ -82,7 +82,6 @@ def _check_value(value: Array) -> None:
         )
 
 
-@dataclass
 class ParametersRecord(TypedDict[str, Array]):
     """Parameters record.
 

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/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
@@ -36,6 +42,11 @@ from flwr.common.constant import (
 )
 from flwr.common.exit_handlers import register_exit_handlers
 from flwr.common.logger import log
+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
 
@@ -354,10 +366,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 +425,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 +516,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 +529,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 +707,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

@@ -29,7 +29,7 @@ from flwr.server.server import Server, init_defaults, run_fl
 from flwr.server.server_config import ServerConfig
 from flwr.server.strategy import Strategy
 
-from ..driver import Driver
+from ..driver import Driver, GrpcDriver
 from .app_utils import start_update_client_manager_thread
 
 DEFAULT_SERVER_ADDRESS_DRIVER = "[::]:9091"
@@ -114,7 +114,7 @@ def start_driver(  # pylint: disable=too-many-arguments, too-many-locals
         # Create the Driver
         if isinstance(root_certificates, str):
             root_certificates = Path(root_certificates).read_bytes()
-        driver = Driver(
+        driver = GrpcDriver(
             driver_service_address=address, root_certificates=root_certificates
         )
 

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
             )