Skip to content

Commit e2d9ba0

Browse files
Michaelvllclaude
andcommitted
[Core] Fix WSL SSH config issues and add unit tests
- Use functools.lru_cache for thread-safe caching of Windows home directory - Fix shell escaping in proxy command conversion using single quotes to prevent expansion of $, backticks, and other metacharacters - Fix SSH key cleanup to use consistent naming ({cluster_name}.pem) - Add comprehensive unit tests for WSL path conversion and SSH config helpers Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
1 parent 0f51416 commit e2d9ba0

2 files changed

Lines changed: 235 additions & 14 deletions

File tree

sky/utils/cluster_utils.py

Lines changed: 20 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212

1313
from sky import sky_logging
1414
from sky.skylet import constants
15+
from sky.utils import annotations
1516
from sky.utils import command_runner
1617
from sky.utils import common_utils
1718
from sky.utils import lock_events
@@ -22,9 +23,6 @@
2223
# called.
2324
SKY_CLUSTER_YAML_REMOTE_PATH = '~/.sky/sky_ray.yml'
2425

25-
# Cache for WSL Windows home directory
26-
_wsl_windows_home_cached: Optional[str] = None
27-
2826

2927
def _convert_windows_path_to_wsl(windows_path: str) -> str:
3028
"""Convert a Windows path to WSL path format.
@@ -80,17 +78,14 @@ def _get_windows_userprofile_via_cmd() -> Optional[str]:
8078
return None
8179

8280

81+
@annotations.lru_cache(scope='global', maxsize=1)
8382
def get_wsl_windows_home() -> Optional[str]:
8483
"""Get the Windows user's home directory path when running in WSL.
8584
8685
Returns:
8786
The path to Windows home directory (e.g., '/mnt/c/Users/username')
8887
or None if not in WSL or cannot determine.
8988
"""
90-
global _wsl_windows_home_cached
91-
if _wsl_windows_home_cached is not None:
92-
return _wsl_windows_home_cached
93-
9489
if not common_utils.is_wsl():
9590
return None
9691

@@ -102,8 +97,7 @@ def get_wsl_windows_home() -> Optional[str]:
10297
if userprofile:
10398
windows_home = _convert_windows_path_to_wsl(userprofile)
10499
if os.path.isdir(windows_home):
105-
_wsl_windows_home_cached = windows_home
106-
return _wsl_windows_home_cached
100+
return windows_home
107101

108102
return None
109103

@@ -136,6 +130,9 @@ class SSHConfigHelper:
136130
ssh_cluster_key_path = constants.SKY_USER_FILE_PATH + '/ssh-keys/{}.key'
137131

138132
# Windows paths (used when running in WSL)
133+
# Note: These flags are best-effort for UX purposes and are not
134+
# thread-safe. The worst case is duplicate log messages, which is
135+
# acceptable for this non-critical functionality.
139136
_windows_ssh_setup_attempted = False
140137
_windows_ssh_setup_warned = False
141138

@@ -204,14 +201,21 @@ def _convert_proxy_command_for_windows(cls, proxy_command: str) -> str:
204201
Wraps the proxy command with wsl.exe so Windows SSH can execute it.
205202
The command runs inside WSL where the original paths are valid.
206203
204+
Uses single quotes to prevent shell expansion of variables and special
205+
characters. Single quotes in the original command are escaped using
206+
the standard shell technique: end quote, escaped single quote, start
207+
quote ('\"'\"').
208+
207209
Args:
208210
proxy_command: The original proxy command from WSL SSH config.
209211
210212
Returns:
211213
A Windows-compatible proxy command using wsl.exe.
212214
"""
213-
escaped_command = proxy_command.replace('"', '\\"')
214-
return f'wsl.exe bash -c "{escaped_command}"'
215+
# Escape single quotes for bash single-quoted string
216+
# 'foo'bar' -> 'foo'"'"'bar'
217+
escaped_command = proxy_command.replace("'", "'\"'\"'")
218+
return f"wsl.exe bash -c '{escaped_command}'"
215219

216220
@classmethod
217221
def _add_cluster_to_windows_ssh_config(
@@ -260,10 +264,12 @@ def _add_cluster_to_windows_ssh_config(
260264
# Copy SSH key to Windows filesystem to avoid UNC path permission
261265
# issues. Windows SSH rejects keys accessed via //wsl$/ paths
262266
# because it cannot verify Unix file permissions.
267+
# Use cluster name as key filename to ensure consistent naming
268+
# between add and remove operations.
263269
windows_ssh_keys_dir = os.path.join(windows_home, '.sky',
264270
'ssh-keys')
265271
os.makedirs(windows_ssh_keys_dir, exist_ok=True, mode=0o700)
266-
key_filename = os.path.basename(key_path)
272+
key_filename = f'{cluster_name}.pem'
267273
windows_key_dest = os.path.join(windows_ssh_keys_dir, key_filename)
268274
# Copy the key file (use copyfile to only copy content, not
269275
# metadata/permissions which fail across WSL->Windows boundary)
@@ -349,10 +355,10 @@ def _remove_cluster_from_windows_ssh_config(cls, cluster_name: str) -> None:
349355

350356
try:
351357
common_utils.remove_file_if_exists(cluster_config_path)
352-
# Also remove the copied SSH key
358+
# Also remove the copied SSH key (named as {cluster_name}.pem)
353359
windows_ssh_keys_dir = os.path.join(windows_home, '.sky',
354360
'ssh-keys')
355-
key_path = os.path.join(windows_ssh_keys_dir, f'{cluster_name}.key')
361+
key_path = os.path.join(windows_ssh_keys_dir, f'{cluster_name}.pem')
356362
common_utils.remove_file_if_exists(key_path)
357363
except (OSError, PermissionError):
358364
# Silently ignore errors - Windows SSH config is optional
Lines changed: 215 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,215 @@
1+
"""Unit tests for sky/utils/cluster_utils.py."""
2+
from unittest import mock
3+
4+
import pytest
5+
6+
from sky.utils import cluster_utils
7+
8+
9+
class TestConvertWindowsPathToWsl:
10+
"""Tests for _convert_windows_path_to_wsl()."""
11+
12+
def test_standard_windows_path(self):
13+
result = cluster_utils._convert_windows_path_to_wsl('C:\\Users\\test')
14+
assert result == '/mnt/c/Users/test'
15+
16+
def test_windows_path_with_forward_slashes(self):
17+
result = cluster_utils._convert_windows_path_to_wsl('C:/Users/test')
18+
assert result == '/mnt/c/Users/test'
19+
20+
def test_drive_letter_case_insensitive(self):
21+
result = cluster_utils._convert_windows_path_to_wsl('D:\\Data')
22+
assert result == '/mnt/d/Data'
23+
24+
def test_root_drive(self):
25+
result = cluster_utils._convert_windows_path_to_wsl('C:\\')
26+
assert result == '/mnt/c/'
27+
28+
def test_already_wsl_path(self):
29+
result = cluster_utils._convert_windows_path_to_wsl('/already/wsl/path')
30+
assert result == '/already/wsl/path'
31+
32+
def test_relative_path(self):
33+
result = cluster_utils._convert_windows_path_to_wsl('relative/path')
34+
assert result == 'relative/path'
35+
36+
def test_empty_string(self):
37+
result = cluster_utils._convert_windows_path_to_wsl('')
38+
assert result == ''
39+
40+
41+
class TestConvertWslPathToWindows:
42+
"""Tests for _convert_wsl_path_to_windows()."""
43+
44+
def test_standard_wsl_path(self):
45+
result = cluster_utils._convert_wsl_path_to_windows('/mnt/c/Users/test')
46+
assert result == 'C:/Users/test'
47+
48+
def test_different_drive(self):
49+
result = cluster_utils._convert_wsl_path_to_windows('/mnt/d/Data')
50+
assert result == 'D:/Data'
51+
52+
def test_root_drive(self):
53+
result = cluster_utils._convert_wsl_path_to_windows('/mnt/c/')
54+
assert result == 'C:/'
55+
56+
def test_non_mnt_path(self):
57+
result = cluster_utils._convert_wsl_path_to_windows('/home/user')
58+
assert result == '/home/user'
59+
60+
def test_short_mnt_path(self):
61+
result = cluster_utils._convert_wsl_path_to_windows('/mnt/')
62+
assert result == '/mnt/'
63+
64+
def test_empty_string(self):
65+
result = cluster_utils._convert_wsl_path_to_windows('')
66+
assert result == ''
67+
68+
69+
class TestConvertProxyCommandForWindows:
70+
"""Tests for SSHConfigHelper._convert_proxy_command_for_windows()."""
71+
72+
def test_simple_proxy_command(self):
73+
cmd = 'ssh -W %h:%p user@host'
74+
result = cluster_utils.SSHConfigHelper._convert_proxy_command_for_windows(
75+
cmd)
76+
assert result == "wsl.exe bash -c 'ssh -W %h:%p user@host'"
77+
78+
def test_proxy_command_with_double_quotes(self):
79+
cmd = 'ssh -o "StrictHostKeyChecking=no" -W %h:%p user@host'
80+
result = cluster_utils.SSHConfigHelper._convert_proxy_command_for_windows(
81+
cmd)
82+
assert result == (
83+
"wsl.exe bash -c 'ssh -o \"StrictHostKeyChecking=no\" -W %h:%p user@host'"
84+
)
85+
86+
def test_proxy_command_with_single_quotes(self):
87+
cmd = "ssh -o 'StrictHostKeyChecking=no' -W %h:%p user@host"
88+
result = cluster_utils.SSHConfigHelper._convert_proxy_command_for_windows(
89+
cmd)
90+
# Single quotes should be escaped with '"'"'
91+
assert result == (
92+
"wsl.exe bash -c 'ssh -o '\"'\"'StrictHostKeyChecking=no'\"'\"' -W %h:%p user@host'"
93+
)
94+
95+
def test_proxy_command_with_dollar_sign(self):
96+
# Dollar signs should not be expanded in single-quoted strings
97+
cmd = 'echo $HOME'
98+
result = cluster_utils.SSHConfigHelper._convert_proxy_command_for_windows(
99+
cmd)
100+
assert result == "wsl.exe bash -c 'echo $HOME'"
101+
102+
def test_proxy_command_with_backticks(self):
103+
# Backticks should not be expanded in single-quoted strings
104+
cmd = 'echo `hostname`'
105+
result = cluster_utils.SSHConfigHelper._convert_proxy_command_for_windows(
106+
cmd)
107+
assert result == "wsl.exe bash -c 'echo `hostname`'"
108+
109+
110+
class TestGetWslWindowsHome:
111+
"""Tests for get_wsl_windows_home()."""
112+
113+
def test_not_wsl_returns_none(self):
114+
# Clear the lru_cache to ensure fresh state
115+
cluster_utils.get_wsl_windows_home.cache_clear()
116+
with mock.patch('sky.utils.common_utils.is_wsl', return_value=False):
117+
result = cluster_utils.get_wsl_windows_home()
118+
assert result is None
119+
120+
def test_wsl_with_userprofile_env(self):
121+
cluster_utils.get_wsl_windows_home.cache_clear()
122+
with mock.patch('sky.utils.common_utils.is_wsl', return_value=True):
123+
with mock.patch.dict('os.environ',
124+
{'USERPROFILE': 'C:\\Users\\testuser'}):
125+
with mock.patch('os.path.isdir', return_value=True):
126+
result = cluster_utils.get_wsl_windows_home()
127+
assert result == '/mnt/c/Users/testuser'
128+
129+
def test_wsl_without_userprofile_uses_cmd(self):
130+
cluster_utils.get_wsl_windows_home.cache_clear()
131+
with mock.patch('sky.utils.common_utils.is_wsl', return_value=True):
132+
with mock.patch.dict('os.environ', {}, clear=True):
133+
with mock.patch(
134+
'sky.utils.cluster_utils._get_windows_userprofile_via_cmd',
135+
return_value='D:\\Users\\cmduser'):
136+
with mock.patch('os.path.isdir', return_value=True):
137+
result = cluster_utils.get_wsl_windows_home()
138+
assert result == '/mnt/d/Users/cmduser'
139+
140+
def test_wsl_with_invalid_home_returns_none(self):
141+
cluster_utils.get_wsl_windows_home.cache_clear()
142+
with mock.patch('sky.utils.common_utils.is_wsl', return_value=True):
143+
with mock.patch.dict('os.environ',
144+
{'USERPROFILE': 'C:\\Users\\testuser'}):
145+
with mock.patch('os.path.isdir', return_value=False):
146+
result = cluster_utils.get_wsl_windows_home()
147+
assert result is None
148+
149+
def test_caching_behavior(self):
150+
"""Test that the function result is cached."""
151+
cluster_utils.get_wsl_windows_home.cache_clear()
152+
call_count = 0
153+
154+
def mock_is_wsl():
155+
nonlocal call_count
156+
call_count += 1
157+
return True
158+
159+
with mock.patch('sky.utils.common_utils.is_wsl',
160+
side_effect=mock_is_wsl):
161+
with mock.patch.dict('os.environ',
162+
{'USERPROFILE': 'C:\\Users\\cached'}):
163+
with mock.patch('os.path.isdir', return_value=True):
164+
# Call multiple times
165+
result1 = cluster_utils.get_wsl_windows_home()
166+
result2 = cluster_utils.get_wsl_windows_home()
167+
result3 = cluster_utils.get_wsl_windows_home()
168+
169+
# Should only call is_wsl once due to caching
170+
assert call_count == 1
171+
assert result1 == result2 == result3 == '/mnt/c/Users/cached'
172+
173+
174+
class TestGetWindowsUserprofileViaCmd:
175+
"""Tests for _get_windows_userprofile_via_cmd()."""
176+
177+
def test_successful_cmd_call(self):
178+
mock_result = mock.Mock()
179+
mock_result.returncode = 0
180+
mock_result.stdout = 'C:\\Users\\testuser\n'
181+
182+
with mock.patch('subprocess.run', return_value=mock_result):
183+
result = cluster_utils._get_windows_userprofile_via_cmd()
184+
assert result == 'C:\\Users\\testuser'
185+
186+
def test_failed_cmd_call(self):
187+
mock_result = mock.Mock()
188+
mock_result.returncode = 1
189+
mock_result.stdout = ''
190+
191+
with mock.patch('subprocess.run', return_value=mock_result):
192+
result = cluster_utils._get_windows_userprofile_via_cmd()
193+
assert result is None
194+
195+
def test_unexpanded_variable(self):
196+
mock_result = mock.Mock()
197+
mock_result.returncode = 0
198+
mock_result.stdout = '%USERPROFILE%\n'
199+
200+
with mock.patch('subprocess.run', return_value=mock_result):
201+
result = cluster_utils._get_windows_userprofile_via_cmd()
202+
assert result is None
203+
204+
def test_timeout_exception(self):
205+
import subprocess
206+
with mock.patch('subprocess.run',
207+
side_effect=subprocess.TimeoutExpired(cmd='cmd.exe',
208+
timeout=5)):
209+
result = cluster_utils._get_windows_userprofile_via_cmd()
210+
assert result is None
211+
212+
def test_file_not_found(self):
213+
with mock.patch('subprocess.run', side_effect=FileNotFoundError()):
214+
result = cluster_utils._get_windows_userprofile_via_cmd()
215+
assert result is None

0 commit comments

Comments
 (0)