Add server auth interceptor

src/py/flwr/common/secure_aggregation/crypto/symmetric_encryption.py

@@ -117,3 +117,18 @@ def verify_hmac(key: bytes, message: bytes, hmac_value: bytes) -> bool:
         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"
+    )

src/py/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,10 @@
 )
 from flwr.common.exit_handlers import register_exit_handlers
 from flwr.common.logger import log
+from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
+    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 +61,7 @@
     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 +365,28 @@ 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
+            interceptors = [
+                AuthenticateServerInterceptor(
+                    client_public_keys,
+                    server_private_key,
+                    server_public_key,
+                )
+            ]
+
         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 +419,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 +510,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 +523,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 +701,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:

src/py/flwr/server/server_test.py

@@ -15,9 +15,22 @@
 """Flower server tests."""
 
 
+import argparse
+import csv
+import tempfile
+from pathlib import Path
 from typing import List, Optional
 
 import numpy as np
+from cryptography.hazmat.primitives.asymmetric import ec
+from cryptography.hazmat.primitives.serialization import (
+    Encoding,
+    NoEncryption,
+    PrivateFormat,
+    PublicFormat,
+    load_ssh_private_key,
+    load_ssh_public_key,
+)
 
 from flwr.common import (
     Code,
@@ -35,8 +48,14 @@
     Status,
     ndarray_to_bytes,
 )
+from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
+    generate_key_pairs,
+    private_key_to_bytes,
+    public_key_to_bytes,
+)
 from flwr.server.client_manager import SimpleClientManager
 
+from .app import _try_setup_client_authentication
 from .client_proxy import ClientProxy
 from .server import Server, evaluate_clients, fit_clients
 
@@ -182,3 +201,71 @@ def test_set_max_workers() -> None:
 
     # Assert
     assert server.max_workers == 42
+
+
+def test_setup_client_auth() -> None:  # pylint: disable=R0914
+    """Test setup client authentication."""
+    # Prepare
+    _, first_public_key = generate_key_pairs()
+    private_key, public_key = generate_key_pairs()
+
+    server_public_key = public_key.public_bytes(
+        encoding=Encoding.OpenSSH, format=PublicFormat.OpenSSH
+    )
+    server_private_key = private_key.private_bytes(
+        Encoding.PEM, PrivateFormat.OpenSSH, NoEncryption()
+    )
+    _, second_public_key = generate_key_pairs()
+
+    # Execute
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Initialize temporary files
+        client_keys_file_path = Path(temp_dir) / "client_keys.csv"
+        server_private_key_path = Path(temp_dir) / "server_private_key"
+        server_public_key_path = Path(temp_dir) / "server_public_key"
+
+        # Fill the files with relevant keys
+        with open(client_keys_file_path, "w", newline="", encoding="utf-8") as csvfile:
+            writer = csv.writer(csvfile)
+            writer.writerow(
+                [
+                    first_public_key.public_bytes(
+                        encoding=Encoding.OpenSSH, format=PublicFormat.OpenSSH
+                    ).decode(),
+                    second_public_key.public_bytes(
+                        encoding=Encoding.OpenSSH, format=PublicFormat.OpenSSH
+                    ).decode(),
+                ]
+            )
+        server_public_key_path.write_bytes(server_public_key)
+        server_private_key_path.write_bytes(server_private_key)
+
+        # Mock argparse with `require-client-authentication`` flag
+        mock_args = argparse.Namespace(
+            require_client_authentication=[
+                str(client_keys_file_path),
+                str(server_private_key_path),
+                str(server_public_key_path),
+            ]
+        )
+
+        # Run _try_setup_client_authentication
+        result = _try_setup_client_authentication(mock_args, (b"", b"", b""))
+
+        expected_private_key = load_ssh_private_key(server_private_key, None)
+        expected_public_key = load_ssh_public_key(server_public_key)
+
+        # Assert
+        assert isinstance(expected_private_key, ec.EllipticCurvePrivateKey)
+        assert isinstance(expected_public_key, ec.EllipticCurvePublicKey)
+        assert result is not None
+        assert result[0] == {
+            public_key_to_bytes(first_public_key),
+            public_key_to_bytes(second_public_key),
+        }
+        assert private_key_to_bytes(result[1]) == private_key_to_bytes(
+            expected_private_key
+        )
+        assert public_key_to_bytes(result[2]) == public_key_to_bytes(
+            expected_public_key
+        )

src/py/flwr/server/superlink/fleet/grpc_bidi/grpc_server.py

@@ -18,7 +18,7 @@
 import concurrent.futures
 import sys
 from logging import ERROR
-from typing import Any, Callable, Optional, Tuple, Union
+from typing import Any, Callable, Optional, Sequence, Tuple, Union
 
 import grpc
 
@@ -162,6 +162,7 @@ def generic_create_grpc_server(  # pylint: disable=too-many-arguments
     max_message_length: int = GRPC_MAX_MESSAGE_LENGTH,
     keepalive_time_ms: int = 210000,
     certificates: Optional[Tuple[bytes, bytes, bytes]] = None,
+    interceptors: Optional[Sequence[grpc.ServerInterceptor]] = None,
 ) -> grpc.Server:
     """Create a gRPC server with a single servicer.
 
@@ -249,6 +250,7 @@ def generic_create_grpc_server(  # pylint: disable=too-many-arguments
         # returning RESOURCE_EXHAUSTED status, or None to indicate no limit.
         maximum_concurrent_rpcs=max_concurrent_workers,
         options=options,
+        interceptors=interceptors,
     )
     add_servicer_to_server_fn(servicer, server)