kraal 0.3.0

//! Clients: connections, their frames, attaching and detaching.

use std::io::Write;
use std::os::unix::net::UnixStream;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::{Duration, Instant};

use crossbeam_channel::{SendTimeoutError, TrySendError};

use super::{Client, ClientId, Event, Mode, Server};
use crate::ansi::Encoder;
use crate::input::InputParser;
use crate::pane::PaneSpec;
use crate::protocol::{self, ClientMsg, ServerMsg};

/// Clients connected at once (attached terminals, CLI commands, event streams).
const MAX_CLIENTS: usize = 256;

impl Server {
    pub(super) fn new_connection(&mut self, stream: UnixStream) {
        // Every connection costs two threads: a runaway script must not open thousands.
        if self.clients.len() + self.subscribers.len() + self.waits.len() >= MAX_CLIENTS {
            let mut stream = stream;
            let refused = ServerMsg::Response { result: Err(format!("too many connections ({MAX_CLIENTS})")) };
            let _ = protocol::write_json(&mut stream, &refused);
            return;
        }
        let id = self.alloc_id();
        let peer = crate::process::socket_peer(&stream).ok().and_then(|(pid, _)| u32::try_from(pid).ok());
        let Ok(read_half) = stream.try_clone() else { return };
        // One that connects without saying anything is let go.
        let _ = read_half.set_read_timeout(Some(Duration::from_secs(10)));
        // Without a thread to read or write it, the connection is closed (dropped).
        let Ok(out) = Outbox::start(stream, id) else { return };
        let tx = self.tx.clone();
        let reading = std::thread::Builder::new().name(format!("client-{id}-read")).spawn(move || {
            let mut r = std::io::BufReader::with_capacity(64 * 1024, read_half);
            let mut first = true;
            while let Ok(Some((kind, payload))) = protocol::read_frame(&mut r) {
                if std::mem::take(&mut first) {
                    let _ = r.get_ref().set_read_timeout(None);
                }
                if tx.send(Event::ClientFrame(id, kind, payload)).is_err() {
                    return;
                }
            }
            let _ = tx.send(Event::ClientGone(id));
        });
        if reading.is_err() {
            return;
        }
        let ws = self.workspaces.first().map(|w| w.id).unwrap_or(0);
        self.clients.insert(
            id,
            Client {
                id,
                out,
                attached: false,
                cols: 80,
                rows: 24,
                workspace: ws,
                parser: InputParser::new(),
                pending_since: None,
                mode: Mode::Normal,
                encoder: Encoder::new(),
                frame: Default::default(),
                memo: Default::default(),
                sidebar: None,
                dirty: false,
                last_input: Instant::now(),
                terminal_focused: true,
                pointer_motion: false,
                keys_full: None,
                keys_by_place: false,
                told_shown: None,
                paste: None,
                drag: None,
                toast: None,
                notice: None,
                last_click: None,
                welcome: false,
                remote: false,
                peer,
            },
        );
    }

    pub(super) fn client_frame(&mut self, id: ClientId, kind: u8, payload: Vec<u8>) {
        if kind == protocol::BYTES {
            self.client_input(id, &payload);
            return;
        }
        let msg: ClientMsg = match serde_json::from_slice(&payload) {
            Ok(m) => m,
            Err(e) => {
                self.send_msg(id, &ServerMsg::Response { result: Err(format!("bad message: {e}")) });
                return;
            }
        };
        match msg {
            ClientMsg::Attach { version, cols, rows, workspace, remote, pane, installed, .. } => {
                if let Some(c) = self.clients.get_mut(&id) {
                    c.remote = remote;
                }
                let (cols, rows) = clamp_size(cols, rows);
                if version != protocol::PROTOCOL_VERSION {
                    self.send_msg(id, &ServerMsg::exit(self.version_mismatch(version)));
                    return;
                }
                self.attach(id, cols, rows, workspace);
                if let Some(installed) = installed {
                    self.version_notice(id, &installed);
                }
                // A pane asked for (an agent chosen from another kraal's sidebar) is shown.
                if let Some(p) = pane {
                    match self.resolve(&p) {
                        Ok(p) => self.reveal(id, p),
                        Err(e) => self.warnings.push(e),
                    }
                }
            }
            ClientMsg::Resize { cols, rows } => {
                let (cols, rows) = clamp_size(cols, rows);
                if let Some(c) = self.clients.get_mut(&id) {
                    c.cols = cols;
                    c.rows = rows;
                    c.encoder.invalidate();
                }
                self.primary = Some(id);
                self.probe_cell_size(id);
                self.relayout();
            }
            ClientMsg::Detach => self.detach(id, "detached"),
            ClientMsg::Request { version, request, .. } => {
                let result = if version != protocol::PROTOCOL_VERSION {
                    Some(Err(self.version_mismatch(version)))
                } else {
                    self.request(id, request)
                };
                if let Some(result) = result {
                    self.send_msg(id, &ServerMsg::Response { result });
                    self.clients.remove(&id);
                }
            }
            // From a newer client: an attached one goes on without it, a connection that came
            // to ask something is told so.
            ClientMsg::Unknown(msg) => {
                if !self.clients.get(&id).is_some_and(|c| c.attached) {
                    self.send_msg(id, &ServerMsg::Response { result: Err(protocol::unsupported(&msg)) });
                    self.clients.remove(&id);
                }
            }
        }
    }

    /// What a client of another protocol generation is told. `kraal restart` and `kraal kill`
    /// stop a server of any version.
    fn version_mismatch(&self, version: u32) -> String {
        let (session, release) = (&self.session, protocol::KRAAL_VERSION);
        if version > protocol::PROTOCOL_VERSION {
            format!(
                "session {session} runs kraal {release}, too old for this kraal: `kraal -s {session} restart` restarts it on this one (its layout is kept)"
            )
        } else {
            format!("session {session} runs kraal {release}, too new for this kraal: update this one (`kraal update`)")
        }
    }

    pub(super) fn attach(&mut self, id: ClientId, cols: u16, rows: u16, workspace: Option<String>) {
        let mut target = None;
        if let Some(name) = workspace {
            match self.open_workspace(&name) {
                Ok(ws) => target = Some(ws),
                Err(e) => self.warnings.push(format!("{e:#}")),
            }
        }
        if self.workspaces.is_empty() {
            let ws = self.new_workspace("main".into(), None);
            if let Err(e) = self.spawn_pane_in(ws, None, PaneSpec::default(), crate::protocol::Placement::Tab, None) {
                self.drop_empty_workspaces();
                self.send_msg(id, &ServerMsg::exit(format!("cannot start a shell: {e:#}")));
                return;
            }
        }
        let ws = target.unwrap_or_else(|| self.workspaces[0].id);
        let welcome = !crate::paths::state_dir().join("welcomed").exists();
        // Before anything is drawn: what this server is.
        let features = Some(protocol::feature_names());
        self.send_msg(id, &ServerMsg::Attached { version: protocol::KRAAL_VERSION.into(), features });
        if let Some(c) = self.clients.get_mut(&id) {
            c.attached = true;
            c.cols = cols;
            c.rows = rows;
            c.workspace = ws;
            c.dirty = true;
            c.welcome = welcome;
            // Ask the outer terminal for its colours so panes querying them get real values.
            let _ = c.out.try_send(frame_bytes(b"\x1b]10;?\x07\x1b]11;?\x07"));
            // The mouse as the server's configuration says, which may have been read again since
            // the client read its own (a reconnection keeps the client's first reading).
            let mouse = if self.config.mouse { crate::client::MOUSE_ON } else { crate::client::MOUSE_OFF };
            let _ = c.out.try_send(frame_bytes(mouse.as_bytes()));
        }
        self.probe_graphics(id);
        self.primary = Some(id);
        self.relayout();
        self.mark_seen_focused(id);
        self.git_soon_for(id);
        self.show_warnings();
        self.mark_dirty();
        // The first frame goes at once: nothing else is coming to batch it with.
        self.next_render = Some(Instant::now());
    }

    pub fn detach(&mut self, id: ClientId, reason: &str) {
        self.graphics_farewell(id);
        self.send_msg(id, &ServerMsg::exit(reason));
        self.drop_client(id);
    }

    pub(super) fn drop_client(&mut self, id: ClientId) {
        self.clients.remove(&id);
        // A `kraal wait` or `kraal events` whose caller went away is dropped with it.
        self.waits.retain(|w| w.client != id);
        self.subscribers.retain(|s| s.client != id);
        // What the client was in the middle of goes with it.
        self.drafts.remove(&id);
        self.md_sessions.remove(&id);
        for p in self.panes.values_mut() {
            if let crate::pane::Content::Files(f) = &mut p.content {
                f.forget_viewer(id);
            }
        }
        if self.primary == Some(id) {
            self.primary = self.clients.values().filter(|c| c.attached).max_by_key(|c| c.last_input).map(|c| c.id);
            self.relayout();
        }
        self.graphics_client_gone(id);
    }

    pub fn send_msg(&self, id: ClientId, msg: &ServerMsg) {
        if let Some(c) = self.clients.get(&id) {
            let mut buf = Vec::new();
            if let Err(e) = protocol::write_json(&mut buf, msg) {
                // Too large to send: the client must not wait for an answer that never comes.
                buf.clear();
                let answer = ServerMsg::Response { result: Err(format!("the answer could not be sent: {e}")) };
                let _ = protocol::write_json(&mut buf, &answer);
            }
            let _ = c.out.send_timeout(buf, Duration::from_millis(200));
        }
    }
}

/// A client's terminal size, within what kraal draws (a bogus size must not allocate gigabytes).
fn clamp_size(cols: u16, rows: u16) -> (u16, u16) {
    (cols.clamp(1, 1000), rows.clamp(1, 500))
}

/// Terminal bytes in a protocol frame.
pub fn frame_bytes(bytes: &[u8]) -> Vec<u8> {
    let mut out = Vec::with_capacity(bytes.len() + 5);
    let _ = protocol::write_frame(&mut out, protocol::BYTES, bytes);
    out
}

/// Where what goes to a client is sent. Written to its socket at once, from the server's
/// thread, as far as the socket takes it without waiting and when nothing waits before it;
/// the rest is queued, in order, for the client's writer thread.
#[derive(Clone)]
pub struct Outbox {
    queue: crossbeam_channel::Sender<Vec<u8>>,
    /// Messages queued and not written yet.
    queued: Arc<AtomicUsize>,
    socket: Arc<UnixStream>,
}

impl Outbox {
    /// The outbox of a client connected on `stream`, and its writer thread.
    fn start(stream: UnixStream, id: ClientId) -> std::io::Result<Outbox> {
        let (tx, rx) = crossbeam_channel::bounded::<Vec<u8>>(64);
        // A client that stops reading is let go rather than holding its writer forever.
        let _ = stream.set_write_timeout(Some(Duration::from_secs(10)));
        let out = Outbox { queue: tx, queued: Arc::new(AtomicUsize::new(0)), socket: Arc::new(stream) };
        let (queued, socket) = (out.queued.clone(), out.socket.clone());
        std::thread::Builder::new().name(format!("client-{id}-write")).spawn(move || {
            let mut w = &*socket;
            for data in rx {
                if w.write_all(&data).and_then(|_| w.flush()).is_err() {
                    break;
                }
                // Written: once nothing is left, the server writes again itself.
                queued.fetch_sub(1, Ordering::AcqRel);
            }
            let _ = socket.shutdown(std::net::Shutdown::Both);
        })?;
        Ok(out)
    }

    pub fn try_send(&self, data: Vec<u8>) -> Result<(), TrySendError<Vec<u8>>> {
        let Some(rest) = self.write_now(data) else { return Ok(()) };
        self.queued.fetch_add(1, Ordering::AcqRel);
        self.queue.try_send(rest).inspect_err(|_| {
            self.queued.fetch_sub(1, Ordering::AcqRel);
        })
    }

    pub fn send_timeout(&self, data: Vec<u8>, wait: Duration) -> Result<(), SendTimeoutError<Vec<u8>>> {
        let Some(rest) = self.write_now(data) else { return Ok(()) };
        self.queued.fetch_add(1, Ordering::AcqRel);
        self.queue.send_timeout(rest, wait).inspect_err(|_| {
            self.queued.fetch_sub(1, Ordering::AcqRel);
        })
    }

    /// Messages waiting for the writer thread.
    pub fn len(&self) -> usize {
        self.queued.load(Ordering::Acquire)
    }

    /// Write `data` to the socket as far as it takes it now, unless something waits for the
    /// writer thread: what is left to queue. A message cut there goes on through the queue,
    /// which is empty then and has room for it.
    #[cfg(target_os = "linux")]
    fn write_now(&self, data: Vec<u8>) -> Option<Vec<u8>> {
        use std::os::fd::AsRawFd;
        if self.len() > 0 {
            return Some(data);
        }
        let flags = libc::MSG_DONTWAIT | libc::MSG_NOSIGNAL;
        let mut done = 0;
        while done < data.len() {
            let rest = &data[done..];
            // SAFETY: sends from a live slice on a socket the outbox holds.
            let n = unsafe { libc::send(self.socket.as_raw_fd(), rest.as_ptr().cast(), rest.len(), flags) };
            match n {
                n if n > 0 => done += n as usize,
                n if n < 0 && std::io::Error::last_os_error().kind() == std::io::ErrorKind::Interrupted => {}
                _ => break,
            }
        }
        match done {
            _ if done == data.len() => None,
            0 => Some(data),
            _ => Some(data[done..].to_vec()),
        }
    }

    /// Everything goes through the writer thread: macOS does not let a send be told not to wait
    /// (`MSG_DONTWAIT` waits all the same, up to the socket's timeout), and the socket itself
    /// must wait for that thread and the reader.
    #[cfg(not(target_os = "linux"))]
    fn write_now(&self, data: Vec<u8>) -> Option<Vec<u8>> {
        Some(data)
    }
}

#[cfg(test)]
mod tests {
    use std::io::Read;

    use super::*;

    /// Wait until `out` has nothing queued.
    fn drained(out: &Outbox) {
        let start = Instant::now();
        while out.len() > 0 {
            assert!(start.elapsed() < Duration::from_secs(5));
            std::thread::sleep(Duration::from_millis(5));
        }
    }

    #[test]
    fn what_a_client_does_not_read_yet_waits_for_it_in_order() {
        use std::os::fd::AsRawFd;
        let (server_end, mut client_end) = UnixStream::pair().unwrap();
        // A small send buffer, as macOS gives a socket: messages are cut often.
        let size: libc::c_int = 8 * 1024;
        // SAFETY: sets an int option on a socket the test holds.
        let set = unsafe {
            libc::setsockopt(
                server_end.as_raw_fd(),
                libc::SOL_SOCKET,
                libc::SO_SNDBUF,
                (&size as *const libc::c_int).cast(),
                std::mem::size_of::<libc::c_int>() as libc::socklen_t,
            )
        };
        assert_eq!(set, 0);
        let out = Outbox::start(server_end, 1).unwrap();
        // Small enough to go at once (through the writer thread where a send cannot be told not
        // to wait).
        out.try_send(b"first".to_vec()).unwrap();
        #[cfg(target_os = "linux")]
        assert_eq!(out.len(), 0);
        let mut first = [0u8; 5];
        client_end.read_exact(&mut first).unwrap();
        assert_eq!(&first, b"first");
        drained(&out);
        // Far more than the socket holds: cut where it is full, the rest queued, nothing waits.
        let messages: Vec<Vec<u8>> = (0..40u8).map(|i| vec![b'a' + i % 26; 64 * 1024 + usize::from(i)]).collect();
        let start = Instant::now();
        for m in &messages {
            out.try_send(m.clone()).unwrap();
        }
        assert!(start.elapsed() < Duration::from_millis(500), "{:?}", start.elapsed());
        assert!(out.len() > 0);
        let want: Vec<u8> = messages.concat();
        let mut got = vec![0u8; want.len()];
        client_end.read_exact(&mut got).unwrap();
        assert!(got == want, "every byte, in order");
        drained(&out);
        out.try_send(b"last".to_vec()).unwrap();
        #[cfg(target_os = "linux")]
        assert_eq!(out.len(), 0, "written at once again");
        drop(out);
        let mut rest = Vec::new();
        client_end.read_to_end(&mut rest).unwrap();
        assert_eq!(rest, b"last", "and the connection closed after it");
    }
}