path: root/src/client/transport.rs

//! Client transport: TCP, rustls, TOFU and the HELLO handshake.
//!
//! The handshake records the server's SPKI fingerprint but does not act on it;
//! the trust decision happens after the handshake and before any credentials
//! are sent, so a frontend can ask the player. A changed key is a hard error.

use std::fmt;
use std::io::{Read, Write};
use std::net::TcpStream;
use std::sync::Arc;
use std::sync::mpsc;
use std::thread;
use std::time::{Duration, Instant};

use parking_lot::Mutex;
use rusqlite::Connection;
use rustls::DigitallySignedStruct;
use rustls::client::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier};

use crate::client::database;
use crate::protocol::{
    FrameBuffer, Packet, parse_packet, request_line, write_frame as write_protocol_frame,
};
use crate::tofu::{TrustDecision, decide, spki_fingerprint};
use crate::version::{HelloReply, VERSION, parse_hello};

/// Reserved request ids for the transport handshake. Application requests
/// start at 1, so the two ranges cannot collide on one connection.
const HELLO_REQUEST_ID: u64 = u64::MAX;
const AUTH_REQUEST_ID: u64 = u64::MAX - 1;

/// How long to wait for TCP + TLS before giving up.
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
/// Read timeout once the stream is usable.
const READ_TIMEOUT: Duration = Duration::from_secs(10);

/// What a frontend shows when a server key is unknown.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TrustPrompt {
    pub addr: String,
    pub fingerprint: String,
}

/// A usable connection after trust, version and optional auto-login.
pub struct Connected {
    pub stream: TlsStream,
    pub user: Option<String>,
    /// Set when stored credentials were rejected; the session is still usable.
    pub auth_error: Option<String>,
    pub server_version: String,
    pub latest: String,
    pub url: String,
    /// Server Unix time reported by `HELLO`, milliseconds.
    pub server_time: u64,
    /// Round trip of the HELLO exchange.
    pub elapsed: Duration,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TransportError {
    Timeout,
    TrustDeclined,
    KeyChanged { expected: String, observed: String },
    VersionMismatch(String),
    Closed,
    Unexpected(String),
    Io(String),
}

impl fmt::Display for TransportError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            TransportError::Timeout => write!(f, "Connection timed out"),
            TransportError::TrustDeclined => write!(f, "Certificate was not trusted"),
            TransportError::KeyChanged { expected, observed } => write!(
                f,
                "HOST KEY CHANGED for this server!\n\
                 expected fingerprint: {expected}\n\
                 observed fingerprint: {observed}\n\
                 If this is intentional (server reinstall), forget this host to reset trust."
            ),
            TransportError::VersionMismatch(message) => write!(f, "{message}"),
            TransportError::Closed => write!(f, "Server closed the connection during HELLO"),
            TransportError::Unexpected(message) => write!(f, "{message}"),
            TransportError::Io(message) => write!(f, "Connection error: {message}"),
        }
    }
}

impl std::error::Error for TransportError {}

fn io_failure(error: std::io::Error) -> TransportError {
    TransportError::Io(error.to_string())
}

/// Connect, verify the key, check versions and try the stored login.
///
/// `trust` is called only for a server we have never seen; returning `false`
/// aborts before any application data is sent.
pub fn connect(
    addr: &str,
    conn: &Connection,
    unsafe_mode: bool,
    trust: &mut dyn FnMut(&TrustPrompt) -> bool,
) -> Result<Connected, TransportError> {
    let stored = database::get_fingerprint(conn, addr);
    let (tls, observed) = tls_connect(addr)?;
    let observed = observed.ok_or_else(|| {
        TransportError::Unexpected("the server presented no certificate fingerprint".to_string())
    })?;

    if !unsafe_mode {
        match decide(stored.as_deref(), &observed) {
            TrustDecision::Trusted => {}
            TrustDecision::TrustNew => {
                let prompt = TrustPrompt {
                    addr: addr.to_string(),
                    fingerprint: observed.clone(),
                };
                if !trust(&prompt) {
                    return Err(TransportError::TrustDeclined);
                }
                database::save_fingerprint(conn, addr, &observed)
                    .map_err(|error| TransportError::Io(error.to_string()))?;
            }
            TrustDecision::Mismatch { expected, actual } => {
                return Err(TransportError::KeyChanged {
                    expected,
                    observed: actual,
                });
            }
        }
    }

    let mut tls = tls;
    tls.set_read_timeout(Some(READ_TIMEOUT))
        .map_err(io_failure)?;

    let start = Instant::now();
    let response = exchange(&mut tls, HELLO_REQUEST_ID, &format!("HELLO {VERSION}"))?;
    let elapsed = start.elapsed();

    let response = response.trim();
    if let Some(message) = response.strip_prefix("ERROR") {
        return Err(TransportError::VersionMismatch(message.trim().to_string()));
    }
    let hello: HelloReply =
        parse_hello(response).ok_or_else(|| TransportError::Unexpected(response.to_string()))?;

    // Credentials are keyed by the pinned fingerprint, so any address that
    // presents this key can log in with the same stored account.
    let mut user = None;
    let mut auth_error = None;
    if let Some((username, password)) = database::get_credentials(conn, &observed) {
        let auth = exchange(
            &mut tls,
            AUTH_REQUEST_ID,
            &format!("AUTH {username} {password}"),
        )?;
        if auth.trim().starts_with("OK") {
            user = Some(username);
        } else {
            auth_error = Some(auth.trim().to_string());
        }
    }

    Ok(Connected {
        stream: tls,
        user,
        auth_error,
        server_version: hello.version,
        latest: hello.latest,
        url: hello.url,
        server_time: hello.time,
        elapsed,
    })
}

/// Send one request and wait for the reply with the same id, ignoring events
/// that arrive in between. Used only for the handshake.
fn exchange(stream: &mut TlsStream, id: u64, payload: &str) -> Result<String, TransportError> {
    stream
        .write_request(id, None, payload)
        .map_err(io_failure)?;
    loop {
        match stream.read_packet().map_err(io_failure)? {
            None => return Err(TransportError::Closed),
            Some(Packet::Reply { id: reply_id, body }) if reply_id == id => return Ok(body),
            Some(Packet::Reply { .. }) | Some(Packet::Event { .. }) => continue,
        }
    }
}

/// TCP + TLS on a thread so a blocked handshake can time out.
fn tls_connect(addr: &str) -> Result<(TlsStream, Option<String>), TransportError> {
    let (sender, receiver) = mpsc::channel();
    let addr = addr.to_string();
    let server_name = addr.split(':').next().unwrap_or("127.0.0.1").to_string();

    thread::spawn(move || {
        let outcome = (|| -> Result<(TlsStream, Option<String>), TransportError> {
            let tcp = TcpStream::connect(&addr).map_err(io_failure)?;
            let connector = TlsConnector::new(&server_name)
                .map_err(|error| TransportError::Io(error.to_string()))?;
            let stream = connector
                .connect(tcp)
                .map_err(|error| TransportError::Io(error.to_string()))?;
            Ok((stream, connector.observed_fingerprint()))
        })();
        let _ = sender.send(outcome);
    });

    match receiver.recv_timeout(CONNECT_TIMEOUT) {
        Ok(outcome) => outcome,
        Err(mpsc::RecvTimeoutError::Timeout) => Err(TransportError::Timeout),
        Err(mpsc::RecvTimeoutError::Disconnected) => Err(TransportError::Unexpected(
            "connection attempt failed".to_string(),
        )),
    }
}

struct TlsConnector {
    config: Arc<rustls::ClientConfig>,
    server_name: rustls::ServerName,
    observed: Arc<Mutex<Option<String>>>,
}

impl TlsConnector {
    fn new(server_name: &str) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
        let observed = Arc::new(Mutex::new(None));
        let verifier = Arc::new(RecordingVerifier {
            observed: Arc::clone(&observed),
        });
        let config = rustls::ClientConfig::builder()
            .with_safe_defaults()
            .with_custom_certificate_verifier(verifier)
            .with_no_client_auth();
        Ok(Self {
            config: Arc::new(config),
            server_name: server_name.try_into()?,
            observed,
        })
    }

    fn connect(
        &self,
        stream: TcpStream,
    ) -> Result<TlsStream, Box<dyn std::error::Error + Send + Sync>> {
        let mut connection =
            rustls::ClientConnection::new(Arc::clone(&self.config), self.server_name.clone())?;
        let mut sock = stream;
        connection.complete_io(&mut sock)?;

        Ok(TlsStream {
            session: Arc::new(Mutex::new(Session {
                sock,
                tls: connection,
                frames: FrameBuffer::new(),
                read_timeout: Some(READ_TIMEOUT),
            })),
        })
    }

    fn observed_fingerprint(&self) -> Option<String> {
        self.observed.lock().clone()
    }
}

/// Certificate verification is record-only: the trust decision is made after
/// the handshake, before anything is sent, so a frontend can ask first.
struct RecordingVerifier {
    observed: Arc<Mutex<Option<String>>>,
}

impl ServerCertVerifier for RecordingVerifier {
    fn verify_server_cert(
        &self,
        end_entity: &rustls::Certificate,
        _intermediates: &[rustls::Certificate],
        _server_name: &rustls::ServerName,
        _scts: &mut dyn Iterator<Item = &[u8]>,
        _ocsp_response: &[u8],
        _now: std::time::SystemTime,
    ) -> Result<ServerCertVerified, rustls::Error> {
        let fingerprint = spki_fingerprint(&end_entity.0).map_err(|error| {
            rustls::Error::General(format!("certificate parse failed: {error}"))
        })?;
        *self.observed.lock() = Some(fingerprint);
        Ok(ServerCertVerified::assertion())
    }

    // Identity is pinned by the fingerprint check above; handshake signatures
    // are not additionally validated against a trust store.
    fn verify_tls12_signature(
        &self,
        _message: &[u8],
        _cert: &rustls::Certificate,
        _dss: &DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, rustls::Error> {
        Ok(HandshakeSignatureValid::assertion())
    }

    fn verify_tls13_signature(
        &self,
        _message: &[u8],
        _cert: &rustls::Certificate,
        _dss: &DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, rustls::Error> {
        Ok(HandshakeSignatureValid::assertion())
    }

    fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
        vec![
            rustls::SignatureScheme::RSA_PKCS1_SHA256,
            rustls::SignatureScheme::RSA_PKCS1_SHA384,
            rustls::SignatureScheme::RSA_PKCS1_SHA512,
            rustls::SignatureScheme::RSA_PSS_SHA256,
            rustls::SignatureScheme::RSA_PSS_SHA384,
            rustls::SignatureScheme::RSA_PSS_SHA512,
            rustls::SignatureScheme::ECDSA_NISTP256_SHA256,
            rustls::SignatureScheme::ECDSA_NISTP384_SHA384,
            rustls::SignatureScheme::ECDSA_NISTP521_SHA512,
        ]
    }
}

/// TCP socket + TLS session behind one mutex so clones share the stream.
struct Session {
    sock: TcpStream,
    tls: rustls::ClientConnection,
    frames: FrameBuffer,
    read_timeout: Option<Duration>,
}

impl Session {
    fn drive(&mut self) -> std::io::Result<()> {
        self.tls.complete_io(&mut self.sock).map(|_| ())
    }
}

#[derive(Clone)]
pub struct TlsStream {
    session: Arc<Mutex<Session>>,
}

impl TlsStream {
    /// Write one enveloped request. Application code allocates `id`.
    pub fn write_request(
        &mut self,
        id: u64,
        idem: Option<&str>,
        payload: &str,
    ) -> std::io::Result<()> {
        self.write_frame(&request_line(id, idem, payload))
    }

    /// Write one raw frame payload.
    pub fn write_frame(&mut self, payload: &str) -> std::io::Result<()> {
        let mut s = self.session.lock();
        write_protocol_frame(&mut s.tls.writer(), payload)?;
        s.drive()
    }

    /// Block until one whole frame arrives.
    pub fn read_frame(&mut self) -> std::io::Result<Option<String>> {
        let mut s = self.session.lock();
        loop {
            match s.frames.next_frame() {
                Some(Ok(frame)) => return Ok(Some(frame)),
                Some(Err(error)) => {
                    return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, error));
                }
                None => {}
            }

            let mut chunk = [0u8; 4096];
            match s.tls.reader().read(&mut chunk) {
                Ok(0) => return Ok(None),
                Ok(n) => s.frames.push(&chunk[..n]),
                Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => s.drive()?,
                Err(error) => return Err(error),
            }
        }
    }

    /// Block until one `RES` or `EVT` arrives.
    pub fn read_packet(&mut self) -> std::io::Result<Option<Packet>> {
        loop {
            let Some(frame) = self.read_frame()? else {
                return Ok(None);
            };
            match parse_packet(&frame) {
                Ok(packet) => return Ok(Some(packet)),
                Err(error) => {
                    return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, error));
                }
            }
        }
    }

    pub fn set_read_timeout(&self, duration: Option<Duration>) -> std::io::Result<()> {
        let mut s = self.session.lock();
        s.read_timeout = duration;
        s.sock.set_read_timeout(duration)
    }

    /// Return one packet if the socket already has one, without blocking.
    /// Used to pick up server pushes between requests.
    pub fn read_available_packet(&mut self) -> std::io::Result<Option<Packet>> {
        let mut s = self.session.lock();

        // `set_read_timeout(Some(0))` is rejected by std, so poll with the
        // shortest timeout the socket accepts instead of a nonblocking mode.
        s.sock
            .set_read_timeout(Some(std::time::Duration::from_millis(1)))?;
        let drained = (|| -> std::io::Result<Option<Packet>> {
            loop {
                match s.frames.next_frame() {
                    Some(Ok(frame)) => {
                        return match parse_packet(&frame) {
                            Ok(packet) => Ok(Some(packet)),
                            Err(error) => {
                                Err(std::io::Error::new(std::io::ErrorKind::InvalidData, error))
                            }
                        };
                    }
                    Some(Err(error)) => {
                        return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, error));
                    }
                    None => {}
                }

                let mut buf = [0u8; 4096];
                match s.tls.reader().read(&mut buf) {
                    Ok(0) => return Ok(None),
                    Ok(n) => s.frames.push(&buf[..n]),
                    // The reader only drains plaintext rustls has already
                    // decrypted; `drive` moves bytes off the socket.
                    Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
                        match s.drive() {
                            Ok(()) => continue,
                            Err(error)
                                if matches!(
                                    error.kind(),
                                    std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut
                                ) =>
                            {
                                return Ok(None);
                            }
                            Err(error) => return Err(error),
                        }
                    }
                    Err(error)
                        if matches!(
                            error.kind(),
                            std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut
                        ) =>
                    {
                        return Ok(None);
                    }
                    Err(error) => return Err(error),
                }
            }
        })();
        s.sock.set_read_timeout(s.read_timeout)?;
        drained
    }
}

impl Read for TlsStream {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        let mut s = self.session.lock();
        loop {
            match s.tls.reader().read(buf) {
                Ok(n) => return Ok(n),
                Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => s.drive()?,
                Err(error) => return Err(error),
            }
        }
    }
}

impl Write for TlsStream {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        let mut s = self.session.lock();
        s.tls.writer().write(buf)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        let mut s = self.session.lock();
        s.tls.writer().flush()?;
        s.drive()
    }
}