flwr-nightly 1.23.0.dev20251007__py3-none-any.whl → 1.23.0.dev20251008__py3-none-any.whl

flwr/server/superlink/fleet/grpc_rere/{server_interceptor.py → node_auth_server_interceptor.py}

@@ -16,7 +16,7 @@
 
 
 import datetime
-from typing import Any, Callable, Optional, cast
+from typing import Any, Callable, cast
 
 import grpc
 from google.protobuf.message import Message as GrpcMessage
@@ -29,10 +29,7 @@ from flwr.common.constant import (
     TIMESTAMP_HEADER,
     TIMESTAMP_TOLERANCE,
 )
-from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
-    CreateNodeRequest,
-    CreateNodeResponse,
-)
+from flwr.proto.fleet_pb2 import CreateNodeRequest  # pylint: disable=E0611
 from flwr.server.superlink.linkstate import LinkStateFactory
 from flwr.supercore.primitives.asymmetric import bytes_to_public_key, verify_signature
 
@@ -50,7 +47,7 @@ def _unary_unary_rpc_terminator(
     return grpc.unary_unary_rpc_method_handler(terminate)
 
 
-class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
+class NodeAuthServerInterceptor(grpc.ServerInterceptor):  # type: ignore
     """Server interceptor for node authentication.
 
     Parameters
@@ -110,50 +107,34 @@ class AuthenticateServerInterceptor(grpc.ServerInterceptor):  # type: ignore
         if not MIN_TIMESTAMP_DIFF < time_diff.total_seconds() < MAX_TIMESTAMP_DIFF:
             return _unary_unary_rpc_terminator("Invalid timestamp")
 
-        # Continue the RPC call
-        expected_node_id = state.get_node_id(node_pk_bytes)
-        if not handler_call_details.method.endswith("CreateNode"):
-            # All calls, except for `CreateNode`, must provide a public key that is
-            # already mapped to a `node_id` (in `LinkState`)
-            if expected_node_id is None:
-                return _unary_unary_rpc_terminator("Invalid node ID")
-        # One of the method handlers in
+        # Continue the RPC call: One of the method handlers in
         # `flwr.server.superlink.fleet.grpc_rere.fleet_server.FleetServicer`
         method_handler: grpc.RpcMethodHandler = continuation(handler_call_details)
-        return self._wrap_method_handler(
-            method_handler, expected_node_id, node_pk_bytes
-        )
+        return self._wrap_method_handler(method_handler, node_pk_bytes)
 
     def _wrap_method_handler(
         self,
         method_handler: grpc.RpcMethodHandler,
-        expected_node_id: Optional[int],
-        node_public_key: bytes,
+        expected_public_key: bytes,
     ) -> grpc.RpcMethodHandler:
         def _generic_method_handler(
             request: GrpcMessage,
             context: grpc.ServicerContext,
         ) -> GrpcMessage:
-            # Verify the node ID
-            if not isinstance(request, CreateNodeRequest):
-                try:
-                    if request.node.node_id != expected_node_id:  # type: ignore
-                        raise ValueError
-                except (AttributeError, ValueError):
-                    context.abort(grpc.StatusCode.UNAUTHENTICATED, "Invalid node ID")
+            # Retrieve the public key
+            if isinstance(request, CreateNodeRequest):
+                actual_public_key = request.public_key
+            else:
+                # Note: This function runs in a different thread
+                # than the `intercept_service` function.
+                actual_public_key = self.state_factory.state().get_node_public_key(
+                    request.node.node_id  # type: ignore
+                )
+            # Verify the public key
+            if actual_public_key != expected_public_key:
+                context.abort(grpc.StatusCode.UNAUTHENTICATED, "Invalid node ID")
 
             response: GrpcMessage = method_handler.unary_unary(request, context)
-
-            # Set the public key after a successful CreateNode request
-            if isinstance(response, CreateNodeResponse):
-                state = self.state_factory.state()
-                try:
-                    state.set_node_public_key(response.node.node_id, node_public_key)
-                except ValueError as e:
-                    # Remove newly created node if setting the public key fails
-                    state.delete_node(response.node.node_id)
-                    context.abort(grpc.StatusCode.UNAUTHENTICATED, str(e))
-
             return response
 
         return grpc.unary_unary_rpc_method_handler(

flwr/server/superlink/fleet/message_handler/message_handler.py

@@ -70,7 +70,7 @@ def create_node(
 ) -> CreateNodeResponse:
     """."""
     # Create node
-    node_id = state.create_node(heartbeat_interval=request.heartbeat_interval)
+    node_id = state.create_node(request.public_key, request.heartbeat_interval)
     return CreateNodeResponse(node=Node(node_id=node_id))
 
 

flwr/server/superlink/fleet/vce/vce_api.py

@@ -16,6 +16,7 @@
 
 
 import json
+import secrets
 import threading
 import time
 import traceback
@@ -53,7 +54,12 @@ def _register_nodes(
     nodes_mapping: NodeToPartitionMapping = {}
     state = state_factory.state()
     for i in range(num_nodes):
-        node_id = state.create_node(heartbeat_interval=HEARTBEAT_MAX_INTERVAL)
+        node_id = state.create_node(
+            # No node authentication in simulation;
+            # use random bytes instead
+            secrets.token_bytes(32),
+            heartbeat_interval=HEARTBEAT_MAX_INTERVAL,
+        )
         nodes_mapping[node_id] = i
     log(DEBUG, "Registered %i nodes", len(nodes_mapping))
     return nodes_mapping

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -17,7 +17,6 @@
 
 import secrets
 import threading
-import time
 from bisect import bisect_right
 from collections import defaultdict
 from dataclasses import dataclass, field
@@ -39,6 +38,7 @@ from flwr.common.constant import (
 )
 from flwr.common.record import ConfigRecord
 from flwr.common.typing import Run, RunStatus, UserConfig
+from flwr.proto.node_pb2 import NodeInfo  # pylint: disable=E0611
 from flwr.server.superlink.linkstate.linkstate import LinkState
 from flwr.server.utils import validate_message
 
@@ -70,7 +70,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
     def __init__(self) -> None:
 
         # Map node_id to (online_until, heartbeat_interval)
-        self.node_ids: dict[int, tuple[float, float]] = {}
+        self.nodes: dict[int, NodeInfo] = {}
         self.public_key_to_node_id: dict[bytes, int] = {}
         self.node_id_to_public_key: dict[int, bytes] = {}
 
@@ -114,7 +114,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             )
             return None
         # Validate destination node ID
-        if message.metadata.dst_node_id not in self.node_ids:
+        if message.metadata.dst_node_id not in self.nodes:
             log(
                 ERROR,
                 "Invalid destination node ID for Message: %s",
@@ -136,7 +136,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
 
         # Find Message for node_id that were not delivered yet
         message_ins_list: list[Message] = []
-        current_time = time.time()
+        current_time = now().timestamp()
         with self.lock:
             for _, msg_ins in self.message_ins_store.items():
                 if (
@@ -190,7 +190,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 return None
 
             ins_metadata = msg_ins.metadata
-            if ins_metadata.created_at + ins_metadata.ttl <= time.time():
+            if ins_metadata.created_at + ins_metadata.ttl <= now().timestamp():
                 log(
                     ERROR,
                     "Failed to store Message: the message it is replying to "
@@ -238,7 +238,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         ret: dict[str, Message] = {}
 
         with self.lock:
-            current = time.time()
+            current = now().timestamp()
 
             # Verify Message IDs
             ret = verify_message_ids(
@@ -256,9 +256,9 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 inquired_in_message_ids=message_ids,
                 found_in_message_dict=self.message_ins_store,
                 node_id_to_online_until={
-                    node_id: self.node_ids[node_id][0]
+                    node_id: self.nodes[node_id].online_until
                     for node_id in dst_node_ids
-                    if node_id in self.node_ids
+                    if node_id in self.nodes
                 },
                 current_time=current,
             )
@@ -330,7 +330,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         """
         return len(self.message_res_store)
 
-    def create_node(self, heartbeat_interval: float) -> int:
+    def create_node(self, public_key: bytes, heartbeat_interval: float) -> int:
         """Create, store in the link state, and return `node_id`."""
         # Sample a random int64 as node_id
         node_id = generate_rand_int_from_bytes(
@@ -338,28 +338,40 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         )
 
         with self.lock:
-            if node_id in self.node_ids:
+            if node_id in self.nodes:
                 log(ERROR, "Unexpected node registration failure.")
                 return 0
+            if public_key in self.public_key_to_node_id:
+                raise ValueError("Public key already in use")
 
-            # Mark the node online until time.time() + heartbeat_interval
-            self.node_ids[node_id] = (
-                time.time() + heartbeat_interval,
-                heartbeat_interval,
+            # Mark the node online until now().timestamp() + heartbeat_interval
+            current = now()
+            self.nodes[node_id] = NodeInfo(
+                node_id=node_id,
+                owner_aid="",  # Unused for now
+                status="created",  # Unused for now
+                created_at=current.isoformat(),  # Unused for now
+                last_activated_at=current.isoformat(),  # Unused for now
+                last_deactivated_at="",  # Unused for now
+                deleted_at="",  # Unused for now
+                online_until=current.timestamp() + heartbeat_interval,
+                heartbeat_interval=heartbeat_interval,
             )
+            self.public_key_to_node_id[public_key] = node_id
+            self.node_id_to_public_key[node_id] = public_key
             return node_id
 
     def delete_node(self, node_id: int) -> None:
         """Delete a node."""
         with self.lock:
-            if node_id not in self.node_ids:
+            if node_id not in self.nodes:
                 raise ValueError(f"Node {node_id} not found")
 
             # Remove node ID <> public key mappings
             if pk := self.node_id_to_public_key.pop(node_id, None):
                 del self.public_key_to_node_id[pk]
 
-            del self.node_ids[node_id]
+            del self.nodes[node_id]
 
     def get_nodes(self, run_id: int) -> set[int]:
         """Return all available nodes.
@@ -372,17 +384,17 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         with self.lock:
             if run_id not in self.run_ids:
                 return set()
-            current_time = time.time()
+            current_time = now().timestamp()
             return {
-                node_id
-                for node_id, (online_until, _) in self.node_ids.items()
-                if online_until > current_time
+                info.node_id
+                for info in self.nodes.values()
+                if info.online_until > current_time
             }
 
     def set_node_public_key(self, node_id: int, public_key: bytes) -> None:
         """Set `public_key` for the specified `node_id`."""
         with self.lock:
-            if node_id not in self.node_ids:
+            if node_id not in self.nodes:
                 raise ValueError(f"Node {node_id} not found")
 
             if public_key in self.public_key_to_node_id:
@@ -394,7 +406,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
     def get_node_public_key(self, node_id: int) -> Optional[bytes]:
         """Get `public_key` for the specified `node_id`."""
         with self.lock:
-            if node_id not in self.node_ids:
+            if node_id not in self.nodes:
                 raise ValueError(f"Node {node_id} not found")
 
             return self.node_id_to_public_key.get(node_id)
@@ -608,13 +620,13 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         the node is marked as offline.
         """
         with self.lock:
-            if node_id in self.node_ids:
-                self.node_ids[node_id] = (
-                    time.time() + HEARTBEAT_PATIENCE * heartbeat_interval,
-                    heartbeat_interval,
+            if info := self.nodes.get(node_id):
+                info.online_until = (
+                    now().timestamp() + HEARTBEAT_PATIENCE * heartbeat_interval
                 )
+                info.heartbeat_interval = heartbeat_interval
                 return True
-        return False
+            return False
 
     def acknowledge_app_heartbeat(self, run_id: int, heartbeat_interval: float) -> bool:
         """Acknowledge a heartbeat received from a ServerApp for a given run.

flwr/server/superlink/linkstate/linkstate.py

@@ -128,7 +128,7 @@ class LinkState(CoreState):  # pylint: disable=R0904
         """Get all instruction Message IDs for the given run_id."""
 
     @abc.abstractmethod
-    def create_node(self, heartbeat_interval: float) -> int:
+    def create_node(self, public_key: bytes, heartbeat_interval: float) -> int:
         """Create, store in the link state, and return `node_id`."""
 
     @abc.abstractmethod

flwr/server/superlink/linkstate/sqlite_linkstate.py

@@ -21,7 +21,6 @@ import json
 import re
 import secrets
 import sqlite3
-import time
 from collections.abc import Sequence
 from logging import DEBUG, ERROR, WARNING
 from typing import Any, Optional, Union, cast
@@ -72,10 +71,16 @@ from .utils import (
 
 SQL_CREATE_TABLE_NODE = """
 CREATE TABLE IF NOT EXISTS node(
-    node_id         INTEGER UNIQUE,
-    online_until    REAL,
-    heartbeat_interval   REAL,
-    public_key      BLOB
+    node_id                 INTEGER UNIQUE,
+    owner_aid               TEXT,
+    status                  TEXT,
+    created_at              TEXT,
+    last_activated_at       TEXT,
+    last_deactivated_at     TEXT,
+    deleted_at              TEXT,
+    online_until            REAL,
+    heartbeat_interval      REAL,
+    public_key              BLOB UNIQUE
 );
 """
 
@@ -451,7 +456,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         ret: dict[str, Message] = {}
 
         # Verify Message IDs
-        current = time.time()
+        current = now().timestamp()
         query = f"""
             SELECT *
             FROM message_ins
@@ -597,7 +602,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         return {row["message_id"] for row in rows}
 
-    def create_node(self, heartbeat_interval: float) -> int:
+    def create_node(self, public_key: bytes, heartbeat_interval: float) -> int:
         """Create, store in the link state, and return `node_id`."""
         # Sample a random uint64 as node_id
         uint64_node_id = generate_rand_int_from_bytes(
@@ -607,24 +612,35 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         # Convert the uint64 value to sint64 for SQLite
         sint64_node_id = convert_uint64_to_sint64(uint64_node_id)
 
-        query = (
-            "INSERT INTO node "
-            "(node_id, online_until, heartbeat_interval, public_key) "
-            "VALUES (?, ?, ?, ?)"
-        )
+        query = """
+            INSERT INTO node
+            (node_id, owner_aid, status, created_at, last_activated_at,
+            last_deactivated_at, deleted_at, online_until, heartbeat_interval,
+            public_key)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+        """
 
-        # Mark the node online util time.time() + heartbeat_interval
+        # Mark the node online until now().timestamp() + heartbeat_interval
         try:
             self.query(
                 query,
                 (
-                    sint64_node_id,
-                    time.time() + heartbeat_interval,
-                    heartbeat_interval,
-                    b"",  # Initialize with an empty public key
+                    sint64_node_id,  # node_id
+                    "",  # owner_aid, unused for now
+                    "created",  # status, unused for now
+                    now().isoformat(),  # created_at, unused for now
+                    now().isoformat(),  # last_activated_at, unused for now
+                    "",  # last_deactivated_at, unused for now
+                    "",  # deleted_at, unused for now
+                    now().timestamp() + heartbeat_interval,  # online_until
+                    heartbeat_interval,  # heartbeat_interval
+                    public_key,  # public_key
                 ),
             )
-        except sqlite3.IntegrityError:
+        except sqlite3.IntegrityError as e:
+            if "UNIQUE constraint failed: node.public_key" in str(e):
+                raise ValueError("Public key already in use.") from None
+            # Must be node ID conflict, almost impossible unless system is compromised
             log(ERROR, "Unexpected node registration failure.")
             return 0
 
@@ -668,7 +684,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         # Get nodes
         query = "SELECT node_id FROM node WHERE online_until > ?;"
-        rows = self.query(query, (time.time(),))
+        rows = self.query(query, (now().timestamp(),))
 
         # Convert sint64 node_ids to uint64
         result: set[int] = {convert_sint64_to_uint64(row["node_id"]) for row in rows}
@@ -1010,7 +1026,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         self.query(
             query,
             (
-                time.time() + HEARTBEAT_PATIENCE * heartbeat_interval,
+                now().timestamp() + HEARTBEAT_PATIENCE * heartbeat_interval,
                 heartbeat_interval,
                 sint64_node_id,
             ),
@@ -1140,7 +1156,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         message_ins = rows[0]
         created_at = message_ins["created_at"]
         ttl = message_ins["ttl"]
-        current_time = time.time()
+        current_time = now().timestamp()
 
         # Check if Message is expired
         if ttl is not None and created_at + ttl <= current_time:

flwr/server/utils/validator.py

@@ -15,10 +15,9 @@
 """Validators."""
 
 
-import time
-
 from flwr.common import Message
 from flwr.common.constant import SUPERLINK_NODE_ID
+from flwr.common.date import now
 
 
 # pylint: disable-next=too-many-branches
@@ -44,7 +43,7 @@ def validate_message(message: Message, is_reply_message: bool) -> list[str]:
         validation_errors.append("`metadata.ttl` must be higher than zero")
 
     # Verify TTL and created_at time
-    current_time = time.time()
+    current_time = now().timestamp()
     if metadata.created_at + metadata.ttl <= current_time:
         validation_errors.append("Message TTL has expired")
 

flwr/superlink/auth_plugin/__init__.py

@@ -15,12 +15,41 @@
 """Account auth plugin for ControlServicer."""
 
 
+from flwr.common.constant import AuthnType, AuthzType
+
 from .auth_plugin import ControlAuthnPlugin, ControlAuthzPlugin
 from .noop_auth_plugin import NoOpControlAuthnPlugin, NoOpControlAuthzPlugin
 
+try:
+    from flwr.ee import get_control_authn_ee_plugins, get_control_authz_ee_plugins
+except ImportError:
+
+    def get_control_authn_ee_plugins() -> dict[str, type[ControlAuthnPlugin]]:
+        """Return all Control API authentication plugins for EE."""
+        return {}
+
+    def get_control_authz_ee_plugins() -> dict[str, type[ControlAuthzPlugin]]:
+        """Return all Control API authorization plugins for EE."""
+        return {}
+
+
+def get_control_authn_plugins() -> dict[str, type[ControlAuthnPlugin]]:
+    """Return all Control API authentication plugins."""
+    ee_dict: dict[str, type[ControlAuthnPlugin]] = get_control_authn_ee_plugins()
+    return ee_dict | {AuthnType.NOOP: NoOpControlAuthnPlugin}
+
+
+def get_control_authz_plugins() -> dict[str, type[ControlAuthzPlugin]]:
+    """Return all Control API authorization plugins."""
+    ee_dict: dict[str, type[ControlAuthzPlugin]] = get_control_authz_ee_plugins()
+    return ee_dict | {AuthzType.NOOP: NoOpControlAuthzPlugin}
+
+
 __all__ = [
     "ControlAuthnPlugin",
     "ControlAuthzPlugin",
     "NoOpControlAuthnPlugin",
     "NoOpControlAuthzPlugin",
+    "get_control_authn_plugins",
+    "get_control_authz_plugins",
 ]

flwr/superlink/servicer/control/control_grpc.py

@@ -31,7 +31,11 @@ from flwr.supercore.ffs import FfsFactory
 from flwr.supercore.license_plugin import LicensePlugin
 from flwr.supercore.object_store import ObjectStoreFactory
 from flwr.superlink.artifact_provider import ArtifactProvider
-from flwr.superlink.auth_plugin import ControlAuthnPlugin, ControlAuthzPlugin
+from flwr.superlink.auth_plugin import (
+    ControlAuthnPlugin,
+    ControlAuthzPlugin,
+    NoOpControlAuthnPlugin,
+)
 
 from .control_account_auth_interceptor import ControlAccountAuthInterceptor
 from .control_event_log_interceptor import ControlEventLogInterceptor
@@ -54,8 +58,8 @@ def run_control_api_grpc(
     objectstore_factory: ObjectStoreFactory,
     certificates: Optional[tuple[bytes, bytes, bytes]],
     is_simulation: bool,
-    authn_plugin: Optional[ControlAuthnPlugin] = None,
-    authz_plugin: Optional[ControlAuthzPlugin] = None,
+    authn_plugin: ControlAuthnPlugin,
+    authz_plugin: ControlAuthzPlugin,
     event_log_plugin: Optional[EventLogWriterPlugin] = None,
     artifact_provider: Optional[ArtifactProvider] = None,
 ) -> grpc.Server:
@@ -72,11 +76,9 @@ def run_control_api_grpc(
         authn_plugin=authn_plugin,
         artifact_provider=artifact_provider,
     )
-    interceptors: list[grpc.ServerInterceptor] = []
+    interceptors = [ControlAccountAuthInterceptor(authn_plugin, authz_plugin)]
     if license_plugin is not None:
         interceptors.append(ControlLicenseInterceptor(license_plugin))
-    if authn_plugin is not None and authz_plugin is not None:
-        interceptors.append(ControlAccountAuthInterceptor(authn_plugin, authz_plugin))
     # Event log interceptor must be added after account auth interceptor
     if event_log_plugin is not None:
         interceptors.append(ControlEventLogInterceptor(event_log_plugin))
@@ -90,7 +92,7 @@ def run_control_api_grpc(
         interceptors=interceptors or None,
     )
 
-    if authn_plugin is None:
+    if isinstance(authn_plugin, NoOpControlAuthnPlugin):
         log(INFO, "Flower Deployment Runtime: Starting Control API on %s", address)
     else:
         log(