From 2ca1fdf74952a4e6c0d08f95f071b179b8b803cb Mon Sep 17 00:00:00 2001
From: Zhanghao Wu <zhanghao.wu@outlook.com>
Date: Sat, 9 Mar 2024 15:54:09 -0800
Subject: [PATCH] [Core] Avoid race condition for key generation (#3292)

* lock for key generation

* move lock to ~/.sky

* permission
---
 sky/authentication.py | 32 +++++++++++++++++++-------------
 1 file changed, 19 insertions(+), 13 deletions(-)

diff --git a/sky/authentication.py b/sky/authentication.py
index f98037b8885..3ca7141719c 100644
--- a/sky/authentication.py
+++ b/sky/authentication.py
@@ -34,6 +34,7 @@
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import rsa
+import filelock
 import yaml
 
 from sky import clouds
@@ -62,6 +63,7 @@
 # TODO(zhwu): Support user specified key pair.
 PRIVATE_SSH_KEY_PATH = '~/.ssh/sky-key'
 PUBLIC_SSH_KEY_PATH = '~/.ssh/sky-key.pub'
+_SSH_KEY_GENERATION_LOCK = '~/.sky/generated/ssh/.__internal-sky-key.lock'
 
 
 def _generate_rsa_key_pair() -> Tuple[str, str]:
@@ -84,8 +86,8 @@ def _generate_rsa_key_pair() -> Tuple[str, str]:
 
 def _save_key_pair(private_key_path: str, public_key_path: str,
                    private_key: str, public_key: str) -> None:
-    private_key_dir = os.path.dirname(private_key_path)
-    os.makedirs(private_key_dir, exist_ok=True)
+    key_dir = os.path.dirname(private_key_path)
+    os.makedirs(key_dir, exist_ok=True, mode=0o700)
 
     with open(
             private_key_path,
@@ -106,17 +108,21 @@ def get_or_generate_keys() -> Tuple[str, str]:
     """Returns the aboslute private and public key paths."""
     private_key_path = os.path.expanduser(PRIVATE_SSH_KEY_PATH)
     public_key_path = os.path.expanduser(PUBLIC_SSH_KEY_PATH)
-    if not os.path.exists(private_key_path):
-        public_key, private_key = _generate_rsa_key_pair()
-        _save_key_pair(private_key_path, public_key_path, private_key,
-                       public_key)
-    else:
-        # FIXME(skypilot): ran into failing this assert once, but forgot the
-        # reproduction (has private key; but has not generated public key).
-        #   AssertionError: /home/ubuntu/.ssh/sky-key.pub
-        assert os.path.exists(public_key_path), (
-            'Private key found, but associated public key '
-            f'{public_key_path} does not exist.')
+
+    key_file_lock = os.path.expanduser(_SSH_KEY_GENERATION_LOCK)
+    lock_dir = os.path.dirname(key_file_lock)
+    # We should have the folder ~/.sky/generated/ssh to have 0o700 permission,
+    # as the ssh configs will be written to this folder as well in
+    # backend_utils.SSHConfigHelper
+    os.makedirs(lock_dir, exist_ok=True, mode=0o700)
+    with filelock.FileLock(key_file_lock, timeout=10):
+        if not os.path.exists(private_key_path):
+            public_key, private_key = _generate_rsa_key_pair()
+            _save_key_pair(private_key_path, public_key_path, private_key,
+                           public_key)
+    assert os.path.exists(public_key_path), (
+        'Private key found, but associated public key '
+        f'{public_key_path} does not exist.')
     return private_key_path, public_key_path