flwr-nightly 1.23.0.dev20250930__py3-none-any.whl → 1.26.0.dev20260121__py3-none-any.whl

flwr/superlink/servicer/control/control_servicer.py

@@ -16,22 +16,30 @@
 
 
 import hashlib
+import json
 import time
-from collections.abc import Generator
+from collections.abc import Generator, Sequence
 from logging import ERROR, INFO
-from typing import Any, Optional, cast
+from typing import Any, cast
 
 import grpc
+import requests
 
 from flwr.cli.config_utils import get_fab_metadata
 from flwr.common import Context, RecordDict, now
 from flwr.common.constant import (
     FAB_MAX_SIZE,
+    HEARTBEAT_DEFAULT_INTERVAL,
     LOG_STREAM_INTERVAL,
+    NO_ACCOUNT_AUTH_MESSAGE,
     NO_ARTIFACT_PROVIDER_MESSAGE,
-    NO_USER_AUTH_MESSAGE,
+    NODE_NOT_FOUND_MESSAGE,
+    PUBLIC_KEY_ALREADY_IN_USE_MESSAGE,
+    PUBLIC_KEY_NOT_VALID,
     PULL_UNFINISHED_RUN_MESSAGE,
     RUN_ID_NOT_FOUND_MESSAGE,
+    SUPERLINK_NODE_ID,
+    TRANSPORT_TYPE_GRPC_ADAPTER,
     Status,
     SubStatus,
 )
@@ -48,24 +56,39 @@ from flwr.proto.control_pb2 import (  # pylint: disable=E0611
     GetAuthTokensResponse,
     GetLoginDetailsRequest,
     GetLoginDetailsResponse,
+    ListFederationsRequest,
+    ListFederationsResponse,
+    ListNodesRequest,
+    ListNodesResponse,
     ListRunsRequest,
     ListRunsResponse,
     PullArtifactsRequest,
     PullArtifactsResponse,
+    RegisterNodeRequest,
+    RegisterNodeResponse,
+    ShowFederationRequest,
+    ShowFederationResponse,
     StartRunRequest,
     StartRunResponse,
     StopRunRequest,
     StopRunResponse,
     StreamLogsRequest,
     StreamLogsResponse,
+    UnregisterNodeRequest,
+    UnregisterNodeResponse,
 )
+from flwr.proto.federation_pb2 import Federation  # pylint: disable=E0611
+from flwr.proto.node_pb2 import NodeInfo  # pylint: disable=E0611
 from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
+from flwr.supercore.constant import PLATFORM_API_URL
 from flwr.supercore.ffs import FfsFactory
 from flwr.supercore.object_store import ObjectStore, ObjectStoreFactory
+from flwr.supercore.primitives.asymmetric import bytes_to_public_key, uses_nist_ec_curve
+from flwr.supercore.utils import parse_app_spec, request_download_link
 from flwr.superlink.artifact_provider import ArtifactProvider
-from flwr.superlink.auth_plugin import ControlAuthPlugin
+from flwr.superlink.auth_plugin import ControlAuthnPlugin
 
-from .control_user_auth_interceptor import shared_account_info
+from .control_account_auth_interceptor import get_current_account_info
 
 
 class ControlServicer(control_pb2_grpc.ControlServicer):
@@ -77,17 +100,19 @@ class ControlServicer(control_pb2_grpc.ControlServicer):
         ffs_factory: FfsFactory,
         objectstore_factory: ObjectStoreFactory,
         is_simulation: bool,
-        auth_plugin: Optional[ControlAuthPlugin] = None,
-        artifact_provider: Optional[ArtifactProvider] = None,
+        authn_plugin: ControlAuthnPlugin,
+        artifact_provider: ArtifactProvider | None = None,
+        fleet_api_type: str | None = None,
     ) -> None:
         self.linkstate_factory = linkstate_factory
         self.ffs_factory = ffs_factory
         self.objectstore_factory = objectstore_factory
         self.is_simulation = is_simulation
-        self.auth_plugin = auth_plugin
+        self.authn_plugin = authn_plugin
         self.artifact_provider = artifact_provider
+        self.fleet_api_type = fleet_api_type
 
-    def StartRun(  # pylint: disable=too-many-locals
+    def StartRun(  # pylint: disable=too-many-locals, too-many-branches, too-many-statements
         self, request: StartRunRequest, context: grpc.ServicerContext
     ) -> StartRunResponse:
         """Create run ID."""
@@ -95,7 +120,15 @@ class ControlServicer(control_pb2_grpc.ControlServicer):
         state = self.linkstate_factory.state()
         ffs = self.ffs_factory.ffs()
 
-        if len(request.fab.content) > FAB_MAX_SIZE:
+        verification_dict: dict[str, str] = {}
+        if request.app_spec:
+            fab_file, verification_dict = _get_remote_fab(
+                self.fleet_api_type, request.app_spec, context
+            )
+        else:
+            fab_file = request.fab.content
+
+        if len(fab_file) > FAB_MAX_SIZE:
             log(
                 ERROR,
                 "FAB size exceeds maximum allowed size of %d bytes.",
@@ -103,10 +136,10 @@ class ControlServicer(control_pb2_grpc.ControlServicer):
             )
             return StartRunResponse()
 
-        flwr_aid = shared_account_info.get().flwr_aid if self.auth_plugin else None
+        flwr_aid = get_current_account_info().flwr_aid
+        flwr_aid = _check_flwr_aid_exists(flwr_aid, context)
         override_config = user_config_from_proto(request.override_config)
         federation_options = config_record_from_proto(request.federation_options)
-        fab_file = request.fab.content
 
         try:
             # Check that num-supernodes is set
@@ -115,9 +148,27 @@ class ControlServicer(control_pb2_grpc.ControlServicer):
                     "Federation options doesn't contain key `num-supernodes`."
                 )
 
+            # Check (1) federation exists and (2) the flwr_aid is a member
+            federation = request.federation
+
+            if not state.federation_manager.exists(federation):
+                raise ValueError(f"Federation '{federation}' does not exist.")
+
+            if not state.federation_manager.has_member(flwr_aid, federation):
+                raise ValueError(
+                    f"Account with ID '{flwr_aid}' is not a member of the "
+                    f"federation '{federation}'. Please log in with another account "
+                    "or request access to this federation."
+                )
+
             # Create run
-            fab = Fab(hashlib.sha256(fab_file).hexdigest(), fab_file)
-            fab_hash = ffs.put(fab.content, {})
+            fab = Fab(
+                hashlib.sha256(fab_file).hexdigest(),
+                fab_file,
+                verification_dict,
+            )
+            fab_hash = ffs.put(fab.content, fab.verifications)
+
             if fab_hash != fab.hash_str:
                 raise RuntimeError(
                     f"FAB ({fab.hash_str}) hash from request doesn't match contents"
@@ -129,6 +180,7 @@ class ControlServicer(control_pb2_grpc.ControlServicer):
                 fab_version,
                 fab_hash,
                 override_config,
+                request.federation,
                 federation_options,
                 flwr_aid,
             )
@@ -144,7 +196,7 @@ class ControlServicer(control_pb2_grpc.ControlServicer):
             # Create an empty context for the Run
             context = Context(
                 run_id=run_id,
-                node_id=0,
+                node_id=SUPERLINK_NODE_ID,
                 # Dict is invariant in mypy
                 node_config=node_config,  # type: ignore[arg-type]
                 state=RecordDict(),
@@ -157,7 +209,10 @@ class ControlServicer(control_pb2_grpc.ControlServicer):
         # pylint: disable-next=broad-except
         except Exception as e:
             log(ERROR, "Could not start run: %s", str(e))
-            return StartRunResponse()
+            context.abort(
+                grpc.StatusCode.FAILED_PRECONDITION,
+                str(e),
+            )
 
         log(INFO, "Created run %s", str(run_id))
         return StartRunResponse(run_id=run_id)
@@ -177,12 +232,9 @@ class ControlServicer(control_pb2_grpc.ControlServicer):
         if not run:
             context.abort(grpc.StatusCode.NOT_FOUND, RUN_ID_NOT_FOUND_MESSAGE)
 
-        # If user auth is enabled, check if `flwr_aid` matches the run's `flwr_aid`
-        if self.auth_plugin:
-            flwr_aid = shared_account_info.get().flwr_aid
-            _check_flwr_aid_in_run(
-                flwr_aid=flwr_aid, run=cast(Run, run), context=context
-            )
+        # Check if `flwr_aid` matches the run's `flwr_aid`
+        flwr_aid = get_current_account_info().flwr_aid
+        _check_flwr_aid_in_run(flwr_aid=flwr_aid, run=cast(Run, run), context=context)
 
         after_timestamp = request.after_timestamp + 1e-6
         while context.is_active():
@@ -218,20 +270,11 @@ class ControlServicer(control_pb2_grpc.ControlServicer):
 
         # Build a set of run IDs for `flwr ls --runs`
         if not request.HasField("run_id"):
-            if self.auth_plugin:
-                # If no `run_id` is specified and user auth is enabled,
-                # return run IDs for the authenticated user
-                flwr_aid = shared_account_info.get().flwr_aid
-                if flwr_aid is None:
-                    context.abort(
-                        grpc.StatusCode.PERMISSION_DENIED,
-                        "️⛔️ User authentication is enabled, but `flwr_aid` is None",
-                    )
-                run_ids = state.get_run_ids(flwr_aid=flwr_aid)
-            else:
-                # If no `run_id` is specified and no user auth is enabled,
-                # return all run IDs
-                run_ids = state.get_run_ids(None)
+            # If no `run_id` is specified and account auth is enabled,
+            # return run IDs for the authenticated account
+            flwr_aid = get_current_account_info().flwr_aid
+            _check_flwr_aid_exists(flwr_aid, context)
+            run_ids = state.get_run_ids(flwr_aid=flwr_aid)
         # Build a set of run IDs for `flwr ls --run-id <run_id>`
         else:
             # Retrieve run ID and run
@@ -241,13 +284,11 @@ class ControlServicer(control_pb2_grpc.ControlServicer):
             # Exit if `run_id` not found
             if not run:
                 context.abort(grpc.StatusCode.NOT_FOUND, RUN_ID_NOT_FOUND_MESSAGE)
+                raise grpc.RpcError()  # This line is unreachable
 
-            # If user auth is enabled, check if `flwr_aid` matches the run's `flwr_aid`
-            if self.auth_plugin:
-                flwr_aid = shared_account_info.get().flwr_aid
-                _check_flwr_aid_in_run(
-                    flwr_aid=flwr_aid, run=cast(Run, run), context=context
-                )
+            # Check if `flwr_aid` matches the run's `flwr_aid`
+            flwr_aid = get_current_account_info().flwr_aid
+            _check_flwr_aid_in_run(flwr_aid=flwr_aid, run=run, context=context)
 
             run_ids = {run_id}
 
@@ -269,13 +310,11 @@ class ControlServicer(control_pb2_grpc.ControlServicer):
         # Exit if `run_id` not found
         if not run:
             context.abort(grpc.StatusCode.NOT_FOUND, RUN_ID_NOT_FOUND_MESSAGE)
+            raise grpc.RpcError()  # This line is unreachable
 
-        # If user auth is enabled, check if `flwr_aid` matches the run's `flwr_aid`
-        if self.auth_plugin:
-            flwr_aid = shared_account_info.get().flwr_aid
-            _check_flwr_aid_in_run(
-                flwr_aid=flwr_aid, run=cast(Run, run), context=context
-            )
+        # Check if `flwr_aid` matches the run's `flwr_aid`
+        flwr_aid = get_current_account_info().flwr_aid
+        _check_flwr_aid_in_run(flwr_aid=flwr_aid, run=run, context=context)
 
         run_status = state.get_run_status({run_id})[run_id]
         if run_status.status == Status.FINISHED:
@@ -284,11 +323,15 @@ class ControlServicer(control_pb2_grpc.ControlServicer):
                 f"Run ID {run_id} is already finished",
             )
 
+        # Update run status to finished:stopped
         update_success = state.update_run_status(
             run_id=run_id,
             new_status=RunStatus(Status.FINISHED, SubStatus.STOPPED, ""),
         )
 
+        # Delete the token associated with the run to stop further operations
+        state.delete_token(run_id)
+
         if update_success:
             message_ids: set[str] = state.get_message_ids_from_run_id(run_id)
 
@@ -305,22 +348,22 @@ class ControlServicer(control_pb2_grpc.ControlServicer):
     ) -> GetLoginDetailsResponse:
         """Start login."""
         log(INFO, "ControlServicer.GetLoginDetails")
-        if self.auth_plugin is None:
+        if self.authn_plugin is None:
             context.abort(
                 grpc.StatusCode.UNIMPLEMENTED,
-                NO_USER_AUTH_MESSAGE,
+                NO_ACCOUNT_AUTH_MESSAGE,
             )
             raise grpc.RpcError()  # This line is unreachable
 
         # Get login details
-        details = self.auth_plugin.get_login_details()
+        details = self.authn_plugin.get_login_details()
 
         # Return empty response if details is None
         if details is None:
             return GetLoginDetailsResponse()
 
         return GetLoginDetailsResponse(
-            auth_type=details.auth_type,
+            authn_type=details.authn_type,
             device_code=details.device_code,
             verification_uri_complete=details.verification_uri_complete,
             expires_in=details.expires_in,
@@ -332,15 +375,15 @@ class ControlServicer(control_pb2_grpc.ControlServicer):
     ) -> GetAuthTokensResponse:
         """Get auth token."""
         log(INFO, "ControlServicer.GetAuthTokens")
-        if self.auth_plugin is None:
+        if self.authn_plugin is None:
             context.abort(
                 grpc.StatusCode.UNIMPLEMENTED,
-                NO_USER_AUTH_MESSAGE,
+                NO_ACCOUNT_AUTH_MESSAGE,
             )
             raise grpc.RpcError()  # This line is unreachable
 
         # Get auth tokens
-        credentials = self.auth_plugin.get_auth_tokens(request.device_code)
+        credentials = self.authn_plugin.get_auth_tokens(request.device_code)
 
         # Return empty response if credentials is None
         if credentials is None:
@@ -383,15 +426,160 @@ class ControlServicer(control_pb2_grpc.ControlServicer):
                 grpc.StatusCode.FAILED_PRECONDITION, PULL_UNFINISHED_RUN_MESSAGE
             )
 
-        # Check if `flwr_aid` matches the run's `flwr_aid` when user auth is enabled
-        if self.auth_plugin:
-            flwr_aid = shared_account_info.get().flwr_aid
-            _check_flwr_aid_in_run(flwr_aid=flwr_aid, run=run, context=context)
+        # Check if `flwr_aid` matches the run's `flwr_aid`
+        flwr_aid = get_current_account_info().flwr_aid
+        _check_flwr_aid_in_run(flwr_aid=flwr_aid, run=run, context=context)
 
         # Call artifact provider
         download_url = self.artifact_provider.get_url(run_id)
         return PullArtifactsResponse(url=download_url)
 
+    def RegisterNode(
+        self, request: RegisterNodeRequest, context: grpc.ServicerContext
+    ) -> RegisterNodeResponse:
+        """Add a SuperNode."""
+        log(INFO, "ControlServicer.RegisterNode")
+
+        # Verify public key
+        try:
+            # Attempt to deserialize public key
+            pub_key = bytes_to_public_key(request.public_key)
+            # Check if it's a NIST EC curve public key
+            if not uses_nist_ec_curve(pub_key):
+                err_msg = "The provided public key is not a NIST EC curve public key."
+                log(ERROR, "%s", err_msg)
+                raise ValueError(err_msg)
+        except (ValueError, AttributeError) as err:
+            log(ERROR, "%s", err)
+            context.abort(grpc.StatusCode.FAILED_PRECONDITION, PUBLIC_KEY_NOT_VALID)
+
+        # Init link state
+        state = self.linkstate_factory.state()
+        node_id = 0
+
+        flwr_aid = get_current_account_info().flwr_aid
+        flwr_aid = _check_flwr_aid_exists(flwr_aid, context)
+        # Account name exists if `flwr_aid` exists
+        account_name = cast(str, get_current_account_info().account_name)
+        try:
+            node_id = state.create_node(
+                owner_aid=flwr_aid,
+                owner_name=account_name,
+                public_key=request.public_key,
+                heartbeat_interval=HEARTBEAT_DEFAULT_INTERVAL,
+            )
+
+        except ValueError:
+            # Public key already in use
+            log(ERROR, PUBLIC_KEY_ALREADY_IN_USE_MESSAGE)
+            context.abort(
+                grpc.StatusCode.FAILED_PRECONDITION, PUBLIC_KEY_ALREADY_IN_USE_MESSAGE
+            )
+        log(INFO, "[ControlServicer.RegisterNode] Created node_id=%s", node_id)
+
+        return RegisterNodeResponse(node_id=node_id)
+
+    def UnregisterNode(
+        self, request: UnregisterNodeRequest, context: grpc.ServicerContext
+    ) -> UnregisterNodeResponse:
+        """Remove a SuperNode."""
+        log(INFO, "ControlServicer.UnregisterNode")
+
+        # Init link state
+        state = self.linkstate_factory.state()
+
+        flwr_aid = get_current_account_info().flwr_aid
+        flwr_aid = _check_flwr_aid_exists(flwr_aid, context)
+        try:
+            state.delete_node(owner_aid=flwr_aid, node_id=request.node_id)
+        except ValueError:
+            log(ERROR, NODE_NOT_FOUND_MESSAGE)
+            context.abort(grpc.StatusCode.NOT_FOUND, NODE_NOT_FOUND_MESSAGE)
+
+        return UnregisterNodeResponse()
+
+    def ListNodes(
+        self, request: ListNodesRequest, context: grpc.ServicerContext
+    ) -> ListNodesResponse:
+        """List all SuperNodes."""
+        log(INFO, "ControlServicer.ListNodes")
+
+        if self.is_simulation:
+            log(ERROR, "ListNodes is not available in simulation mode.")
+            context.abort(
+                grpc.StatusCode.UNIMPLEMENTED,
+                "ListNodes is not available in simulation mode.",
+            )
+            raise grpc.RpcError()  # This line is unreachable
+
+        nodes_info: Sequence[NodeInfo] = []
+        # Init link state
+        state = self.linkstate_factory.state()
+
+        flwr_aid = get_current_account_info().flwr_aid
+        flwr_aid = _check_flwr_aid_exists(flwr_aid, context)
+        # Retrieve all nodes for the account
+        nodes_info = state.get_node_info(owner_aids=[flwr_aid])
+
+        return ListNodesResponse(nodes_info=nodes_info, now=now().isoformat())
+
+    def ListFederations(
+        self, request: ListFederationsRequest, context: grpc.ServicerContext
+    ) -> ListFederationsResponse:
+        """List all SuperNodes."""
+        log(INFO, "ControlServicer.ListFederations")
+
+        # Init link state
+        state = self.linkstate_factory.state()
+
+        flwr_aid = get_current_account_info().flwr_aid
+        flwr_aid = _check_flwr_aid_exists(flwr_aid, context)
+
+        # Get federations the account is a member of
+        federations = state.federation_manager.get_federations(flwr_aid=flwr_aid)
+
+        return ListFederationsResponse(
+            federations=[Federation(name=fed) for fed in federations]
+        )
+
+    def ShowFederation(
+        self, request: ShowFederationRequest, context: grpc.ServicerContext
+    ) -> ShowFederationResponse:
+        """Show details of a specific Federation."""
+        log(INFO, "ControlServicer.ShowFederation")
+
+        # Init link state
+        state = self.linkstate_factory.state()
+
+        flwr_aid = get_current_account_info().flwr_aid
+        flwr_aid = _check_flwr_aid_exists(flwr_aid, context)
+
+        # Get federations the account is a member of
+        federations = state.federation_manager.get_federations(flwr_aid=flwr_aid)
+
+        # Ensure flwr_aid is a member of the requested federation
+        federation = request.federation_name
+        if federation not in federations:
+            context.abort(
+                grpc.StatusCode.FAILED_PRECONDITION,
+                f"Federation '{federation}' does not exist or you are "
+                "not a member of it.",
+            )
+
+        # Fetch federation details
+        details = state.federation_manager.get_details(federation)
+
+        # Build Federation proto object
+        federation_proto = Federation(
+            name=federation,
+            member_aids=details.member_aids,
+            nodes=details.nodes,
+            runs=[run_to_proto(run) for run in details.runs],
+        )
+        return ShowFederationResponse(
+            federation=federation_proto, now=now().isoformat()
+        )
+
 
 def _create_list_runs_response(
     run_ids: set[int], state: LinkState, store: ObjectStore
@@ -410,29 +598,101 @@ def _create_list_runs_response(
     )
 
 
-def _check_flwr_aid_in_run(
-    flwr_aid: Optional[str], run: Run, context: grpc.ServicerContext
-) -> None:
-    """Guard clause to check if `flwr_aid` matches the run's `flwr_aid`."""
-    # `flwr_aid` must not be None. Abort if it is None.
+def _check_flwr_aid_exists(flwr_aid: str | None, context: grpc.ServicerContext) -> str:
+    """Guard clause to check if `flwr_aid` exists."""
     if flwr_aid is None:
         context.abort(
             grpc.StatusCode.PERMISSION_DENIED,
-            "️⛔️ User authentication is enabled, but `flwr_aid` is None",
+            "️⛔️ Failed to fetch the account information.",
         )
+        raise RuntimeError  # This line is unreachable
+    return flwr_aid
+
 
+def _check_flwr_aid_in_run(
+    flwr_aid: str | None, run: Run, context: grpc.ServicerContext
+) -> None:
+    """Guard clause to check if `flwr_aid` matches the run's `flwr_aid`."""
+    _check_flwr_aid_exists(flwr_aid, context)
     # `run.flwr_aid` must not be an empty string. Abort if it is empty.
     run_flwr_aid = run.flwr_aid
     if not run_flwr_aid:
         context.abort(
             grpc.StatusCode.PERMISSION_DENIED,
-            "⛔️ User authentication is enabled, but the run is not associated "
-            "with a `flwr_aid`.",
+            "⛔️ Run is not associated with a `flwr_aid`.",
         )
 
     # Exit if `flwr_aid` does not match the run's `flwr_aid`
     if run_flwr_aid != flwr_aid:
         context.abort(
             grpc.StatusCode.PERMISSION_DENIED,
-            "⛔️ Run ID does not belong to the user",
+            "⛔️ Run ID does not belong to the account",
+        )
+
+
+def _format_verification(verifications: list[dict[str, str]]) -> dict[str, str]:
+    """Format verification information for FAB."""
+    # Convert verifications to dict[str, str] type
+    verification_dict = {
+        item["public_key_id"]: json.dumps(
+            {k: v for k, v in item.items() if k != "public_key_id"}
+        )
+        for item in verifications
+    }
+    verification_dict.update({"valid_license": "Valid"})
+
+    return verification_dict
+
+
+def _get_remote_fab(
+    fleet_api_type: str | None,
+    app_spec: str,
+    context: grpc.ServicerContext,
+) -> tuple[bytes, dict[str, str]]:
+    """Get remote FAB from Flower platform API."""
+    if fleet_api_type == TRANSPORT_TYPE_GRPC_ADAPTER:
+        context.abort(
+            grpc.StatusCode.FAILED_PRECONDITION,
+            "The selected SuperLink transport type is not "
+            "supported for connecting to Flower Platform.",
+        )
+
+    # Parse and validate app specification
+    try:
+        app_id, app_version = parse_app_spec(app_spec)
+    except ValueError as e:
+        context.abort(
+            grpc.StatusCode.FAILED_PRECONDITION,
+            f"{e}",
+        )
+
+    # Request download link and verification information
+    url = f"{PLATFORM_API_URL}/hub/fetch-fab"
+    try:
+        presigned_url, verifications = request_download_link(
+            app_id, app_version, url, "fab_url"
+        )
+    except ValueError as e:
+        context.abort(
+            grpc.StatusCode.FAILED_PRECONDITION,
+            f"{e}",
+        )
+
+    # Format verification information
+    verification_dict = (
+        _format_verification(verifications)
+        if verifications is not None
+        else {"valid_license": ""}
+    )
+
+    # Download FAB from Flower platform API
+    try:
+        r = requests.get(presigned_url, timeout=60)
+        r.raise_for_status()
+    except requests.RequestException as e:
+        context.abort(
+            grpc.StatusCode.FAILED_PRECONDITION,
+            f"FAB download failed: {str(e)}",
         )
+    fab_file = r.content
+    return fab_file, verification_dict