dosh/src/client.rs

use crate::config::{
    ClientConfig, HostConfig, expand_tilde, load_client_config, load_hosts_config,
};
use crate::crypto;
use crate::native::{
    self, EnvVar, ForwardingKind, ForwardingRequest, KnownHostStatus, NativeClientHello,
    derive_native_session_key, generate_native_ephemeral, sign_user_auth_with_private_key,
    supported_user_key_algorithms, trust_host, verify_known_host, verify_server_hello,
};
use crate::protocol::{
    self, AttachReject, CLIENT_TO_SERVER, NativeAuthOkBody, NativeClientHelloBody,
    NativeServerHelloBody, NativeUserAuthBody, PacketKind, SERVER_TO_CLIENT,
};
use crate::ssh_agent;
use crate::transport::{DoshTransport, SessionRole, SessionTransportConfig, TransportConfig};
use anyhow::{Context, Result, anyhow, bail};
use std::net::{SocketAddr, ToSocketAddrs};
use std::path::PathBuf;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::net::UdpSocket;

#[derive(Debug, Clone)]
pub struct DoshClient {
    config: ClientConfig,
    hosts: crate::config::HostsConfig,
}

impl DoshClient {
    pub fn load() -> Result<Self> {
        Ok(Self {
            config: load_client_config(None)?,
            hosts: load_hosts_config(None)?,
        })
    }

    pub fn load_from_paths(
        client_config: Option<PathBuf>,
        hosts_config: Option<PathBuf>,
    ) -> Result<Self> {
        Ok(Self {
            config: load_client_config(client_config)?,
            hosts: load_hosts_config(hosts_config)?,
        })
    }

    pub fn with_config(config: ClientConfig, hosts: crate::config::HostsConfig) -> Self {
        Self { config, hosts }
    }

    pub fn connect(&self, host: impl Into<String>) -> DoshClientBuilder {
        DoshClientBuilder::new(self.clone(), host.into())
    }
}

#[derive(Debug, Clone)]
pub struct DoshClientBuilder {
    client: DoshClient,
    host: String,
    services: Vec<String>,
    identity_files: Vec<PathBuf>,
    session: Option<String>,
    user: Option<String>,
    udp_host: Option<String>,
    udp_port: Option<u16>,
    trust_on_first_use: Option<bool>,
    use_ssh_agent: Option<bool>,
    timeout: Option<Duration>,
    env: Vec<EnvVar>,
}

impl DoshClientBuilder {
    pub fn new(client: DoshClient, host: String) -> Self {
        Self {
            client,
            host,
            services: Vec::new(),
            identity_files: Vec::new(),
            session: None,
            user: None,
            udp_host: None,
            udp_port: None,
            trust_on_first_use: None,
            use_ssh_agent: None,
            timeout: None,
            env: Vec::new(),
        }
    }

    pub fn service(mut self, name: impl Into<String>) -> Self {
        self.services.push(name.into());
        self
    }

    pub fn services(mut self, names: impl IntoIterator<Item = impl Into<String>>) -> Self {
        self.services.extend(names.into_iter().map(Into::into));
        self
    }

    pub fn identity_file(mut self, path: impl Into<PathBuf>) -> Self {
        self.identity_files.push(path.into());
        self
    }

    pub fn session(mut self, session: impl Into<String>) -> Self {
        self.session = Some(session.into());
        self
    }

    pub fn user(mut self, user: impl Into<String>) -> Self {
        self.user = Some(user.into());
        self
    }

    pub fn udp_host(mut self, host: impl Into<String>) -> Self {
        self.udp_host = Some(host.into());
        self
    }

    pub fn udp_port(mut self, port: u16) -> Self {
        self.udp_port = Some(port);
        self
    }

    pub fn trust_on_first_use(mut self, trust: bool) -> Self {
        self.trust_on_first_use = Some(trust);
        self
    }

    pub fn use_ssh_agent(mut self, value: bool) -> Self {
        self.use_ssh_agent = Some(value);
        self
    }

    pub fn timeout(mut self, timeout: Duration) -> Self {
        self.timeout = Some(timeout);
        self
    }

    pub fn env(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
        self.env.push(EnvVar {
            name: name.into(),
            value: value.into(),
        });
        self
    }

    pub async fn connect(self) -> Result<ConnectedDoshClient> {
        let host_config = self
            .client
            .hosts
            .hosts
            .get(&self.host)
            .cloned()
            .unwrap_or_default();
        let raw_server = host_config.ssh.clone().unwrap_or_else(|| self.host.clone());
        let requested_user = self
            .user
            .clone()
            .or_else(|| host_config.user.clone())
            .or_else(|| user_from_destination(&raw_server))
            .or_else(|| std::env::var("USER").ok())
            .unwrap_or_else(|| "unknown".to_string());
        let udp_host = self
            .udp_host
            .clone()
            .or_else(|| host_config.dosh_host.clone())
            .or_else(|| self.client.config.dosh_host.clone())
            .unwrap_or_else(|| destination_host(&raw_server));
        let udp_port = self
            .udp_port
            .or(host_config.port)
            .unwrap_or(self.client.config.dosh_port);
        let peer_addr = resolve_addr(&udp_host, udp_port)?;
        let timeout = self.timeout.unwrap_or_else(|| {
            Duration::from_millis(self.client.config.native_auth_timeout_ms.max(1))
        });
        let session = self.session.unwrap_or_else(default_sdk_session);
        let requested_forwardings = self
            .services
            .iter()
            .map(|service| {
                Ok(ForwardingRequest {
                    kind: ForwardingKind::Local,
                    bind_host: None,
                    listen_port: 0,
                    target_host: Some(crate::transport::service_target(service)?),
                    target_port: Some(0),
                })
            })
            .collect::<Result<Vec<_>>>()?;
        let socket = UdpSocket::bind("0.0.0.0:0").await?;
        let (client_secret, client_public) = generate_native_ephemeral();
        let hello = NativeClientHello {
            protocol_version: native::NATIVE_PROTOCOL_VERSION,
            client_random: crypto::random_32(),
            client_ephemeral_public: client_public,
            requested_host: self.host.clone(),
            requested_user,
            requested_session: session.clone(),
            requested_mode: "forward-only".to_string(),
            terminal_size: (80, 24),
            supported_aead: vec!["chacha20poly1305".to_string()],
            supported_user_key_algorithms: supported_user_key_algorithms(),
            cached_host_key_fingerprint: None,
            attach_ticket_envelope: None,
            requested_env: self.env,
        };
        let packet = protocol::encode_plain(
            PacketKind::NativeClientHello,
            [0u8; 16],
            1,
            0,
            &protocol::to_body(&NativeClientHelloBody {
                hello: hello.clone(),
            })?,
        )?;
        socket.send_to(&packet, peer_addr).await?;

        let mut buf = vec![0u8; 65535];
        let (n, _) = tokio::time::timeout(timeout, socket.recv_from(&mut buf)).await??;
        let packet = protocol::decode(&buf[..n])?;
        if packet.header.kind != PacketKind::NativeServerHello {
            if packet.header.kind == PacketKind::AttachReject {
                let reject: AttachReject = protocol::from_body(&packet.body)?;
                bail!("native auth rejected: {}", reject.reason);
            }
            bail!("native auth received unexpected server response");
        }
        let server_hello: NativeServerHelloBody = protocol::from_body(&packet.body)?;
        verify_server_hello(&hello, &server_hello.hello)?;
        verify_or_trust_host(
            &self.client.config,
            &self.host,
            &server_hello.hello.host_key,
            self.trust_on_first_use,
        )?;

        let session_key = derive_native_session_key(
            &client_secret,
            server_hello.hello.server_ephemeral_public,
            &hello,
            &server_hello.hello,
        )?;
        let auth = sign_auth(
            &self.client.config,
            &host_config,
            &hello,
            &server_hello.hello,
            requested_forwardings,
            self.identity_files,
            self.use_ssh_agent,
        )?;
        let mut pending_id = [0u8; 16];
        pending_id.copy_from_slice(&server_hello.hello.auth_challenge[..16]);
        let auth_packet = protocol::encode_encrypted(
            PacketKind::NativeUserAuth,
            pending_id,
            2,
            1,
            &session_key,
            CLIENT_TO_SERVER,
            &protocol::to_body(&NativeUserAuthBody { auth })?,
        )?;
        socket.send_to(&auth_packet, peer_addr).await?;
        let (n, _) = tokio::time::timeout(timeout, socket.recv_from(&mut buf)).await??;
        let packet = protocol::decode(&buf[..n])?;
        if packet.header.kind != PacketKind::NativeAuthOk {
            if packet.header.kind == PacketKind::AttachReject {
                let reject: AttachReject = protocol::from_body(&packet.body)?;
                bail!("native auth rejected: {}", reject.reason);
            }
            bail!("native auth received unexpected auth response");
        }
        let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT)?;
        let ok: NativeAuthOkBody = protocol::from_body(&plain)?;
        let transport = DoshTransport::new_owned(
            socket,
            SessionTransportConfig {
                role: SessionRole::Client,
                conn_id: ok.ok.client_id,
                session_key: ok.ok.session_key,
                peer_addr,
                initial_send_seq: 2,
                initial_ack: ok.ok.initial_seq,
                stream: TransportConfig::default(),
            },
        );
        Ok(ConnectedDoshClient {
            host: self.host,
            session: ok.ok.session,
            transport,
        })
    }
}

pub struct ConnectedDoshClient {
    pub host: String,
    pub session: String,
    pub transport: DoshTransport,
}

impl ConnectedDoshClient {
    pub fn into_transport(self) -> DoshTransport {
        self.transport
    }
}

fn verify_or_trust_host(
    config: &ClientConfig,
    host: &str,
    host_key: &native::HostPublicKey,
    trust_override: Option<bool>,
) -> Result<()> {
    let known_hosts = expand_tilde(&config.known_hosts);
    match verify_known_host(&known_hosts, host, host_key)? {
        KnownHostStatus::Trusted => Ok(()),
        KnownHostStatus::Unknown if trust_override.unwrap_or(config.trust_on_first_use) => {
            trust_host(&known_hosts, host, host_key, "sdk-tofu", false)?;
            Ok(())
        }
        KnownHostStatus::Unknown => Err(anyhow!(
            "Dosh host key for {host} is not trusted; run `dosh trust {host}` first or enable trust_on_first_use"
        )),
        KnownHostStatus::Mismatch { expected, actual } => Err(anyhow!(
            "Dosh host key mismatch for {host}: expected {expected}, got {actual}"
        )),
    }
}

fn sign_auth(
    config: &ClientConfig,
    host_config: &HostConfig,
    hello: &NativeClientHello,
    server_hello: &native::NativeServerHello,
    requested_forwardings: Vec<ForwardingRequest>,
    explicit_identity_files: Vec<PathBuf>,
    use_ssh_agent: Option<bool>,
) -> Result<native::NativeUserAuth> {
    let use_agent = use_ssh_agent.unwrap_or(config.use_ssh_agent);
    let mut errors = Vec::new();
    if use_agent {
        match ssh_agent::sign_user_auth_with_agent(
            hello,
            server_hello,
            requested_forwardings.clone(),
        ) {
            Ok(auth) => return Ok(auth),
            Err(err) => errors.push(format!("ssh-agent: {err:#}")),
        }
    }

    let mut paths = explicit_identity_files;
    if paths.is_empty() {
        paths.extend(config.identity_files.iter().map(|path| expand_tilde(path)));
    }
    if paths.is_empty() {
        paths.extend(default_identity_paths());
    }
    for path in paths {
        match native::load_native_identity(&path).and_then(|identity| {
            sign_user_auth_with_private_key(
                &identity,
                hello,
                server_hello,
                requested_forwardings.clone(),
            )
        }) {
            Ok(auth) => return Ok(auth),
            Err(err) => errors.push(format!("{}: {err:#}", path.display())),
        }
    }
    let _ = host_config;
    Err(anyhow!(
        "native auth found no usable identity: {}",
        errors.join("; ")
    ))
}

fn default_identity_paths() -> Vec<PathBuf> {
    ["~/.ssh/id_ed25519", "~/.ssh/id_ecdsa", "~/.ssh/id_rsa"]
        .into_iter()
        .map(expand_tilde)
        .collect()
}

fn resolve_addr(host: &str, port: u16) -> Result<SocketAddr> {
    (host, port)
        .to_socket_addrs()
        .with_context(|| format!("resolve UDP target {host}:{port}"))?
        .next()
        .ok_or_else(|| anyhow!("no UDP address resolved for {host}:{port}"))
}

fn destination_host(destination: &str) -> String {
    destination
        .rsplit('@')
        .next()
        .unwrap_or(destination)
        .split(':')
        .next()
        .unwrap_or(destination)
        .to_string()
}

fn user_from_destination(destination: &str) -> Option<String> {
    destination
        .rsplit_once('@')
        .map(|(user, _)| user.to_string())
        .filter(|user| !user.is_empty())
}

fn default_sdk_session() -> String {
    let millis = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis();
    format!("sdk-{millis}-{}", std::process::id())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parses_user_and_host_from_destination() {
        assert_eq!(
            user_from_destination("palav@example.com").as_deref(),
            Some("palav")
        );
        assert_eq!(destination_host("palav@example.com"), "example.com");
        assert_eq!(destination_host("example.com:2222"), "example.com");
    }

    #[test]
    fn default_identity_paths_are_expanded() {
        assert!(
            default_identity_paths()
                .iter()
                .any(|path| path.ends_with(".ssh/id_ed25519"))
        );
    }
}