From 87725f7f6d3ec5e51fa9b9662ac63860f5c3bb20 Mon Sep 17 00:00:00 2001 From: Chong Shen Ng Date: Thu, 19 Dec 2024 10:54:39 +0000 Subject: [PATCH 1/2] refactor(framework) Add `.credentials` to `.gitignore` (#4745) --- src/py/flwr/cli/utils.py | 47 ++++++++++++++++++++++++++++++++++++++-- 1 file changed, 45 insertions(+), 2 deletions(-) diff --git a/src/py/flwr/cli/utils.py b/src/py/flwr/cli/utils.py index 8cb89255ed40..9968895ef73b 100644 --- a/src/py/flwr/cli/utils.py +++ b/src/py/flwr/cli/utils.py @@ -161,10 +161,53 @@ def get_sha256_hash(file_path: Path) -> str: def get_user_auth_config_path(root_dir: Path, federation: str) -> Path: - """Return the path to the user auth config file.""" + """Return the path to the user auth config file. + + Additionally, a `.gitignore` file will be created in the Flower directory to + include the `.credentials` folder to be excluded from git. If the `.gitignore` + file already exists, a warning will be displayed if the `.credentials` entry is + not found. + """ # Locate the credentials directory - credentials_dir = root_dir.absolute() / FLWR_DIR / CREDENTIALS_DIR + abs_flwr_dir = root_dir.absolute() / FLWR_DIR + credentials_dir = abs_flwr_dir / CREDENTIALS_DIR credentials_dir.mkdir(parents=True, exist_ok=True) + + # Determine the absolute path of the Flower directory for .gitignore + gitignore_path = abs_flwr_dir / ".gitignore" + credential_entry = CREDENTIALS_DIR + + try: + if gitignore_path.exists(): + with open(gitignore_path, encoding="utf-8") as gitignore_file: + lines = gitignore_file.read().splitlines() + + # Warn if .credentials is not already in .gitignore + if credential_entry not in lines: + typer.secho( + f"`.gitignore` exists, but `{credential_entry}` entry not found. " + "Consider adding it to your `.gitignore` to exclude Flower " + "credentials from git.", + fg=typer.colors.YELLOW, + bold=True, + ) + else: + typer.secho( + f"Creating a new `.gitignore` with `{credential_entry}` entry...", + fg=typer.colors.BLUE, + ) + # Create a new .gitignore with .credentials + with open(gitignore_path, "w", encoding="utf-8") as gitignore_file: + gitignore_file.write(f"{credential_entry}\n") + except Exception as err: + typer.secho( + "❌ An error occurred while handling `.gitignore.` " + f"Please check the permissions of `{gitignore_path}` and try again.", + fg=typer.colors.RED, + bold=True, + ) + raise typer.Exit(code=1) from err + return credentials_dir / f"{federation}.json" From 08d02f0038502b9e95364316aeb8b344cdaba75e Mon Sep 17 00:00:00 2001 From: Javier Date: Thu, 19 Dec 2024 11:10:13 +0000 Subject: [PATCH 2/2] fix(framework) Clear node auth keys upon `SuperLink` start (#4736) --- src/py/flwr/server/app.py | 1 + .../linkstate/in_memory_linkstate.py | 7 ++++++ .../server/superlink/linkstate/linkstate.py | 4 ++++ .../superlink/linkstate/linkstate_test.py | 23 +++++++++++++++++++ .../superlink/linkstate/sqlite_linkstate.py | 6 +++++ 5 files changed, 41 insertions(+) diff --git a/src/py/flwr/server/app.py b/src/py/flwr/server/app.py index 91514f845651..b5c7ae95e224 100644 --- a/src/py/flwr/server/app.py +++ b/src/py/flwr/server/app.py @@ -374,6 +374,7 @@ def run_superlink() -> None: server_public_key, ) = maybe_keys state = state_factory.state() + state.clear_supernode_auth_keys_and_credentials() state.store_node_public_keys(node_public_keys) state.store_server_private_public_key( private_key_to_bytes(server_private_key), diff --git a/src/py/flwr/server/superlink/linkstate/in_memory_linkstate.py b/src/py/flwr/server/superlink/linkstate/in_memory_linkstate.py index f26bb11a4bdb..d22072b41621 100644 --- a/src/py/flwr/server/superlink/linkstate/in_memory_linkstate.py +++ b/src/py/flwr/server/superlink/linkstate/in_memory_linkstate.py @@ -430,6 +430,13 @@ def get_server_public_key(self) -> Optional[bytes]: """Retrieve `server_public_key` in urlsafe bytes.""" return self.server_public_key + def clear_supernode_auth_keys_and_credentials(self) -> None: + """Clear stored `node_public_keys` and credentials in the link state if any.""" + with self.lock: + self.server_private_key = None + self.server_public_key = None + self.node_public_keys.clear() + def store_node_public_keys(self, public_keys: set[bytes]) -> None: """Store a set of `node_public_keys` in the link state.""" with self.lock: diff --git a/src/py/flwr/server/superlink/linkstate/linkstate.py b/src/py/flwr/server/superlink/linkstate/linkstate.py index ae9d1710f069..4f3c16a5460a 100644 --- a/src/py/flwr/server/superlink/linkstate/linkstate.py +++ b/src/py/flwr/server/superlink/linkstate/linkstate.py @@ -284,6 +284,10 @@ def get_server_private_key(self) -> Optional[bytes]: def get_server_public_key(self) -> Optional[bytes]: """Retrieve `server_public_key` in urlsafe bytes.""" + @abc.abstractmethod + def clear_supernode_auth_keys_and_credentials(self) -> None: + """Clear stored `node_public_keys` and credentials in the link state if any.""" + @abc.abstractmethod def store_node_public_keys(self, public_keys: set[bytes]) -> None: """Store a set of `node_public_keys` in the link state.""" diff --git a/src/py/flwr/server/superlink/linkstate/linkstate_test.py b/src/py/flwr/server/superlink/linkstate/linkstate_test.py index 93f5d94daef7..3edaf72ec20c 100644 --- a/src/py/flwr/server/superlink/linkstate/linkstate_test.py +++ b/src/py/flwr/server/superlink/linkstate/linkstate_test.py @@ -820,6 +820,29 @@ def test_store_server_private_public_key_twice(self) -> None: new_private_key_bytes, new_public_key_bytes ) + def test_clear_supernode_auth_keys_and_credentials(self) -> None: + """Test clear_supernode_auth_keys_and_credentials from linkstate.""" + # Prepare + state: LinkState = self.state_factory() + key_pairs = [generate_key_pairs() for _ in range(3)] + public_keys = {public_key_to_bytes(pair[1]) for pair in key_pairs} + + # Execute (store) + state.store_node_public_keys(public_keys) + private_key, public_key = generate_key_pairs() + private_key_bytes = private_key_to_bytes(private_key) + public_key_bytes = public_key_to_bytes(public_key) + state.store_server_private_public_key(private_key_bytes, public_key_bytes) + + # Execute (clear) + state.clear_supernode_auth_keys_and_credentials() + node_public_keys = state.get_node_public_keys() + + # Assert + assert node_public_keys == set() + assert state.get_server_private_key() is None + assert state.get_server_public_key() is None + def test_node_public_keys(self) -> None: """Test store_node_public_keys and get_node_public_keys from state.""" # Prepare diff --git a/src/py/flwr/server/superlink/linkstate/sqlite_linkstate.py b/src/py/flwr/server/superlink/linkstate/sqlite_linkstate.py index 99334233319d..e8311dfaac5e 100644 --- a/src/py/flwr/server/superlink/linkstate/sqlite_linkstate.py +++ b/src/py/flwr/server/superlink/linkstate/sqlite_linkstate.py @@ -818,6 +818,12 @@ def get_server_public_key(self) -> Optional[bytes]: public_key = None return public_key + def clear_supernode_auth_keys_and_credentials(self) -> None: + """Clear stored `node_public_keys` and credentials in the link state if any.""" + queries = ["DELETE FROM public_key;", "DELETE FROM credential;"] + for query in queries: + self.query(query) + def store_node_public_keys(self, public_keys: set[bytes]) -> None: """Store a set of `node_public_keys` in the link state.""" query = "INSERT INTO public_key (public_key) VALUES (?)"