crabshell_services/
socket_server.rs

1//! Unix socket servers a config runs, as Quickshell's `SocketServer`: each
2//! client that connects is a [`Socket`].
3
4#[cfg(test)]
5mod tests;
6
7use std::cell::{Cell, RefCell};
8use std::io::ErrorKind;
9use std::os::fd::AsFd;
10use std::os::unix::fs::{FileTypeExt, MetadataExt};
11use std::os::unix::net::{UnixListener, UnixStream};
12use std::panic::{self, AssertUnwindSafe, Location};
13use std::path::{Path, PathBuf};
14use std::rc::Rc;
15use std::sync::Once;
16use std::time::{Duration, Instant};
17
18use calloop::timer::TimeoutAction;
19use calloop::{Interest, RegistrationToken};
20use crabshell_ui::Cx;
21use rustix::io::Errno;
22use rustix::net::{AddressFamily, SocketAddrUnix, SocketFlags, SocketType};
23
24use crate::Shared;
25use crate::lines::panic_message;
26use crate::socket::{self, Socket, Sockets};
27
28/// The most connections open at once. A client past it is closed at once.
29pub(crate) const MAX_CONNECTIONS: usize = 64;
30
31/// How long the server rests after an accept that failed and left the
32/// client waiting, as when the shell is out of file descriptors.
33const BACK_OFF: Duration = Duration::from_secs(1);
34
35/// A server on a Unix socket, from `cx.socket_server(path)`, as
36/// Quickshell's `SocketServer`. It listens from the start and hands each
37/// client that connects to [`on_connection`](Self::on_connection) as a
38/// [`Socket`]. A clone is cheap and shares the server.
39///
40/// It stops when the context that made it goes, or at
41/// [`close`](Self::close): every connection closes, and the socket file
42/// goes if it is still the one the server made. A socket file already at
43/// the path is replaced if nothing listens on it. One that something
44/// listens on, or anything that isn't a socket, is left alone and logged,
45/// and the server never listens. The file gets the permissions the
46/// shell's umask leaves, so keep it in a directory only the user can
47/// reach, such as `$XDG_RUNTIME_DIR`.
48#[derive(Clone)]
49pub struct SocketServer {
50  inner: Rc<Inner>,
51}
52
53/// What a config does with each connection.
54type OnConnection = Rc<dyn Fn(&Cx, Socket)>;
55
56struct Inner {
57  /// `None` when the shell provided no services.
58  shared: Option<Rc<Shared>>,
59  path: PathBuf,
60  /// The server's own context. Each connection's is made inside it, so
61  /// disposing it ends them all.
62  cx: Cx,
63  /// Taken when the server stops, so a client still waiting is turned
64  /// away.
65  listener: RefCell<Option<UnixListener>>,
66  /// The socket file's device and inode, so only ours is ever removed.
67  file: Option<(u64, u64)>,
68  /// The loop source that waits for clients, or the timer that brings it
69  /// back after a back-off.
70  source: Cell<Option<RegistrationToken>>,
71  /// How many connections' contexts are alive.
72  open: Rc<Cell<usize>>,
73  on_connection: RefCell<Option<(OnConnection, &'static Location<'static>)>>,
74}
75
76/// `cx.socket_server(path)`.
77pub(crate) fn socket_server(cx: &Cx, path: PathBuf) -> SocketServer {
78  let shared = cx.service::<Sockets>().map(|sockets| Rc::clone(&sockets.0));
79  if shared.is_none() {
80    static LOGGED: Once = Once::new();
81    LOGGED.call_once(|| {
82      tracing::warn!("cx.socket_server() needs the services the shell provides; it never listens");
83    });
84  }
85  let bound = shared.as_ref().and_then(|_| {
86    bind(&path)
87      .inspect_err(|error| {
88        let path = path.display();
89        tracing::warn!(%path, %error, "can't listen on the socket");
90      })
91      .ok()
92  });
93  let (listener, file) = bound.unzip();
94  let cx = cx.child();
95  let inner = Rc::new(Inner {
96    shared,
97    path,
98    cx: cx.clone(),
99    listener: RefCell::new(listener),
100    file,
101    source: Cell::new(None),
102    open: Rc::default(),
103    on_connection: RefCell::new(None),
104  });
105  // Dropped with the server's context, which stops it.
106  cx.owner().signal(StopOnDrop(Rc::clone(&inner)));
107  inner.listen();
108  SocketServer { inner }
109}
110
111/// Stops the server when its context's owner drops it.
112struct StopOnDrop(Rc<Inner>);
113
114impl Drop for StopOnDrop {
115  fn drop(&mut self) {
116    // Dropped outside the borrow: it may hold this server, and its drop
117    // may use it.
118    let on_connection = self.0.on_connection.borrow_mut().take();
119    drop(on_connection);
120    self.0.stop();
121  }
122}
123
124/// Counts a connection while its context lives.
125struct Open(Rc<Cell<usize>>);
126
127impl Open {
128  fn count(open: &Rc<Cell<usize>>) -> Open {
129    open.set(open.get() + 1);
130    Open(Rc::clone(open))
131  }
132}
133
134impl Drop for Open {
135  fn drop(&mut self) {
136    self.0.set(self.0.get() - 1);
137  }
138}
139
140impl SocketServer {
141  /// Calls `on_connection` with each client that connects: a context of
142  /// the connection's own, which goes on the loop's turn after the
143  /// connection ends, and the connection as a [`Socket`], connected, for
144  /// `on_line` or `on_split`, `write` and `close`. Lines are capped as a
145  /// `Socket`'s are. The socket's `connected` signal goes with that
146  /// context too, so code outside it reads it with `try_get`. A client
147  /// that connects while nothing handles it is closed at once.
148  ///
149  /// A client that only stops sending, as `printf cmd | socat -
150  /// UNIX-CONNECT:path` does, has its last line passed on, and the
151  /// connection stays open for the reply until the config closes it, a
152  /// write fails or 5 s pass, as the shell's own IPC socket waits. One
153  /// that hangs up while still sending ends the connection at once; one
154  /// that hangs up after it stopped sending, as socat does once its wait
155  /// for a reply ends, keeps its place until a write fails or the 5 s
156  /// pass, so a config that answers should close the connection.
157  ///
158  /// It runs on the main thread. If it panics, that is logged with the
159  /// server's path and where `on_connection` was called, that connection
160  /// is closed, and it isn't called again; the server and its other
161  /// connections go on.
162  #[track_caller]
163  pub fn on_connection(self, on_connection: impl Fn(&Cx, Socket) + 'static) -> SocketServer {
164    // Kept only while the server runs, so a handler that holds it can't
165    // keep it alive.
166    if !self.inner.cx.owner().is_disposed() {
167      let handler: OnConnection = Rc::new(on_connection);
168      let older = (self.inner.on_connection)
169        .borrow_mut()
170        .replace((handler, Location::caller()));
171      drop(older);
172    }
173    self
174  }
175
176  /// Stops it, as its context going does: every connection closes, and
177  /// the socket file goes. Nothing opens it again.
178  pub fn close(&self) {
179    self.inner.cx.owner().dispose();
180  }
181}
182
183impl Inner {
184  fn event_loop(&self) -> Option<Box<dyn crate::EventLoop>> {
185    self.shared.as_ref()?.event_loop()
186  }
187
188  /// Waits for clients, unless the server stopped.
189  fn listen(self: &Rc<Self>) {
190    let Some(event_loop) = self.event_loop() else {
191      return;
192    };
193    let fd = match self
194      .listener
195      .borrow()
196      .as_ref()
197      .map(|listener| listener.as_fd().try_clone_to_owned())
198    {
199      Some(Ok(fd)) => fd,
200      Some(Err(error)) => {
201        let path = self.path.display();
202        return tracing::warn!(%path, %error, "can't listen on the socket");
203      }
204      None => return,
205    };
206    let inner = Rc::clone(self);
207    let ready = Box::new(move || inner.accept());
208    self
209      .source
210      .set(event_loop.insert_fd(fd, Interest::READ, ready));
211  }
212
213  /// Takes every client waiting.
214  fn accept(self: &Rc<Self>) {
215    loop {
216      // Not borrowed while a handler runs: it may close the server.
217      let accepted = self.listener.borrow().as_ref().map(UnixListener::accept);
218      match accepted {
219        Some(Ok((stream, _))) => self.serve(stream),
220        Some(Err(error)) if error.kind() == ErrorKind::WouldBlock => return,
221        Some(Err(error))
222          if matches!(
223            error.kind(),
224            ErrorKind::Interrupted | ErrorKind::ConnectionAborted
225          ) => {}
226        Some(Err(error)) => {
227          let path = self.path.display();
228          tracing::warn!(%path, %error, "can't accept a connection; trying again in 1 s");
229          return self.back_off();
230        }
231        None => return,
232      }
233    }
234  }
235
236  /// Stops waiting for clients for [`BACK_OFF`], since the one waiting
237  /// would wake the loop again at once.
238  fn back_off(self: &Rc<Self>) {
239    let Some(event_loop) = self.event_loop() else {
240      return;
241    };
242    if let Some(token) = self.source.take() {
243      event_loop.remove(token);
244    }
245    let inner = Rc::clone(self);
246    let again = Box::new(move || {
247      inner.source.set(None);
248      inner.listen();
249      TimeoutAction::Drop
250    });
251    self
252      .source
253      .set(event_loop.insert_timer(Instant::now() + BACK_OFF, again));
254  }
255
256  /// Hands a client to the config's handler, under a context of its own.
257  fn serve(self: &Rc<Self>, stream: UnixStream) {
258    let path = self.path.display();
259    if self.open.get() >= MAX_CONNECTIONS {
260      return tracing::warn!(%path, "refused a connection: {MAX_CONNECTIONS} are open already");
261    }
262    let handler = self.on_connection.borrow().clone();
263    let (Some((handler, at)), Some(shared)) = (handler, &self.shared) else {
264      return tracing::info!(%path, "closed a connection, as nothing handles it");
265    };
266    let cx = self.cx.child();
267    cx.owner().signal(Open::count(&self.open));
268    let Some(socket) = socket::accepted(&cx, shared, self.path.clone(), stream) else {
269      return cx.owner().dispose();
270    };
271    let handed = socket.clone();
272    if let Err(payload) = panic::catch_unwind(AssertUnwindSafe(|| handler(&cx, handed))) {
273      let message = panic_message(payload.as_ref());
274      tracing::error!("socket server {path} on_connection (called at {at}) panicked: {message}");
275      let mut current = self.on_connection.borrow_mut();
276      // Unless it set a new one meanwhile.
277      let disabled = current
278        .as_ref()
279        .is_some_and(|(now, _)| Rc::ptr_eq(now, &handler))
280        .then(|| current.take());
281      drop(current);
282      drop(disabled);
283      socket.close();
284    }
285  }
286
287  /// Stops listening and removes the socket file if it is still ours.
288  fn stop(&self) {
289    if let (Some(token), Some(event_loop)) = (self.source.take(), self.event_loop()) {
290      event_loop.remove(token);
291    }
292    let listener = self.listener.borrow_mut().take();
293    if let (Some(_), Some((dev, ino))) = (listener, self.file) {
294      let ours = std::fs::symlink_metadata(&self.path)
295        .is_ok_and(|meta| (meta.dev(), meta.ino()) == (dev, ino));
296      if ours {
297        let _ = std::fs::remove_file(&self.path);
298      }
299    }
300  }
301}
302
303/// Listens at `path`, in place of a socket file nobody listens on.
304fn bind(path: &Path) -> std::io::Result<(UnixListener, (u64, u64))> {
305  match std::fs::symlink_metadata(path) {
306    Ok(meta) if !meta.file_type().is_socket() => {
307      let why = "something that isn't a socket is there";
308      return Err(std::io::Error::new(ErrorKind::AlreadyExists, why));
309    }
310    Ok(_) if listening(path)? => {
311      let why = "another program listens there";
312      return Err(std::io::Error::new(ErrorKind::AddrInUse, why));
313    }
314    Ok(_) => {
315      std::fs::remove_file(path)?;
316      let path = path.display();
317      tracing::info!(%path, "replaced a socket nobody listened on");
318    }
319    Err(error) if error.kind() == ErrorKind::NotFound => {}
320    Err(error) => return Err(error),
321  }
322  let listener = UnixListener::bind(path)?;
323  listener.set_nonblocking(true)?;
324  let meta = std::fs::symlink_metadata(path)?;
325  Ok((listener, (meta.dev(), meta.ino())))
326}
327
328/// Whether something takes connections at `path`, found without waiting:
329/// a full backlog counts as listening.
330fn listening(path: &Path) -> std::io::Result<bool> {
331  let flags = SocketFlags::NONBLOCK | SocketFlags::CLOEXEC;
332  let probe = rustix::net::socket_with(AddressFamily::UNIX, SocketType::STREAM, flags, None)?;
333  let refused = rustix::net::connect(&probe, &SocketAddrUnix::new(path)?);
334  Ok(!matches!(refused, Err(Errno::CONNREFUSED | Errno::NOENT)))
335}
crabshell_services/socket_server.rs

crabshell_services/
socket_server.rs
