Add server auth interceptor

src/py/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,21 @@ 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

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

@@ -0,0 +1,47 @@
+# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Symmetric encryption tests."""
+
+
+from .symmetric_encryption import (
+    compute_hmac,
+    generate_key_pairs,
+    generate_shared_key,
+    verify_hmac,
+)
+
+
+def test_generate_shared_key() -> None:
+    """Test util function generate_shared_key."""
+    client_keys = generate_key_pairs()
+    server_keys = generate_key_pairs()
+
+    client_shared_secret = generate_shared_key(client_keys[0], server_keys[1])
+    server_shared_secret = generate_shared_key(server_keys[0], client_keys[1])
+
+    assert client_shared_secret == server_shared_secret
+
+
+def test_hmac() -> None:
+    """Test util function compute and verify hmac."""
+    client_keys = generate_key_pairs()
+    server_keys = generate_key_pairs()
+    client_shared_secret = generate_shared_key(client_keys[0], server_keys[1])
+    server_shared_secret = generate_shared_key(server_keys[0], client_keys[1])
+    message = b"Flower is the future of AI"
+
+    client_compute_hmac = compute_hmac(client_shared_secret, message)
+
+    assert verify_hmac(server_shared_secret, message, client_compute_hmac)

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,9 @@
 )
 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,
+)
 from flwr.proto.fleet_pb2_grpc import (  # pylint: disable=E0611
     add_FleetServicer_to_server,
 )
@@ -44,6 +53,7 @@
 from .history import History
 from .server import Server, init_defaults, run_fl
 from .server_config import ServerConfig
+from .server_interceptor import AuthenticateServerInterceptor
 from .strategy import Strategy
 from .superlink.driver.driver_grpc import run_driver_api_grpc
 from .superlink.fleet.grpc_bidi.grpc_server import (
@@ -352,10 +362,29 @@ 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}"
+
+        data = _try_setup_client_authentication(args, certificates)
+        interceptors: Optional[Sequence[grpc.ServerInterceptor]] = None
+        if data is not None:
+            (
+                client_public_keys,
+                server_public_key,
+                server_private_key,
+            ) = data
+            interceptors = [
+                AuthenticateServerInterceptor(
+                    state_factory,
+                    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:
@@ -388,6 +417,66 @@ 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.EllipticCurvePublicKey, ec.EllipticCurvePrivateKey]]:
+    if args.require_client_authentication:
+        if certificates is None:
+            sys.exit(
+                "Certificates are required to enable client authentication. "
+                "Please provide certificate paths with '--certificates' before "
+                "enabling '--require-client-authentication'."
+            )
+        client_keys_file_path = Path(args.require_client_authentication[0])
+        if client_keys_file_path.exists():
+            client_public_keys: Set[bytes] = set()
+            public_key = load_ssh_public_key(
+                Path(args.require_client_authentication[1]).read_bytes()
+            )
+            private_key = load_ssh_private_key(
+                Path(args.require_client_authentication[2]).read_bytes(),
+                None,
+            )
+            log(INFO, type(public_key))
+            log(INFO, type(private_key))
+            if not isinstance(public_key, ec.EllipticCurvePublicKey):
+                sys.exit(
+                    "An eliptic curve public and private key pair is required for "
+                    "client authentication. Please provide the file path containing "
+                    "valid public and private key to '--require-client-authentication'."
+                )
+            server_public_key = public_key
+            if not isinstance(private_key, ec.EllipticCurvePrivateKey):
+                sys.exit(
+                    "An eliptic curve public and private key pair is required for "
+                    "client authentication. Please provide the file path containing "
+                    "valid public and private key to '--require-client-authentication'."
+                )
+            server_private_key = private_key
+
+            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))
+                return (
+                    client_public_keys,
+                    server_public_key,
+                    server_private_key,
+                )
+        else:
+            sys.exit(
+                "Client public keys csv file are required for client authentication. "
+                "Please provide the csv file path containing known client public keys "
+                "to '--require-client-authentication'."
+            )
+    else:
+        return None
+
+
 def _try_obtain_certificates(
     args: argparse.Namespace,
 ) -> Optional[Tuple[bytes, bytes, bytes]]:
@@ -415,6 +504,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
@@ -427,6 +517,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 +697,14 @@ 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_PUBLIC_KEY", "SERVER_PRIVATE_KEY"),
+        type=str,
+        help="Paths to .csv file containing list of known client public keys for "
+        "authentication, server public key, and server private key, in that order.",
+    )
 
 
 def _add_args_driver_api(parser: argparse.ArgumentParser) -> None: