Vectrum
Sign In
186 lines · 7.1 KB · Rust
//! The client side: a connection to another device's app on its own thread,
//! reconnecting when it drops. Messages from the host reach the workspace as
//! [`ClientEvent`]s.

use std::{
    io::ErrorKind,
    net::{TcpStream, ToSocketAddrs},
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
        mpsc,
    },
    thread,
    time::Duration,
};

use tungstenite::{Message, WebSocket};

use super::{ClientMessage, PROTOCOL, ServerMessage, with_port};

const CONNECT_TIMEOUT: Duration = Duration::from_secs(6);
const POLL: Duration = Duration::from_millis(40);

#[derive(Clone, Debug, PartialEq)]
pub enum ConnectionStatus {
    Connecting,
    /// Connected to the host of that name.
    Connected(String),
    Failed(String),
    /// The host refused the key; no more retries until the settings change.
    Denied(String),
}

pub enum ClientEvent {
    Status(ConnectionStatus),
    Message(ServerMessage),
}

pub struct RemoteClient {
    outbox: mpsc::Sender<String>,
    stop: Arc<AtomicBool>,
    pub events: smol::channel::Receiver<ClientEvent>,
}

/// A way to the host that works off the UI thread.
#[derive(Clone)]
pub struct ClientSender(mpsc::Sender<String>);

impl ClientSender {
    pub fn send(&self, message: &ClientMessage) {
        if let Ok(text) = serde_json::to_string(message) {
            let _ = self.0.send(text);
        }
    }
}

impl RemoteClient {
    pub fn sender(&self) -> ClientSender {
        ClientSender(self.outbox.clone())
    }

    /// Connects to `address` (`host` or `host:port`) and keeps reconnecting
    /// until dropped.
    pub fn connect(address: &str, key: String, name: String) -> Self {
        let address = with_port(address);
        let stop = Arc::new(AtomicBool::new(false));
        let (outbox, inbox) = mpsc::channel::<String>();
        let (events_tx, events) = smol::channel::unbounded();
        {
            let stop = stop.clone();
            let _ = thread::Builder::new().name("remote-client".into()).spawn(move || {
                let mut backoff = Duration::from_secs(2);
                while !stop.load(Ordering::Relaxed) {
                    let _ = events_tx.send_blocking(ClientEvent::Status(ConnectionStatus::Connecting));
                    match open(&address, &key, &name) {
                        Ok((mut socket, host)) => {
                            backoff = Duration::from_secs(2);
                            // Whatever was asked while offline is stale now.
                            while inbox.try_recv().is_ok() {}
                            let _ = events_tx.send_blocking(ClientEvent::Status(ConnectionStatus::Connected(host)));
                            let reason = run(&mut socket, &inbox, &events_tx, &stop);
                            let _ = socket.close(None);
                            if stop.load(Ordering::Relaxed) {
                                break;
                            }
                            let _ = events_tx.send_blocking(ClientEvent::Status(ConnectionStatus::Failed(reason)));
                        }
                        Err(Failure::Denied(reason)) => {
                            let _ = events_tx.send_blocking(ClientEvent::Status(ConnectionStatus::Denied(reason)));
                            break;
                        }
                        Err(Failure::Other(reason)) => {
                            let _ = events_tx.send_blocking(ClientEvent::Status(ConnectionStatus::Failed(reason)));
                        }
                    }
                    // Wait before retrying, a little longer each time.
                    let mut waited = Duration::ZERO;
                    while waited < backoff && !stop.load(Ordering::Relaxed) {
                        thread::sleep(Duration::from_millis(200));
                        waited += Duration::from_millis(200);
                    }
                    backoff = (backoff * 2).min(Duration::from_secs(30));
                }
            });
        }
        Self { outbox, stop, events }
    }

    pub fn send(&self, message: &ClientMessage) {
        if let Ok(text) = serde_json::to_string(message) {
            let _ = self.outbox.send(text);
        }
    }
}

impl Drop for RemoteClient {
    fn drop(&mut self) {
        self.stop.store(true, Ordering::Relaxed);
    }
}

enum Failure {
    Denied(String),
    Other(String),
}

/// Connects, says hello, and waits for the welcome.
fn open(address: &str, key: &str, name: &str) -> Result<(WebSocket<TcpStream>, String), Failure> {
    let other = |err: &dyn std::fmt::Display| Failure::Other(err.to_string());
    let target = address.to_socket_addrs().map_err(|err| other(&err))?.next().ok_or_else(|| other(&"no such host"))?;
    let stream = TcpStream::connect_timeout(&target, CONNECT_TIMEOUT).map_err(|err| other(&err))?;
    let _ = stream.set_nodelay(true);
    stream.set_read_timeout(Some(Duration::from_secs(15))).map_err(|err| other(&err))?;
    let (mut socket, _) = tungstenite::client(format!("ws://{address}/"), stream).map_err(|err| other(&err))?;
    let hello = ClientMessage::Hello { key: key.to_string(), name: name.to_string(), protocol: PROTOCOL };
    socket
        .send(Message::text(serde_json::to_string(&hello).unwrap_or_default()))
        .map_err(|err| other(&err))?;
    loop {
        match socket.read().map_err(|err| other(&err))? {
            Message::Text(text) => {
                return match serde_json::from_str::<ServerMessage>(text.as_str()) {
                    Ok(ServerMessage::Welcome { host, .. }) => {
                        let _ = socket.get_mut().set_read_timeout(Some(POLL));
                        Ok((socket, host))
                    }
                    Ok(ServerMessage::Denied { reason }) => Err(Failure::Denied(reason)),
                    _ => Err(other(&"the host answered something unexpected")),
                };
            }
            Message::Close(_) => return Err(Failure::Denied("The host closed the connection.".into())),
            _ => {}
        }
    }
}

/// Trades messages until the connection drops; returns why it did.
fn run(
    socket: &mut WebSocket<TcpStream>,
    inbox: &mpsc::Receiver<String>,
    events: &smol::channel::Sender<ClientEvent>,
    stop: &AtomicBool,
) -> String {
    while !stop.load(Ordering::Relaxed) {
        while let Ok(text) = inbox.try_recv() {
            if let Err(err) = socket.send(Message::text(text)) {
                return err.to_string();
            }
        }
        match socket.read() {
            Ok(Message::Text(text)) => match serde_json::from_str::<ServerMessage>(text.as_str()) {
                Ok(message) => {
                    if events.send_blocking(ClientEvent::Message(message)).is_err() {
                        return "closed".into();
                    }
                }
                Err(err) => eprintln!("vectrum-code: unreadable message from host: {err}"),
            },
            Ok(Message::Close(_)) => return "The host closed the connection.".into(),
            Ok(_) => {}
            Err(tungstenite::Error::Io(err)) if matches!(err.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut) => {}
            Err(err) => return err.to_string(),
        }
    }
    "closed".into()
}