bssh 2.0.1

// Copyright 2025 Lablup Inc. and Jeongkyu Shin
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! SSH connection management and establishment.
//!
//! This module handles the low-level SSH connection establishment,
//! including address resolution, connection attempts, and initial handshake.

use russh::client::{Config, Handle, Handler};
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;
use std::{fmt::Debug, io};

use super::authentication::{AuthMethod, ServerCheckMethod};

/// Default keepalive interval in seconds.
/// Sends keepalive packets every 60 seconds to detect dead connections.
pub const DEFAULT_KEEPALIVE_INTERVAL: u64 = 60;

/// Default maximum keepalive attempts before considering connection dead.
/// With 60s interval and 3 max, connection failure is detected within 180s.
pub const DEFAULT_KEEPALIVE_MAX: usize = 3;

/// SSH connection configuration for keepalive and timeout settings.
///
/// This struct provides a centralized way to configure SSH connection
/// parameters, particularly for keepalive functionality which prevents
/// idle connections from being terminated by firewalls or NAT devices.
///
/// # Example
///
/// ```no_run
/// use bssh::ssh::tokio_client::SshConnectionConfig;
///
/// // Use defaults (60s interval, 3 max attempts)
/// let config = SshConnectionConfig::default();
///
/// // Custom configuration
/// let config = SshConnectionConfig::new()
///     .with_keepalive_interval(Some(30))
///     .with_keepalive_max(5);
///
/// // Disable keepalive
/// let config = SshConnectionConfig::new()
///     .with_keepalive_interval(None);
/// ```
#[derive(Debug, Clone)]
pub struct SshConnectionConfig {
    /// Interval in seconds between keepalive packets.
    /// None disables keepalive.
    /// Default: 60 seconds
    pub keepalive_interval: Option<u64>,

    /// Maximum number of keepalive packets to send without response
    /// before considering the connection dead.
    /// Default: 3
    pub keepalive_max: usize,
}

impl Default for SshConnectionConfig {
    fn default() -> Self {
        Self {
            keepalive_interval: Some(DEFAULT_KEEPALIVE_INTERVAL),
            keepalive_max: DEFAULT_KEEPALIVE_MAX,
        }
    }
}

impl SshConnectionConfig {
    /// Create a new configuration with default values.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the keepalive interval in seconds.
    /// Pass None to disable keepalive.
    #[must_use]
    pub fn with_keepalive_interval(mut self, interval: Option<u64>) -> Self {
        self.keepalive_interval = interval;
        self
    }

    /// Set the maximum number of keepalive attempts.
    #[must_use]
    pub fn with_keepalive_max(mut self, max: usize) -> Self {
        self.keepalive_max = max;
        self
    }

    /// Convert this configuration to a russh client Config.
    ///
    /// When keepalive is enabled, `inactivity_timeout` is set to `None` so the
    /// keepalive mechanism is the sole dead-peer detector. russh's default
    /// `inactivity_timeout` is 10 minutes and would otherwise tear down an
    /// otherwise-healthy idle session at that mark regardless of keepalive
    /// liveness. When keepalive is disabled, we preserve a generous
    /// inactivity timeout so truly dead sockets are still reaped.
    pub fn to_russh_config(&self) -> Config {
        let inactivity_timeout = if self.keepalive_interval.is_some() {
            None
        } else {
            Some(Duration::from_secs(3600))
        };
        Config {
            keepalive_interval: self.keepalive_interval.map(Duration::from_secs),
            keepalive_max: self.keepalive_max,
            inactivity_timeout,
            ..Default::default()
        }
    }

    /// Derive a TCP-level keepalive configuration from this SSH keepalive
    /// configuration. Returns `None` if SSH keepalive is disabled.
    ///
    /// TCP keepalive is a belt-and-suspenders mechanism: it lets the kernel
    /// detect a broken TCP path even when no application data is flowing and
    /// even if SSH-level keepalive replies are dropped by a middlebox.
    pub fn to_tcp_keepalive(&self) -> Option<socket2::TcpKeepalive> {
        let interval = self.keepalive_interval?;
        // Start probing after `interval` seconds of idleness, probe every
        // half-interval, up to keepalive_max retries.
        let probe_interval = (interval / 2).max(1);
        let ka = socket2::TcpKeepalive::new()
            .with_time(Duration::from_secs(interval))
            .with_interval(Duration::from_secs(probe_interval));
        #[cfg(any(
            target_os = "linux",
            target_os = "macos",
            target_os = "freebsd",
            target_os = "netbsd",
            target_os = "tvos",
            target_os = "watchos",
            target_os = "ios",
        ))]
        let ka = ka.with_retries(self.keepalive_max.max(1) as u32);
        Some(ka)
    }
}
use super::ToSocketAddrsWithHostname;

/// A ssh connection to a remote server.
///
/// After creating a `Client` by [`connect`]ing to a remote host,
/// use [`execute`] to send commands and receive results through the connections.
///
/// [`connect`]: Client::connect
/// [`execute`]: Client::execute
///
/// # Examples
///
/// ```no_run
/// use bssh::ssh::tokio_client::{Client, AuthMethod, ServerCheckMethod};
/// #[tokio::main]
/// async fn main() -> Result<(), bssh::ssh::tokio_client::Error> {
///     let mut client = Client::connect(
///         ("10.10.10.2", 22),
///         "root",
///         AuthMethod::with_password("root"),
///         ServerCheckMethod::NoCheck,
///     ).await?;
///
///     let result = client.execute("echo Hello SSH").await?;
///     assert_eq!(result.stdout, "Hello SSH\n");
///     assert_eq!(result.exit_status, 0);
///
///     Ok(())
/// }
#[derive(Clone)]
pub struct Client {
    pub(super) connection_handle: Arc<Handle<ClientHandler>>,
    pub(super) username: String,
    pub(super) address: SocketAddr,
    /// Public access to the SSH session for jump host operations
    #[allow(private_interfaces)]
    pub session: Arc<Handle<ClientHandler>>,
}

impl Client {
    /// Open a ssh connection to a remote host with default keepalive settings.
    ///
    /// `addr` is an address of the remote host. Anything which implements
    /// [`ToSocketAddrsWithHostname`] trait can be supplied for the address;
    /// ToSocketAddrsWithHostname reimplements all of [`ToSocketAddrs`];
    /// see this trait's documentation for concrete examples.
    ///
    /// If `addr` yields multiple addresses, `connect` will be attempted with
    /// each of the addresses until a connection is successful.
    /// Authentification is tried on the first successful connection and the whole
    /// process aborted if this fails.
    ///
    /// This method uses default keepalive settings (60s interval, 3 max attempts)
    /// to prevent idle connection timeouts.
    pub async fn connect(
        addr: impl ToSocketAddrsWithHostname,
        username: &str,
        auth: AuthMethod,
        server_check: ServerCheckMethod,
    ) -> Result<Self, super::Error> {
        Self::connect_with_ssh_config(
            addr,
            username,
            auth,
            server_check,
            &SshConnectionConfig::default(),
        )
        .await
    }

    /// Connect with custom SSH connection configuration.
    ///
    /// This method allows specifying keepalive settings and other connection
    /// parameters through [`SshConnectionConfig`].
    ///
    /// # Example
    ///
    /// ```no_run
    /// use bssh::ssh::tokio_client::{Client, AuthMethod, ServerCheckMethod, SshConnectionConfig};
    ///
    /// #[tokio::main]
    /// async fn main() -> Result<(), bssh::ssh::tokio_client::Error> {
    ///     let ssh_config = SshConnectionConfig::new()
    ///         .with_keepalive_interval(Some(30))
    ///         .with_keepalive_max(5);
    ///
    ///     let client = Client::connect_with_ssh_config(
    ///         ("example.com", 22),
    ///         "user",
    ///         AuthMethod::with_key_file("~/.ssh/id_rsa", None),
    ///         ServerCheckMethod::DefaultKnownHostsFile,
    ///         &ssh_config,
    ///     ).await?;
    ///
    ///     Ok(())
    /// }
    /// ```
    pub async fn connect_with_ssh_config(
        addr: impl ToSocketAddrsWithHostname,
        username: &str,
        auth: AuthMethod,
        server_check: ServerCheckMethod,
        ssh_config: &SshConnectionConfig,
    ) -> Result<Self, super::Error> {
        let config = ssh_config.to_russh_config();
        let tcp_keepalive = ssh_config.to_tcp_keepalive();
        Self::connect_with_config_inner(
            addr,
            username,
            auth,
            server_check,
            config,
            tcp_keepalive.as_ref(),
        )
        .await
    }

    /// Same as `connect`, but with the option to specify a non default
    /// [`russh::client::Config`].
    ///
    /// For most use cases, prefer [`connect_with_ssh_config`] which provides
    /// a higher-level API with sensible defaults.
    pub async fn connect_with_config(
        addr: impl ToSocketAddrsWithHostname,
        username: &str,
        auth: AuthMethod,
        server_check: ServerCheckMethod,
        config: Config,
    ) -> Result<Self, super::Error> {
        Self::connect_with_config_inner(addr, username, auth, server_check, config, None).await
    }

    async fn connect_with_config_inner(
        addr: impl ToSocketAddrsWithHostname,
        username: &str,
        auth: AuthMethod,
        server_check: ServerCheckMethod,
        config: Config,
        tcp_keepalive: Option<&socket2::TcpKeepalive>,
    ) -> Result<Self, super::Error> {
        let config = Arc::new(config);

        // Connection code inspired from std::net::TcpStream::connect and std::net::each_addr
        let socket_addrs = addr
            .to_socket_addrs()
            .map_err(super::Error::AddressInvalid)?;
        let mut connect_res: Result<
            (SocketAddr, russh::client::Handle<ClientHandler>),
            super::Error,
        > = Err(super::Error::AddressInvalid(io::Error::new(
            io::ErrorKind::InvalidInput,
            "could not resolve to any addresses",
        )));
        for socket_addr in socket_addrs {
            let handler = ClientHandler {
                hostname: addr.hostname(),
                host: socket_addr,
                server_check: server_check.clone(),
            };

            let stream = match tokio::net::TcpStream::connect(socket_addr).await {
                Ok(s) => s,
                Err(e) => {
                    connect_res = Err(super::Error::IoError(e));
                    continue;
                }
            };

            if let Some(ka) = tcp_keepalive {
                let sock_ref = socket2::SockRef::from(&stream);
                if let Err(e) = sock_ref.set_tcp_keepalive(ka) {
                    tracing::debug!(
                        "Failed to set TCP keepalive on socket to {}: {}",
                        socket_addr,
                        e
                    );
                }
            }

            match russh::client::connect_stream(config.clone(), stream, handler).await {
                Ok(h) => {
                    connect_res = Ok((socket_addr, h));
                    break;
                }
                Err(e) => connect_res = Err(e),
            }
        }
        let (address, mut handle) = connect_res?;
        let username = username.to_string();

        super::authentication::authenticate(&mut handle, &username, auth).await?;

        let connection_handle = Arc::new(handle);
        Ok(Self {
            connection_handle: connection_handle.clone(),
            username,
            address,
            session: connection_handle,
        })
    }

    /// Create a Client from an existing russh handle and address.
    ///
    /// This is used internally for jump host connections where we already have
    /// an authenticated russh handle from connect_stream.
    pub fn from_handle_and_address(
        handle: Arc<Handle<ClientHandler>>,
        username: String,
        address: SocketAddr,
    ) -> Self {
        Self {
            connection_handle: handle.clone(),
            username,
            address,
            session: handle,
        }
    }

    /// A debugging function to get the username this client is connected as.
    pub fn get_connection_username(&self) -> &String {
        &self.username
    }

    /// A debugging function to get the address this client is connected to.
    pub fn get_connection_address(&self) -> &SocketAddr {
        &self.address
    }

    /// Disconnect from the remote host.
    pub async fn disconnect(&self) -> Result<(), super::Error> {
        self.connection_handle
            .disconnect(russh::Disconnect::ByApplication, "", "")
            .await
            .map_err(super::Error::SshError)
    }

    /// Check if the connection is closed.
    pub fn is_closed(&self) -> bool {
        self.connection_handle.is_closed()
    }

    /// Request remote port forwarding (tcpip-forward) - Future Implementation Placeholder
    ///
    /// **TODO**: This method needs to be implemented once russh provides
    /// global request functionality or we find the appropriate API.
    ///
    /// This sends a global request to the SSH server to bind a port on the remote end
    /// and forward connections back to the client. This is used for remote port forwarding (-R).
    ///
    /// # Arguments
    /// * `bind_address` - Address to bind on the remote server (e.g., "localhost", "0.0.0.0")
    /// * `bind_port` - Port to bind on the remote server (0 to let server choose)
    ///
    /// # Returns
    /// The actual port number that was bound by the server (useful when bind_port is 0)
    pub async fn request_port_forward(
        &self,
        _bind_address: String,
        _bind_port: u32,
    ) -> Result<u32, super::Error> {
        // **TODO**: Implement actual tcpip-forward global request
        // For now, return an error indicating this is not yet implemented
        tracing::warn!("Remote port forwarding request not yet implemented - TODO");
        Err(super::Error::PortForwardingNotSupported)
    }

    /// Cancel remote port forwarding (cancel-tcpip-forward) - Future Implementation Placeholder
    ///
    /// **TODO**: This method needs to be implemented once russh provides
    /// global request functionality or we find the appropriate API.
    ///
    /// This sends a global request to cancel a previously established remote port forward.
    ///
    /// # Arguments
    /// * `bind_address` - Address that was bound on the remote server
    /// * `bind_port` - Port that was bound on the remote server
    pub async fn cancel_port_forward(
        &self,
        _bind_address: String,
        _bind_port: u32,
    ) -> Result<(), super::Error> {
        // **TODO**: Implement actual cancel-tcpip-forward global request
        // For now, return an error indicating this is not yet implemented
        tracing::warn!("Cancel remote port forwarding not yet implemented - TODO");
        Err(super::Error::PortForwardingNotSupported)
    }
}

impl Debug for Client {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Client")
            .field("username", &self.username)
            .field("address", &self.address)
            .field("connection_handle", &"Handle<ClientHandler>")
            .finish()
    }
}

/// SSH client handler for managing server key verification.
#[derive(Debug, Clone)]
pub struct ClientHandler {
    hostname: String,
    host: SocketAddr,
    server_check: ServerCheckMethod,
}

impl ClientHandler {
    /// Create a new client handler.
    pub fn new(hostname: String, host: SocketAddr, server_check: ServerCheckMethod) -> Self {
        Self {
            hostname,
            host,
            server_check,
        }
    }
}

impl Handler for ClientHandler {
    type Error = super::Error;

    async fn check_server_key(
        &mut self,
        server_public_key: &russh::keys::PublicKey,
    ) -> Result<bool, Self::Error> {
        match &self.server_check {
            ServerCheckMethod::NoCheck => Ok(true),
            ServerCheckMethod::PublicKey(key) => {
                let pk = russh::keys::parse_public_key_base64(key)
                    .map_err(|_| super::Error::ServerCheckFailed)?;

                Ok(pk == *server_public_key)
            }
            ServerCheckMethod::PublicKeyFile(key_file_name) => {
                let pk = russh::keys::load_public_key(key_file_name)
                    .map_err(|_| super::Error::ServerCheckFailed)?;

                Ok(pk == *server_public_key)
            }
            ServerCheckMethod::KnownHostsFile(known_hosts_path) => {
                let result = russh::keys::check_known_hosts_path(
                    &self.hostname,
                    self.host.port(),
                    server_public_key,
                    known_hosts_path,
                )
                .map_err(|_| super::Error::ServerCheckFailed)?;

                Ok(result)
            }
            ServerCheckMethod::DefaultKnownHostsFile => {
                let result = russh::keys::check_known_hosts(
                    &self.hostname,
                    self.host.port(),
                    server_public_key,
                )
                .map_err(|_| super::Error::ServerCheckFailed)?;

                Ok(result)
            }
        }
    }
}