support conn sec for cell pipe

yanchengnv · Dec 12, 2024 · e001178 · e001178
1 parent f73e127
commit e001178
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 4 deletions.
diff --git a/nvflare/app_common/executors/client_api_launcher_executor.py b/nvflare/app_common/executors/client_api_launcher_executor.py
@@ -15,7 +15,7 @@
 import os
 from typing import Optional
 
-from nvflare.apis.fl_constant import FLMetaKey
+from nvflare.apis.fl_constant import FLMetaKey, SecureTrainConst
 from nvflare.apis.fl_context import FLContext
 from nvflare.app_common.app_constant import AppConstants
 from nvflare.app_common.executors.launcher_executor import LauncherExecutor
@@ -138,6 +138,10 @@ def prepare_config_for_launch(self, fl_ctx: FLContext):
             FLMetaKey.AUTH_TOKEN_SIGNATURE: signature,
         }
 
+        conn_sec = get_scope_property(site_name, SecureTrainConst.CONNECTION_SECURITY)
+        if conn_sec:
+            config_data[SecureTrainConst.CONNECTION_SECURITY] = conn_sec
+
         config_file_path = self._get_external_config_file_path(fl_ctx)
         write_config_to_file(config_data=config_data, config_file_path=config_file_path)
 

diff --git a/nvflare/client/config.py b/nvflare/client/config.py
@@ -16,7 +16,7 @@
 import os
 from typing import Dict, Optional
 
-from nvflare.apis.fl_constant import FLMetaKey
+from nvflare.apis.fl_constant import FLMetaKey, SecureTrainConst
 from nvflare.fuel.utils.config_factory import ConfigFactory
 
 
@@ -165,6 +165,9 @@ def get_auth_token(self):
     def get_auth_token_signature(self):
         return self.config.get(FLMetaKey.AUTH_TOKEN_SIGNATURE)
 
+    def get_connection_security(self):
+        return self.config.get(SecureTrainConst.CONNECTION_SECURITY)
+
     def to_json(self, config_file: str):
         with open(config_file, "w") as f:
             json.dump(self.config, f, indent=2)

diff --git a/nvflare/client/ex_process/api.py b/nvflare/client/ex_process/api.py
@@ -17,7 +17,7 @@
 from typing import Any, Dict, Optional, Tuple
 
 from nvflare.apis.analytix import AnalyticsDataType
-from nvflare.apis.fl_constant import FLMetaKey
+from nvflare.apis.fl_constant import FLMetaKey, SecureTrainConst
 from nvflare.apis.utils.analytix_utils import create_analytic_dxo
 from nvflare.app_common.abstract.fl_model import FLModel
 from nvflare.client.api_spec import APISpec
@@ -45,6 +45,11 @@ def _create_client_config(config: str) -> ClientConfig:
     site_name = client_config.get_site_name()
     set_scope_property(scope_name=site_name, key=FLMetaKey.AUTH_TOKEN, value=auth_token)
     set_scope_property(scope_name=site_name, key=FLMetaKey.AUTH_TOKEN_SIGNATURE, value=signature)
+
+    conn_sec = client_config.get_connection_security()
+    if conn_sec:
+        set_scope_property(site_name, SecureTrainConst.CONNECTION_SECURITY, conn_sec)
+
     return client_config
 
 

diff --git a/nvflare/fuel/utils/pipe/cell_pipe.py b/nvflare/fuel/utils/pipe/cell_pipe.py
@@ -18,7 +18,7 @@
 import time
 from typing import Tuple, Union
 
-from nvflare.apis.fl_constant import CellMessageAuthHeaderKey, FLMetaKey, SystemVarName
+from nvflare.apis.fl_constant import CellMessageAuthHeaderKey, FLMetaKey, SecureTrainConst, SystemVarName
 from nvflare.fuel.data_event.utils import get_scope_property
 from nvflare.fuel.f3.cellnet.cell import Cell
 from nvflare.fuel.f3.cellnet.cell import Message as CellMessage
@@ -149,6 +149,10 @@ def _build_cell(cls, mode, root_url, site_name, token, secure_mode, workspace_di
                         DriverParams.CA_CERT.value: root_cert_path,
                     }
 
+                    conn_sec = get_scope_property(site_name, SecureTrainConst.CONNECTION_SECURITY)
+                    if conn_sec:
+                        credentials[DriverParams.CONNECTION_SECURITY.value] = conn_sec
+
                 cell = Cell(
                     fqcn=_cell_fqcn(mode, site_name, token),
                     root_url=root_url,

diff --git a/nvflare/private/fed/client/fed_client_base.py b/nvflare/private/fed/client/fed_client_base.py
@@ -196,6 +196,7 @@ def _create_cell(self, location, scheme):
             conn_security = self.client_args.get(SecureTrainConst.CONNECTION_SECURITY)
             if conn_security:
                 credentials[DriverParams.CONNECTION_SECURITY.value] = conn_security
+                set_scope_prop(self.client_name, SecureTrainConst.CONNECTION_SECURITY, conn_security)
         else:
             credentials = {}
         self.cell = Cell(