crabshell-services 0.2.0

//! Unix sockets a config opens, as Quickshell's `Socket`: lines in, bytes
//! out.

#[cfg(test)]
mod tests;

use std::cell::RefCell;
use std::io::{ErrorKind, Read, Write};
use std::os::fd::AsFd;
use std::os::unix::net::UnixStream;
use std::panic::{self, AssertUnwindSafe, Location};
use std::path::PathBuf;
use std::rc::Rc;
use std::sync::Once;

use calloop::{Interest, RegistrationToken};
use crabshell_reactive::{ReadSignal, Signal, untrack};
use crabshell_ui::Cx;

use crate::Shared;
use crate::lines::{Lines, panic_message};

/// The most bytes left to write before the socket is closed, so a program
/// that reads nothing can't grow the shell without end.
pub(crate) const MAX_OUTBOX: usize = 4 * 1024 * 1024;

/// What `cx.socket` needs from the shell: the loop.
pub(crate) struct Sockets(pub(crate) Rc<Shared>);

/// A connection to a Unix socket, from `cx.socket(path)`. It connects at
/// once, passes each line it reads to [`on_line`](Self::on_line), and
/// writes what [`write`](Self::write) is given. A clone is cheap and
/// shares the connection.
///
/// It closes when the context that opened it goes. When the other end
/// closes it, [`connected`](Self::connected) turns `false` and it stays
/// closed until [`reconnect`](Self::reconnect): nothing retries on its
/// own, so a socket that isn't there never wakes the shell.
#[derive(Clone)]
pub struct Socket {
  inner: Rc<Inner>,
}

struct Inner {
  /// `None` when the shell provided no services.
  shared: Option<Rc<Shared>>,
  path: PathBuf,
  /// Lives under the opening context's owner, so it says when the socket
  /// must close for good.
  connected: Signal<bool>,
  state: RefCell<State>,
}

/// What a config does with each line.
type OnLine = Rc<dyn Fn(&str)>;

#[derive(Default)]
struct State {
  connection: Option<Connection>,
  on_line: Option<OnLine>,
  /// Where `on_line` was set, for the log if it panics.
  at: Option<&'static Location<'static>>,
  lines: Lines,
  /// Counts connects and closes the config asked for, so lines read
  /// before one aren't passed on after it.
  serial: u64,
}

struct Connection {
  stream: UnixStream,
  /// The loop source that waits for something to read.
  reader: RegistrationToken,
  /// The loop source that waits for room to write, while the socket
  /// holds back part of `outbox`.
  writer: Option<RegistrationToken>,
  outbox: Vec<u8>,
  /// Why writing failed, after which only reading goes on.
  write_failed: Option<String>,
}

/// `cx.socket(path)`.
pub(crate) fn socket(cx: &Cx, path: PathBuf) -> Socket {
  let shared = cx.service::<Sockets>().map(|sockets| Rc::clone(&sockets.0));
  if shared.is_none() {
    static LOGGED: Once = Once::new();
    LOGGED.call_once(|| {
      tracing::warn!("cx.socket() needs the services the shell provides; it never connects");
    });
  }
  let inner = Rc::new(Inner {
    shared,
    path,
    connected: cx.signal(false),
    state: RefCell::default(),
  });
  // Dropped with the opening context's owner, which closes the socket.
  cx.owner().signal(CloseOnDrop(Rc::clone(&inner)));
  inner.connect();
  Socket { inner }
}

/// Closes the socket when the opening context's owner drops it.
struct CloseOnDrop(Rc<Inner>);

impl Drop for CloseOnDrop {
  fn drop(&mut self) {
    // Dropped outside the borrow: `on_line` may hold this socket, which
    // would keep it alive, and its drop may use it.
    let on_line = self.0.state.borrow_mut().on_line.take();
    drop(on_line);
    self.0.close();
  }
}

impl Socket {
  /// Whether it is connected: `true` from a connect that worked until the
  /// other end closes it, it breaks, or [`close`](Self::close).
  pub fn connected(&self) -> ReadSignal<bool> {
    self.inner.connected.read_only()
  }

  /// Calls `on_line` with each line read, without its `\n`, decoded as
  /// UTF-8 with bad bytes replaced. What is left when the other end closes
  /// counts as a last line. A line over 1 MiB is dropped and logged.
  ///
  /// It runs on the main thread. If it panics, that is logged with the
  /// socket's path and where `on_line` was called, and it isn't called
  /// again; the socket stays open.
  #[track_caller]
  pub fn on_line(self, on_line: impl Fn(&str) + 'static) -> Socket {
    let mut state = self.inner.state.borrow_mut();
    state.on_line = Some(Rc::new(on_line));
    state.at = Some(Location::caller());
    drop(state);
    self
  }

  /// Sends `data`, all of it, as the other end takes it. While not
  /// connected it is dropped, and that is logged. If more than 4 MiB wait
  /// because the other end reads nothing, the socket is closed.
  pub fn write(&self, data: impl AsRef<[u8]>) {
    self.inner.write(data.as_ref());
  }

  /// Closes it and connects again, as after the other end restarted.
  /// Lines the old connection read and `on_line` hasn't had yet are
  /// dropped.
  pub fn reconnect(&self) {
    self.inner.close();
    self.inner.connect();
  }

  /// Closes it. [`reconnect`](Self::reconnect) opens it again. Lines read
  /// and not yet passed on are dropped.
  pub fn close(&self) {
    self.inner.close();
  }
}

impl Inner {
  /// Whether the opening context is still there. Untracked, so an effect
  /// that opens or reconnects a socket doesn't run again for it.
  fn alive(&self) -> bool {
    untrack(|| self.connected.try_with(|_| ()).is_some())
  }

  fn connect(self: &Rc<Self>) {
    let Some(event_loop) = self.shared.as_ref().and_then(|shared| shared.event_loop()) else {
      return;
    };
    if !self.alive() {
      return;
    }
    let opened = UnixStream::connect(&self.path).and_then(|stream| {
      stream.set_nonblocking(true)?;
      let fd = stream.as_fd().try_clone_to_owned()?;
      Ok((stream, fd))
    });
    let (stream, fd) = match opened {
      Ok(opened) => opened,
      Err(error) => {
        let path = self.path.display();
        return tracing::warn!(%path, %error, "can't connect to the socket");
      }
    };
    let inner = Rc::clone(self);
    let read = Box::new(move || inner.on_readable());
    let Some(reader) = event_loop.insert_fd(fd, Interest::READ, read) else {
      return;
    };
    let mut state = self.state.borrow_mut();
    state.lines = Lines::default();
    state.serial += 1;
    state.connection = Some(Connection {
      stream,
      reader,
      writer: None,
      outbox: Vec::new(),
      write_failed: None,
    });
    drop(state);
    self.connected.try_set(true);
  }

  /// Reads what came, and passes each whole line on.
  fn on_readable(self: &Rc<Self>) {
    let mut chunk = [0; 64 * 1024];
    let mut state = self.state.borrow_mut();
    let Some(connection) = state.connection.as_mut() else {
      return;
    };
    let (lines, closed) = match connection.stream.read(&mut chunk) {
      Ok(0) => {
        let last = state.lines.finish();
        (
          last.into_iter().collect(),
          Some(String::from("the other end closed it")),
        )
      }
      Ok(read) => (state.lines.push(&chunk[..read], &self.path.display()), None),
      Err(error) if matches!(error.kind(), ErrorKind::WouldBlock | ErrorKind::Interrupted) => {
        return;
      }
      // As a reset from a peer that closed with input unread: the end.
      Err(error) => {
        let last = state.lines.finish();
        (last.into_iter().collect(), Some(error.to_string()))
      }
    };
    let serial = state.serial;
    drop(state);
    // Closed first, so `on_line` may reconnect.
    if let Some(why) = closed {
      self.drop_connection(&why);
    }
    self.deliver(lines, serial);
  }

  /// Calls `on_line` with each line, with nothing borrowed, until the
  /// config closes or reconnects the socket.
  fn deliver(&self, lines: Vec<String>, serial: u64) {
    let (on_line, at) = {
      let state = self.state.borrow();
      (state.on_line.clone(), state.at)
    };
    let Some(on_line) = on_line else {
      return;
    };
    for line in lines {
      if self.state.borrow().serial != serial {
        return;
      }
      if let Err(payload) = panic::catch_unwind(AssertUnwindSafe(|| on_line(&line))) {
        let path = self.path.display();
        let at = at.map_or_else(String::new, |at| format!(" (called at {at})"));
        let message = panic_message(payload.as_ref());
        tracing::error!("socket {path} on_line{at} panicked: {message}");
        let mut state = self.state.borrow_mut();
        // Unless it set a new one meanwhile.
        if state
          .on_line
          .as_ref()
          .is_some_and(|now| Rc::ptr_eq(now, &on_line))
        {
          state.on_line = None;
        }
        return;
      }
    }
  }

  fn write(self: &Rc<Self>, data: &[u8]) {
    let mut state = self.state.borrow_mut();
    let Some(connection) = state
      .connection
      .as_mut()
      .filter(|connection| connection.write_failed.is_none())
    else {
      let path = self.path.display();
      return tracing::warn!(%path, "the socket takes no writes now, so a write was dropped");
    };
    connection.outbox.extend_from_slice(data);
    drop(state);
    self.flush();
  }

  /// Writes what the socket takes, and waits for room for the rest.
  fn flush(self: &Rc<Self>) {
    let Some(event_loop) = self.shared.as_ref().and_then(|shared| shared.event_loop()) else {
      return;
    };
    let mut state = self.state.borrow_mut();
    let Some(connection) = state.connection.as_mut() else {
      return;
    };
    let mut broken = None;
    while !connection.outbox.is_empty() {
      match connection.stream.write(&connection.outbox) {
        Ok(written) => {
          connection.outbox.drain(..written);
        }
        Err(error) if error.kind() == ErrorKind::Interrupted => {}
        Err(error) if error.kind() == ErrorKind::WouldBlock => break,
        Err(error) => {
          // The other end may still have sent something: stop writing
          // and read on to its end.
          let path = self.path.display();
          tracing::info!(%path, %error, "the socket takes no more writes");
          connection.write_failed = Some(error.to_string());
          connection.outbox = Vec::new();
          break;
        }
      }
    }
    if connection.outbox.len() > MAX_OUTBOX {
      broken = Some(String::from("the other end reads nothing it is sent"));
    }
    if broken.is_none() && connection.outbox.is_empty() {
      // A big write's room isn't kept.
      connection.outbox.shrink_to(64 * 1024);
      if let Some(writer) = connection.writer.take() {
        event_loop.remove(writer);
      }
    } else if broken.is_none() && connection.writer.is_none() {
      match connection.stream.as_fd().try_clone_to_owned() {
        Ok(fd) => {
          let inner = Rc::clone(self);
          let write = Box::new(move || inner.flush());
          connection.writer = event_loop.insert_fd(fd, Interest::WRITE, write);
        }
        Err(error) => broken = Some(error.to_string()),
      }
    }
    drop(state);
    if let Some(why) = broken {
      self.drop_connection(&why);
    }
  }

  /// Ends the connection the other end ended or broke, and logs why.
  /// Lines it read before still reach `on_line`.
  fn drop_connection(&self, why: &str) {
    if self.state.borrow().connection.is_some() {
      let path = self.path.display();
      tracing::info!(%path, why, "the socket closed");
    }
    self.end(false);
  }

  /// Closes it as the config asked, so lines read before are dropped.
  fn close(&self) {
    self.end(true);
  }

  fn end(&self, asked: bool) {
    let connection = {
      let mut state = self.state.borrow_mut();
      state.lines = Lines::default();
      if asked {
        state.serial += 1;
      }
      state.connection.take()
    };
    if let (Some(connection), Some(event_loop)) = (
      connection,
      self.shared.as_ref().and_then(|shared| shared.event_loop()),
    ) {
      event_loop.remove(connection.reader);
      if let Some(writer) = connection.writer {
        event_loop.remove(writer);
      }
    }
    if untrack(|| self.connected.try_with(|connected| *connected)) == Some(true) {
      self.connected.try_set(false);
    }
  }
}