Remote Dev Plugin Part 3 - Enable remote ssh connection

[sfc-gh-zzhu]

Pre-review checklist

• I've confirmed that instructions included in README.md are still correct after my changes in the codebase.
• I've added or updated automated unit tests to verify correctness of my new code.
• I've added or updated integration tests to verify correctness of my new code.
• I've confirmed that my changes are working by executing CLI's commands manually on MacOS.
• I've confirmed that my changes are working by executing CLI's commands manually on Windows.
• I've confirmed that my changes are up-to-date with the target branch.
• I've described my changes in the release notes.
• I've described my changes in the section below.

Changes description

This PR adds SSH support to the remote development environment feature, allowing users to securely connect to their remote environments via SSH. New command options to snow remote start includes:

• --ssh flag to set up SSH configuration for connecting to remote environments
• --no-ssh-key flag to use token-only authentication when SSH key generation is not desired

Users can now easily connect to their remote environments using standard SSH clients with the command: snow remote start myproject --ssh

[sfc-gh-zzhu]

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

• Add default block volume for persistent storage #2660
• Address feedbacks #2636
• Remote Dev Plugin Part 4 - bug fix and feature improvement #2569
• Remote Dev Plugin Part 3 - Enable remote ssh connection #2547 👈 (View in Graphite)
• Remote Dev Plugin Part 2 - Add basic remote commands #2546
• Remote Dev Plugin Part 1 - Register the remote dev plugin and add utils #2545
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[sfc-gh-dhung]

Do we have a way of testing this?

[sfc-gh-zzhu]

Currently no - I plan to explore e2e test later which may help the testing.

[sfc-gh-dhung]

Is it possible the endpoint uses a different protocol?

[sfc-gh-zzhu]

In our implementation it will always be wss

[sfc-gh-dhung]

Is this a fatal error, i.e. should we throw here?

[sfc-gh-zzhu]

Change to use warning instead of normal cc.step.

[graphite-app]

KeyboardInterrupt handling bug: The time.sleep(1) call in _wait_with_countdown() will not immediately respond to KeyboardInterrupt. Python's time.sleep() can delay interrupt handling, making the SSH session management unresponsive to Ctrl+C. This should use signal handling or shorter sleep intervals with interrupt checks.

Suggested change

	while time.time() < end_time:
	remaining = int(end_time - time.time())
	if remaining > 0 and remaining % SSH_COUNTDOWN_INTERVAL == 0:
	log.debug(
	"⏳ Next refresh in %d seconds... (Press Ctrl+C to stop)", remaining
	)
	time.sleep(1)
	while time.time() < end_time:
	remaining = int(end_time - time.time())
	if remaining > 0 and remaining % SSH_COUNTDOWN_INTERVAL == 0:
	log.debug(
	"⏳ Next refresh in %d seconds... (Press Ctrl+C to stop)", remaining
	)
	# Use shorter sleep intervals to be more responsive to keyboard interrupts
	try:
	time.sleep(0.1)
	except KeyboardInterrupt:
	raise

Spotted by Diamond



Is this helpful? React 👍 or 👎 to let us know.

[graphite-app]

The connection created in _get_fresh_token() is never explicitly closed, which could lead to connection pool exhaustion during long-running SSH sessions. While the comment mentions allowing "session to expire naturally," the connection object itself should still be properly closed in a finally block to ensure resource cleanup even when exceptions occur. Consider adding:

try:
    # existing code...
except Exception as e:
    # existing code...
finally:
    if fresh_connection and not fresh_connection.is_closed():
        fresh_connection.close()

This ensures proper connection cleanup regardless of execution path.

Suggested change

	fresh_connection = None
	try:
	log.debug("Creating fresh connection for SSH token...")
	current_context = get_cli_context().connection_context
	# Create connection with natural session expiration for SSH token refresh
	fresh_connection = connect_to_snowflake(
	connection_name=current_context.connection_name,
	temporary_connection=current_context.temporary_connection,
	# Allow session to expire naturally - don't keep it alive artificially
	using_session_keep_alive=False,
	)
	fresh_connection.cursor().execute(
	"ALTER SESSION SET python_connector_query_result_format = 'JSON'"
	)
	token = fresh_connection.rest.token
	if token:
	log.debug("Fresh token obtained successfully")
	return token
	else:
	log.debug("No token available from fresh connection")
	return None
	except Exception as e:
	log.debug("Failed to create fresh connection: %s", str(e))
	return None
	fresh_connection = None
	try:
	log.debug("Creating fresh connection for SSH token...")
	current_context = get_cli_context().connection_context
	# Create connection with natural session expiration for SSH token refresh
	fresh_connection = connect_to_snowflake(
	connection_name=current_context.connection_name,
	temporary_connection=current_context.temporary_connection,
	# Allow session to expire naturally - don't keep it alive artificially
	using_session_keep_alive=False,
	)
	fresh_connection.cursor().execute(
	"ALTER SESSION SET python_connector_query_result_format = 'JSON'"
	)
	token = fresh_connection.rest.token
	if token:
	log.debug("Fresh token obtained successfully")
	return token
	else:
	log.debug("No token available from fresh connection")
	return None
	except Exception as e:
	log.debug("Failed to create fresh connection: %s", str(e))
	return None
	finally:
	if fresh_connection and not fresh_connection.is_closed():
	fresh_connection.close()

Spotted by Diamond



Is this helpful? React 👍 or 👎 to let us know.

[graphite-app]

Security Concern: Potential Command Injection Risk

The current token escaping approach creates a potential security vulnerability. While shlex.quote() is used to escape the token, it's then embedded in a complex string with additional manual escaping (\\") which could compromise the escaping logic.

Consider refactoring this to use a more robust approach:

# Instead of:
f'ProxyCommand {escaped_websocat_path} --binary wss://{hostname}/ -H "Authorization: Snowflake Token=\\"{escaped_token}\\""'

# Consider a cleaner approach that doesn't mix escaping methods:
f'ProxyCommand {escaped_websocat_path} --binary wss://{hostname}/ -H "Authorization: Snowflake Token={escaped_token}"'

This would ensure the token remains properly contained within the shell command structure without risking escape sequence conflicts.

Suggested change

	# Properly escape values to prevent command injection
	escaped_websocat_path = shlex.quote(websocat_path)
	# For the token, we need to escape it for shell safety but also ensure it's properly quoted
	# within the Authorization header. We escape the token and then add quotes around it.
	escaped_token = shlex.quote(token)
	config_lines = [
	f"Host {service_name}",
	f" HostName {hostname}",
	f" Port {SSH_DEFAULT_PORT}",
	f" User {SSH_DEFAULT_USER}",
	f' ProxyCommand {escaped_websocat_path} --binary wss://{hostname}/ -H "Authorization: Snowflake Token=\\"{escaped_token}\\""',
	# Properly escape values to prevent command injection
	escaped_websocat_path = shlex.quote(websocat_path)
	# For the token, we need to escape it for shell safety
	escaped_token = shlex.quote(token)
	config_lines = [
	f"Host {service_name}",
	f" HostName {hostname}",
	f" Port {SSH_DEFAULT_PORT}",
	f" User {SSH_DEFAULT_USER}",
	f' ProxyCommand {escaped_websocat_path} --binary wss://{hostname}/ -H "Authorization: Snowflake Token={escaped_token}"',

Spotted by Diamond



Is this helpful? React 👍 or 👎 to let us know.

[graphite-app]

Consider adding a maximum retry limit or exponential backoff to the while True loop. Currently, if token retrieval consistently fails, the loop will continue indefinitely with fixed 30-second intervals, potentially consuming resources unnecessarily. A retry counter or increasing backoff periods would gracefully handle persistent authentication failures.

Suggested change

	while True:
	log.debug("Token refresh cycle #%d", token_refresh_count + 1)
	token = self._get_fresh_token()
	if not token:
	cc.step("❌ Unable to get token. Retrying in 30 seconds...")
	self._wait_with_shutdown_check(SSH_RETRY_INTERVAL)
	continue
	# Update SSH configuration
	auth_method = "SSH key" if private_key_path else "token-only"
	log.debug("Updating SSH config with %s authentication", auth_method)
	setup_ssh_config_with_token(
	service_name, ssh_hostname, token, private_key_path
	)
	token_refresh_count += 1
	cc.step(f"SSH configuration updated (refresh #{token_refresh_count})")
	# Wait for next refresh with countdown
	self._wait_with_countdown(refresh_interval)
	max_consecutive_failures = 5
	consecutive_failures = 0
	base_retry_interval = SSH_RETRY_INTERVAL # 30 seconds
	while True:
	log.debug("Token refresh cycle #%d", token_refresh_count + 1)
	token = self._get_fresh_token()
	if not token:
	consecutive_failures += 1
	if consecutive_failures > max_consecutive_failures:
	cc.step(f"❌ Failed to get token after {max_consecutive_failures} consecutive attempts. Exiting...")
	break
	retry_interval = base_retry_interval * (2 ** (consecutive_failures - 1)) # Exponential backoff
	cc.step(f"❌ Unable to get token. Retrying in {retry_interval} seconds... (attempt {consecutive_failures}/{max_consecutive_failures})")
	self._wait_with_shutdown_check(retry_interval)
	continue
	# Reset failure counter on success
	consecutive_failures = 0
	# Update SSH configuration
	auth_method = "SSH key" if private_key_path else "token-only"
	log.debug("Updating SSH config with %s authentication", auth_method)
	setup_ssh_config_with_token(
	service_name, ssh_hostname, token, private_key_path
	)
	token_refresh_count += 1
	cc.step(f"SSH configuration updated (refresh #{token_refresh_count})")
	# Wait for next refresh with countdown
	self._wait_with_countdown(refresh_interval)

Spotted by Diamond



Is this helpful? React 👍 or 👎 to let us know.

[graphite-app]

The _setup_ssh_key method should handle potential failures from generate_ssh_key_pair(). Currently, if key generation fails, the method will propagate None to service creation which could cause unexpected behavior. Consider adding error handling to either:

1. Catch exceptions from generate_ssh_key_pair() and log appropriate errors
2. Verify ssh_public_key is not None after generation
3. Raise a specific exception with a clear message if key generation fails

This would prevent silent failures and provide better diagnostics when SSH setup doesn't work as expected.

Suggested change

	def _setup_ssh_key(self, service_name: str, generate_key: bool) -> Optional[str]:
	"""Set up SSH key for a service and return the public key.
	Args:
	service_name: Name of the service
	generate_key: Whether to generate or use SSH keys for the service.
	If False, no SSH key operations are performed.
	Returns:
	SSH public key content if generate_key is True and either:
	- An existing SSH key pair is found for the service, or
	- A new SSH key pair is successfully generated
	None if generate_key is False (no SSH key operations performed)
	"""
	if not generate_key:
	return None
	ssh_key_result = get_existing_ssh_key(service_name)
	if ssh_key_result:
	_, ssh_public_key = ssh_key_result
	log.debug("Using existing SSH key pair for service %s", service_name)
	return ssh_public_key
	else:
	log.debug("Generating SSH key pair for service %s", service_name)
	_, ssh_public_key = generate_ssh_key_pair(service_name)
	return ssh_public_key
	def _setup_ssh_key(self, service_name: str, generate_key: bool) -> Optional[str]:
	"""Set up SSH key for a service and return the public key.
	Args:
	service_name: Name of the service
	generate_key: Whether to generate or use SSH keys for the service.
	If False, no SSH key operations are performed.
	Returns:
	SSH public key content if generate_key is True and either:
	- An existing SSH key pair is found for the service, or
	- A new SSH key pair is successfully generated
	None if generate_key is False (no SSH key operations performed)
	Raises:
	SnowflakeSSHKeyError: If SSH key generation fails
	"""
	if not generate_key:
	return None
	ssh_key_result = get_existing_ssh_key(service_name)
	if ssh_key_result:
	_, ssh_public_key = ssh_key_result
	log.debug("Using existing SSH key pair for service %s", service_name)
	return ssh_public_key
	else:
	log.debug("Generating SSH key pair for service %s", service_name)
	try:
	_, ssh_public_key = generate_ssh_key_pair(service_name)
	if not ssh_public_key:
	raise SnowflakeSSHKeyError(f"Failed to generate SSH public key for service {service_name}")
	return ssh_public_key
	except Exception as e:
	log.error("SSH key generation failed for service %s: %s", service_name, str(e))
	raise SnowflakeSSHKeyError(
	f"Failed to generate SSH key pair for service {service_name}: {str(e)}"
	) from e

Spotted by Diamond



Is this helpful? React 👍 or 👎 to let us know.