crabshell_services/
socket.rs

1//! Unix sockets a config opens, as Quickshell's `Socket`: lines in, bytes
2//! out.
3
4#[cfg(test)]
5mod tests;
6
7use std::cell::RefCell;
8use std::io::{ErrorKind, Read, Write};
9use std::os::fd::AsFd;
10use std::os::unix::net::UnixStream;
11use std::panic::{self, AssertUnwindSafe, Location};
12use std::path::PathBuf;
13use std::rc::Rc;
14use std::sync::Once;
15
16use calloop::{Interest, RegistrationToken};
17use crabshell_reactive::{ReadSignal, Signal, untrack};
18use crabshell_ui::Cx;
19
20use crate::Shared;
21use crate::lines::{Lines, panic_message};
22
23/// The most bytes left to write before the socket is closed, so a program
24/// that reads nothing can't grow the shell without end.
25pub(crate) const MAX_OUTBOX: usize = 4 * 1024 * 1024;
26
27/// What `cx.socket` needs from the shell: the loop.
28pub(crate) struct Sockets(pub(crate) Rc<Shared>);
29
30/// A connection to a Unix socket, from `cx.socket(path)`. It connects at
31/// once, passes each line it reads to [`on_line`](Self::on_line), and
32/// writes what [`write`](Self::write) is given. A clone is cheap and
33/// shares the connection.
34///
35/// It closes when the context that opened it goes. When the other end
36/// closes it, [`connected`](Self::connected) turns `false` and it stays
37/// closed until [`reconnect`](Self::reconnect): nothing retries on its
38/// own, so a socket that isn't there never wakes the shell.
39#[derive(Clone)]
40pub struct Socket {
41  inner: Rc<Inner>,
42}
43
44struct Inner {
45  /// `None` when the shell provided no services.
46  shared: Option<Rc<Shared>>,
47  path: PathBuf,
48  /// Lives under the opening context's owner, so it says when the socket
49  /// must close for good.
50  connected: Signal<bool>,
51  state: RefCell<State>,
52}
53
54/// What a config does with each line.
55type OnLine = Rc<dyn Fn(&str)>;
56
57#[derive(Default)]
58struct State {
59  connection: Option<Connection>,
60  on_line: Option<OnLine>,
61  /// Where `on_line` was set, for the log if it panics.
62  at: Option<&'static Location<'static>>,
63  lines: Lines,
64  /// Counts connects and closes the config asked for, so lines read
65  /// before one aren't passed on after it.
66  serial: u64,
67}
68
69struct Connection {
70  stream: UnixStream,
71  /// The loop source that waits for something to read.
72  reader: RegistrationToken,
73  /// The loop source that waits for room to write, while the socket
74  /// holds back part of `outbox`.
75  writer: Option<RegistrationToken>,
76  outbox: Vec<u8>,
77  /// Why writing failed, after which only reading goes on.
78  write_failed: Option<String>,
79}
80
81/// `cx.socket(path)`.
82pub(crate) fn socket(cx: &Cx, path: PathBuf) -> Socket {
83  let shared = cx.service::<Sockets>().map(|sockets| Rc::clone(&sockets.0));
84  if shared.is_none() {
85    static LOGGED: Once = Once::new();
86    LOGGED.call_once(|| {
87      tracing::warn!("cx.socket() needs the services the shell provides; it never connects");
88    });
89  }
90  let inner = Rc::new(Inner {
91    shared,
92    path,
93    connected: cx.signal(false),
94    state: RefCell::default(),
95  });
96  // Dropped with the opening context's owner, which closes the socket.
97  cx.owner().signal(CloseOnDrop(Rc::clone(&inner)));
98  inner.connect();
99  Socket { inner }
100}
101
102/// Closes the socket when the opening context's owner drops it.
103struct CloseOnDrop(Rc<Inner>);
104
105impl Drop for CloseOnDrop {
106  fn drop(&mut self) {
107    // Dropped outside the borrow: `on_line` may hold this socket, which
108    // would keep it alive, and its drop may use it.
109    let on_line = self.0.state.borrow_mut().on_line.take();
110    drop(on_line);
111    self.0.close();
112  }
113}
114
115impl Socket {
116  /// Whether it is connected: `true` from a connect that worked until the
117  /// other end closes it, it breaks, or [`close`](Self::close).
118  pub fn connected(&self) -> ReadSignal<bool> {
119    self.inner.connected.read_only()
120  }
121
122  /// Calls `on_line` with each line read, without its `\n`, decoded as
123  /// UTF-8 with bad bytes replaced. What is left when the other end closes
124  /// counts as a last line. A line over 1 MiB is dropped and logged.
125  ///
126  /// It runs on the main thread. If it panics, that is logged with the
127  /// socket's path and where `on_line` was called, and it isn't called
128  /// again; the socket stays open.
129  #[track_caller]
130  pub fn on_line(self, on_line: impl Fn(&str) + 'static) -> Socket {
131    let mut state = self.inner.state.borrow_mut();
132    state.on_line = Some(Rc::new(on_line));
133    state.at = Some(Location::caller());
134    drop(state);
135    self
136  }
137
138  /// Sends `data`, all of it, as the other end takes it. While not
139  /// connected it is dropped, and that is logged. If more than 4 MiB wait
140  /// because the other end reads nothing, the socket is closed.
141  pub fn write(&self, data: impl AsRef<[u8]>) {
142    self.inner.write(data.as_ref());
143  }
144
145  /// Closes it and connects again, as after the other end restarted.
146  /// Lines the old connection read and `on_line` hasn't had yet are
147  /// dropped.
148  pub fn reconnect(&self) {
149    self.inner.close();
150    self.inner.connect();
151  }
152
153  /// Closes it. [`reconnect`](Self::reconnect) opens it again. Lines read
154  /// and not yet passed on are dropped.
155  pub fn close(&self) {
156    self.inner.close();
157  }
158}
159
160impl Inner {
161  /// Whether the opening context is still there. Untracked, so an effect
162  /// that opens or reconnects a socket doesn't run again for it.
163  fn alive(&self) -> bool {
164    untrack(|| self.connected.try_with(|_| ()).is_some())
165  }
166
167  fn connect(self: &Rc<Self>) {
168    let Some(event_loop) = self.shared.as_ref().and_then(|shared| shared.event_loop()) else {
169      return;
170    };
171    if !self.alive() {
172      return;
173    }
174    let opened = UnixStream::connect(&self.path).and_then(|stream| {
175      stream.set_nonblocking(true)?;
176      let fd = stream.as_fd().try_clone_to_owned()?;
177      Ok((stream, fd))
178    });
179    let (stream, fd) = match opened {
180      Ok(opened) => opened,
181      Err(error) => {
182        let path = self.path.display();
183        return tracing::warn!(%path, %error, "can't connect to the socket");
184      }
185    };
186    let inner = Rc::clone(self);
187    let read = Box::new(move || inner.on_readable());
188    let Some(reader) = event_loop.insert_fd(fd, Interest::READ, read) else {
189      return;
190    };
191    let mut state = self.state.borrow_mut();
192    state.lines = Lines::default();
193    state.serial += 1;
194    state.connection = Some(Connection {
195      stream,
196      reader,
197      writer: None,
198      outbox: Vec::new(),
199      write_failed: None,
200    });
201    drop(state);
202    self.connected.try_set(true);
203  }
204
205  /// Reads what came, and passes each whole line on.
206  fn on_readable(self: &Rc<Self>) {
207    let mut chunk = [0; 64 * 1024];
208    let mut state = self.state.borrow_mut();
209    let Some(connection) = state.connection.as_mut() else {
210      return;
211    };
212    let (lines, closed) = match connection.stream.read(&mut chunk) {
213      Ok(0) => {
214        let last = state.lines.finish();
215        (
216          last.into_iter().collect(),
217          Some(String::from("the other end closed it")),
218        )
219      }
220      Ok(read) => (state.lines.push(&chunk[..read], &self.path.display()), None),
221      Err(error) if matches!(error.kind(), ErrorKind::WouldBlock | ErrorKind::Interrupted) => {
222        return;
223      }
224      // As a reset from a peer that closed with input unread: the end.
225      Err(error) => {
226        let last = state.lines.finish();
227        (last.into_iter().collect(), Some(error.to_string()))
228      }
229    };
230    let serial = state.serial;
231    drop(state);
232    // Closed first, so `on_line` may reconnect.
233    if let Some(why) = closed {
234      self.drop_connection(&why);
235    }
236    self.deliver(lines, serial);
237  }
238
239  /// Calls `on_line` with each line, with nothing borrowed, until the
240  /// config closes or reconnects the socket.
241  fn deliver(&self, lines: Vec<String>, serial: u64) {
242    let (on_line, at) = {
243      let state = self.state.borrow();
244      (state.on_line.clone(), state.at)
245    };
246    let Some(on_line) = on_line else {
247      return;
248    };
249    for line in lines {
250      if self.state.borrow().serial != serial {
251        return;
252      }
253      if let Err(payload) = panic::catch_unwind(AssertUnwindSafe(|| on_line(&line))) {
254        let path = self.path.display();
255        let at = at.map_or_else(String::new, |at| format!(" (called at {at})"));
256        let message = panic_message(payload.as_ref());
257        tracing::error!("socket {path} on_line{at} panicked: {message}");
258        let mut state = self.state.borrow_mut();
259        // Unless it set a new one meanwhile.
260        if state
261          .on_line
262          .as_ref()
263          .is_some_and(|now| Rc::ptr_eq(now, &on_line))
264        {
265          state.on_line = None;
266        }
267        return;
268      }
269    }
270  }
271
272  fn write(self: &Rc<Self>, data: &[u8]) {
273    let mut state = self.state.borrow_mut();
274    let Some(connection) = state
275      .connection
276      .as_mut()
277      .filter(|connection| connection.write_failed.is_none())
278    else {
279      let path = self.path.display();
280      return tracing::warn!(%path, "the socket takes no writes now, so a write was dropped");
281    };
282    connection.outbox.extend_from_slice(data);
283    drop(state);
284    self.flush();
285  }
286
287  /// Writes what the socket takes, and waits for room for the rest.
288  fn flush(self: &Rc<Self>) {
289    let Some(event_loop) = self.shared.as_ref().and_then(|shared| shared.event_loop()) else {
290      return;
291    };
292    let mut state = self.state.borrow_mut();
293    let Some(connection) = state.connection.as_mut() else {
294      return;
295    };
296    let mut broken = None;
297    while !connection.outbox.is_empty() {
298      match connection.stream.write(&connection.outbox) {
299        Ok(written) => {
300          connection.outbox.drain(..written);
301        }
302        Err(error) if error.kind() == ErrorKind::Interrupted => {}
303        Err(error) if error.kind() == ErrorKind::WouldBlock => break,
304        Err(error) => {
305          // The other end may still have sent something: stop writing
306          // and read on to its end.
307          let path = self.path.display();
308          tracing::info!(%path, %error, "the socket takes no more writes");
309          connection.write_failed = Some(error.to_string());
310          connection.outbox = Vec::new();
311          break;
312        }
313      }
314    }
315    if connection.outbox.len() > MAX_OUTBOX {
316      broken = Some(String::from("the other end reads nothing it is sent"));
317    }
318    if broken.is_none() && connection.outbox.is_empty() {
319      // A big write's room isn't kept.
320      connection.outbox.shrink_to(64 * 1024);
321      if let Some(writer) = connection.writer.take() {
322        event_loop.remove(writer);
323      }
324    } else if broken.is_none() && connection.writer.is_none() {
325      match connection.stream.as_fd().try_clone_to_owned() {
326        Ok(fd) => {
327          let inner = Rc::clone(self);
328          let write = Box::new(move || inner.flush());
329          connection.writer = event_loop.insert_fd(fd, Interest::WRITE, write);
330        }
331        Err(error) => broken = Some(error.to_string()),
332      }
333    }
334    drop(state);
335    if let Some(why) = broken {
336      self.drop_connection(&why);
337    }
338  }
339
340  /// Ends the connection the other end ended or broke, and logs why.
341  /// Lines it read before still reach `on_line`.
342  fn drop_connection(&self, why: &str) {
343    if self.state.borrow().connection.is_some() {
344      let path = self.path.display();
345      tracing::info!(%path, why, "the socket closed");
346    }
347    self.end(false);
348  }
349
350  /// Closes it as the config asked, so lines read before are dropped.
351  fn close(&self) {
352    self.end(true);
353  }
354
355  fn end(&self, asked: bool) {
356    let connection = {
357      let mut state = self.state.borrow_mut();
358      state.lines = Lines::default();
359      if asked {
360        state.serial += 1;
361      }
362      state.connection.take()
363    };
364    if let (Some(connection), Some(event_loop)) = (
365      connection,
366      self.shared.as_ref().and_then(|shared| shared.event_loop()),
367    ) {
368      event_loop.remove(connection.reader);
369      if let Some(writer) = connection.writer {
370        event_loop.remove(writer);
371      }
372    }
373    if untrack(|| self.connected.try_with(|connected| *connected)) == Some(true) {
374      self.connected.try_set(false);
375    }
376  }
377}
crabshell_services/socket.rs

crabshell_services/
socket.rs
