Krolyxon
Commit e515e5a

Branches

Tags

Refs →

remagesrcipc.rs

557 lines · 18.0 KB · Rust

Raw Blame History Download
1//! Unix-socket server speaking newline-delimited JSON, with event subscriptions, and the
2//! `remage msg` client.
3//!
4//! Requests: `{"action":"goto 5"}` (or an array of actions), `{"query":"current"}` and
5//! `{"subscribe":["image","marks","mode"]}`, which turns the connection into an event stream.
6//! Replies: `{"ok":true,"data":…}` or `{"ok":false,"error":"…"}`; events:
7//! `{"event":"image","data":…}`.
8
9use std::io::{BufRead, BufReader, Read, Write};
10use std::os::unix::fs::{DirBuilderExt, PermissionsExt};
11use std::os::unix::net::{UnixListener, UnixStream};
12use std::path::{Path, PathBuf};
13use std::sync::Arc;
14use std::sync::atomic::{AtomicUsize, Ordering};
15use std::time::{Duration, Instant, SystemTime};
16
17use crossbeam_channel::{Receiver, Sender};
18use serde_json::{Value, json};
19
20use crate::action::Action;
21
22/// Longest request line accepted.
23const MAX_REQUEST: u64 = 1 << 20;
24/// How long a connection waits for the event loop to answer.
25const REPLY_TIMEOUT: Duration = Duration::from_secs(5);
26
27#[derive(Clone, Copy, Debug, PartialEq, Eq)]
28pub enum Query {
29 Current,
30 Marked,
31 Queue,
32 State,
33}
34
35#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
36pub enum Topic {
37 Image,
38 Marks,
39 Mode,
40}
41
42impl Topic {
43 pub fn name(self) -> &'static str {
44 match self {
45 Topic::Image => "image",
46 Topic::Marks => "marks",
47 Topic::Mode => "mode",
48 }
49 }
50}
51
52#[derive(Clone, Debug, PartialEq)]
53pub enum Request {
54 Actions(Vec<Action>),
55 Query(Query),
56 Subscribe(Vec<Topic>),
57}
58
59fn word<T: Copy>(what: &str, s: &str, opts: &[(&str, T)]) -> Result<T, String> {
60 opts.iter()
61 .find(|(k, _)| *k == s)
62 .map(|(_, v)| *v)
63 .ok_or_else(|| {
64 let names: Vec<_> = opts.iter().map(|(k, _)| *k).collect();
65 format!("unknown {what} `{s}`: expected {}", names.join(", "))
66 })
67}
68
69const QUERIES: &[(&str, Query)] = &[
70 ("current", Query::Current),
71 ("marked", Query::Marked),
72 ("queue", Query::Queue),
73 ("state", Query::State),
74];
75
76const TOPICS: &[(&str, Topic)] = &[
77 ("image", Topic::Image),
78 ("marks", Topic::Marks),
79 ("mode", Topic::Mode),
80];
81
82/// A string or an array of strings.
83fn strings(v: &Value, what: &str) -> Result<Vec<String>, String> {
84 match v {
85 Value::String(s) => Ok(vec![s.clone()]),
86 Value::Array(a) => a
87 .iter()
88 .map(|e| {
89 e.as_str()
90 .map(str::to_string)
91 .ok_or_else(|| format!("`{what}` must be a string or an array of strings"))
92 })
93 .collect(),
94 _ => Err(format!("`{what}` must be a string or an array of strings")),
95 }
96}
97
98pub fn parse_request(line: &str) -> Result<Request, String> {
99 let v: Value = serde_json::from_str(line).map_err(|e| format!("invalid JSON: {e}"))?;
100 let Value::Object(obj) = v else {
101 return Err("expected a JSON object".into());
102 };
103 let mut keys = obj.keys();
104 let (Some(key), None) = (keys.next(), keys.next()) else {
105 return Err("expected exactly one of `action`, `query`, `subscribe`".into());
106 };
107 let val = &obj[key];
108 match key.as_str() {
109 "action" => {
110 let actions = strings(val, "action")?
111 .iter()
112 .map(|s| s.parse::<Action>().map_err(|e| e.to_string()))
113 .collect::<Result<Vec<_>, _>>()?;
114 if actions.is_empty() {
115 return Err("no actions".into());
116 }
117 Ok(Request::Actions(actions))
118 }
119 "query" => {
120 let q = val.as_str().ok_or("`query` must be a string")?;
121 Ok(Request::Query(word("query", q, QUERIES)?))
122 }
123 "subscribe" => {
124 let mut topics = Vec::new();
125 for t in strings(val, "subscribe")? {
126 let t = word("event", &t, TOPICS)?;
127 if !topics.contains(&t) {
128 topics.push(t);
129 }
130 }
131 if topics.is_empty() {
132 return Err("no events to subscribe to".into());
133 }
134 Ok(Request::Subscribe(topics))
135 }
136 other => Err(format!(
137 "unknown request `{other}`: expected `action`, `query` or `subscribe`"
138 )),
139 }
140}
141
142// Envelopes are written by hand so `ok` and `event` come first (serde_json sorts keys).
143
144pub fn ok(data: Value) -> String {
145 if data.is_null() {
146 r#"{"ok":true}"#.into()
147 } else {
148 format!(r#"{{"ok":true,"data":{data}}}"#)
149 }
150}
151
152pub fn error(msg: &str) -> String {
153 format!(r#"{{"ok":false,"error":{}}}"#, Value::from(msg))
154}
155
156pub fn event(topic: Topic, data: Value) -> String {
157 format!(r#"{{"event":"{}","data":{data}}}"#, topic.name())
158}
159
160pub fn path_json(p: &Path) -> Value {
161 Value::String(p.to_string_lossy().into_owned())
162}
163
164/// `$XDG_RUNTIME_DIR/remage`, or a per-user directory in `/tmp`.
165pub fn socket_dir() -> PathBuf {
166 match std::env::var_os("XDG_RUNTIME_DIR").filter(|d| !d.is_empty()) {
167 Some(d) => PathBuf::from(d).join("remage"),
168 None => {
169 let user = std::env::var("USER").unwrap_or_else(|_| "user".into());
170 std::env::temp_dir().join(format!("remage-{user}"))
171 }
172 }
173}
174
175fn pid_of(path: &Path) -> Option<u32> {
176 path.file_name()?
177 .to_str()?
178 .strip_suffix(".sock")?
179 .parse()
180 .ok()
181}
182
183fn alive(pid: u32) -> bool {
184 Path::new("/proc").join(pid.to_string()).exists()
185}
186
187/// Live instances' sockets, newest first.
188fn instances(dir: &Path) -> Vec<PathBuf> {
189 let Ok(entries) = std::fs::read_dir(dir) else {
190 return Vec::new();
191 };
192 let mut found: Vec<(SystemTime, PathBuf)> = entries
193 .flatten()
194 .map(|e| e.path())
195 .filter(|p| pid_of(p).is_some_and(alive))
196 .map(|p| {
197 let t = std::fs::metadata(&p)
198 .and_then(|m| m.modified())
199 .unwrap_or(SystemTime::UNIX_EPOCH);
200 (t, p)
201 })
202 .collect();
203 found.sort_by_key(|f| std::cmp::Reverse(f.0));
204 found.into_iter().map(|(_, p)| p).collect()
205}
206
207/// A request from a connection, with where its reply goes.
208pub struct Incoming {
209 pub request: Request,
210 pub reply: Sender<String>,
211 /// For subscriptions: where events go after the reply.
212 pub events: Option<Sender<String>>,
213}
214
215/// The listening socket. Dropping it removes the socket file once replies have gone out.
216pub struct Server {
217 path: PathBuf,
218 /// Replies handed to connection threads but not yet written.
219 in_flight: Arc<AtomicUsize>,
220}
221
222impl Server {
223 pub fn path(&self) -> &Path {
224 &self.path
225 }
226}
227
228impl Drop for Server {
229 fn drop(&mut self) {
230 // A `quit` sent over the socket still gets its reply.
231 let until = Instant::now() + Duration::from_millis(200);
232 while self.in_flight.load(Ordering::Acquire) > 0 && Instant::now() < until {
233 std::thread::sleep(Duration::from_millis(2));
234 }
235 let _ = std::fs::remove_file(&self.path);
236 }
237}
238
239/// Listen on `<socket_dir>/<pid>.sock`; `deliver` hands each request to the event loop.
240pub fn serve(deliver: impl Fn(Incoming) -> bool + Send + Sync + 'static) -> anyhow::Result<Server> {
241 let dir = socket_dir();
242 std::fs::DirBuilder::new()
243 .recursive(true)
244 .mode(0o700)
245 .create(&dir)?;
246 // Sockets of instances that died without cleaning up.
247 if let Ok(entries) = std::fs::read_dir(&dir) {
248 for p in entries.flatten().map(|e| e.path()) {
249 if pid_of(&p).is_some_and(|pid| !alive(pid)) {
250 let _ = std::fs::remove_file(&p);
251 }
252 }
253 }
254 let path = dir.join(format!("{}.sock", std::process::id()));
255 let _ = std::fs::remove_file(&path);
256 let listener = UnixListener::bind(&path)?;
257 std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600))?;
258 let in_flight = Arc::new(AtomicUsize::new(0));
259 let deliver = Arc::new(deliver);
260 let flight = in_flight.clone();
261 std::thread::Builder::new()
262 .name("remage-ipc".into())
263 .spawn(move || {
264 for stream in listener.incoming() {
265 let Ok(stream) = stream else { continue };
266 let deliver = deliver.clone();
267 let flight = flight.clone();
268 let _ = std::thread::Builder::new()
269 .name("remage-ipc-conn".into())
270 .spawn(move || connection(stream, &*deliver, &flight));
271 }
272 })?;
273 Ok(Server { path, in_flight })
274}
275
276fn connection(stream: UnixStream, deliver: &dyn Fn(Incoming) -> bool, in_flight: &AtomicUsize) {
277 let Ok(read) = stream.try_clone() else { return };
278 let mut reader = BufReader::new(read.take(u64::MAX));
279 let mut out = stream;
280 let mut line = String::new();
281 loop {
282 line.clear();
283 reader.get_mut().set_limit(MAX_REQUEST);
284 match reader.read_line(&mut line) {
285 Ok(0) | Err(_) => return,
286 Ok(_) if !line.ends_with('\n') && line.len() as u64 >= MAX_REQUEST => {
287 let _ = writeln!(out, "{}", error("request too long"));
288 return;
289 }
290 Ok(_) => {}
291 }
292 if line.trim().is_empty() {
293 continue;
294 }
295 let request = match parse_request(line.trim()) {
296 Ok(r) => r,
297 Err(e) => {
298 if writeln!(out, "{}", error(&e)).is_err() {
299 return;
300 }
301 continue;
302 }
303 };
304 let subscribe = matches!(request, Request::Subscribe(_));
305 let (reply_tx, reply_rx) = crossbeam_channel::bounded(1);
306 let (events_tx, events_rx) = if subscribe {
307 let (tx, rx) = crossbeam_channel::unbounded();
308 (Some(tx), Some(rx))
309 } else {
310 (None, None)
311 };
312 in_flight.fetch_add(1, Ordering::AcqRel);
313 let sent = deliver(Incoming {
314 request,
315 reply: reply_tx,
316 events: events_tx,
317 });
318 let reply = if sent {
319 reply_rx
320 .recv_timeout(REPLY_TIMEOUT)
321 .unwrap_or_else(|_| error("no reply from the viewer"))
322 } else {
323 error("the viewer is shutting down")
324 };
325 let written = writeln!(out, "{reply}").is_ok();
326 in_flight.fetch_sub(1, Ordering::AcqRel);
327 if !written {
328 return;
329 }
330 if let Some(rx) = events_rx {
331 stream_events(&mut out, &rx);
332 return;
333 }
334 }
335}
336
337fn stream_events(out: &mut UnixStream, rx: &Receiver<String>) {
338 while let Ok(ev) = rx.recv() {
339 if writeln!(out, "{ev}").is_err() {
340 return;
341 }
342 }
343}
344
345/// `remage msg` arguments.
346pub struct ClientArgs {
347 pub pid: Option<u32>,
348 pub query: Option<String>,
349 pub subscribe: Vec<String>,
350 pub action: Vec<String>,
351}
352
353fn find_socket(pid: Option<u32>) -> Result<PathBuf, String> {
354 if let Some(pid) = pid {
355 let p = socket_dir().join(format!("{pid}.sock"));
356 return if p.exists() {
357 Ok(p)
358 } else {
359 Err(format!(
360 "no remage with pid {pid} ({} missing)",
361 p.display()
362 ))
363 };
364 }
365 if let Some(s) = std::env::var_os("REMAGE_SOCKET").filter(|s| !s.is_empty()) {
366 return Ok(PathBuf::from(s));
367 }
368 instances(&socket_dir())
369 .into_iter()
370 .next()
371 .ok_or_else(|| "no running remage found".into())
372}
373
374/// Run the client; returns the process exit code.
375pub fn client(args: ClientArgs) -> i32 {
376 match client_inner(args) {
377 Ok(code) => code,
378 Err(e) => {
379 eprintln!("remage msg: {e}");
380 1
381 }
382 }
383}
384
385fn client_inner(args: ClientArgs) -> Result<i32, String> {
386 let request = if let Some(q) = &args.query {
387 json!({ "query": q })
388 } else if !args.subscribe.is_empty() {
389 let topics: Vec<&str> = args
390 .subscribe
391 .iter()
392 .flat_map(|s| s.split(','))
393 .map(str::trim)
394 .filter(|s| !s.is_empty())
395 .collect();
396 json!({ "subscribe": topics })
397 } else if !args.action.is_empty() {
398 json!({ "action": args.action.join(" ") })
399 } else {
400 return Err("nothing to send: give an action, --query or --subscribe".into());
401 };
402 let path = find_socket(args.pid)?;
403 let mut stream = UnixStream::connect(&path).map_err(|e| format!("{}: {e}", path.display()))?;
404 writeln!(stream, "{request}").map_err(|e| e.to_string())?;
405 let mut reader = BufReader::new(stream);
406 let mut line = String::new();
407 if reader.read_line(&mut line).map_err(|e| e.to_string())? == 0 {
408 return Err("the viewer closed the connection".into());
409 }
410 let reply: Value = serde_json::from_str(&line).map_err(|e| format!("bad reply: {e}"))?;
411 if reply["ok"] != Value::Bool(true) {
412 let msg = reply["error"].as_str().unwrap_or("request failed");
413 return Err(msg.to_string());
414 }
415 let stdout = std::io::stdout();
416 let mut out = stdout.lock();
417 if !args.subscribe.is_empty() {
418 loop {
419 line.clear();
420 match reader.read_line(&mut line) {
421 Ok(0) | Err(_) => return Ok(0),
422 Ok(_) => {
423 if out.write_all(line.as_bytes()).is_err() || out.flush().is_err() {
424 return Ok(0);
425 }
426 }
427 }
428 }
429 }
430 print_data(&mut out, &reply["data"]).map_err(|e| e.to_string())?;
431 Ok(0)
432}
433
434/// Strings and lists of strings print one per line, for shell pipelines; the rest as JSON.
435fn print_data(out: &mut impl Write, data: &Value) -> std::io::Result<()> {
436 match data {
437 Value::Null => Ok(()),
438 Value::String(s) => writeln!(out, "{s}"),
439 Value::Array(a) if a.iter().all(Value::is_string) => {
440 for s in a.iter().filter_map(Value::as_str) {
441 writeln!(out, "{s}")?;
442 }
443 Ok(())
444 }
445 other => writeln!(out, "{other}"),
446 }
447}
448
449#[cfg(test)]
450mod tests {
451 use super::*;
452
453 #[test]
454 fn requests_parse() {
455 assert_eq!(
456 parse_request(r#"{"action":"goto 5"}"#),
457 Ok(Request::Actions(vec![Action::Goto(5)]))
458 );
459 assert_eq!(
460 parse_request(r#"{"action":["next","mark on"]}"#),
461 Ok(Request::Actions(vec![
462 Action::Next(1),
463 Action::Mark(crate::action::Switch::On)
464 ]))
465 );
466 assert_eq!(
467 parse_request(r#"{"query":"marked"}"#),
468 Ok(Request::Query(Query::Marked))
469 );
470 assert_eq!(
471 parse_request(r#"{"subscribe":["image","mode","image"]}"#),
472 Ok(Request::Subscribe(vec![Topic::Image, Topic::Mode]))
473 );
474 assert_eq!(
475 parse_request(r#"{"subscribe":"marks"}"#),
476 Ok(Request::Subscribe(vec![Topic::Marks]))
477 );
478 let e = |s: &str| parse_request(s).unwrap_err();
479 assert!(e("nope").contains("invalid JSON"));
480 assert!(e("[]").contains("object"));
481 assert!(e(r#"{"action":"nxt"}"#).contains("did you mean `next`"));
482 assert!(e(r#"{"query":"all"}"#).contains("expected current"));
483 assert!(e(r#"{"action":"next","query":"state"}"#).contains("exactly one"));
484 assert!(e(r#"{"subscribe":[]}"#).contains("no events"));
485 assert!(e(r#"{"frob":1}"#).contains("unknown request"));
486 }
487
488 #[test]
489 fn replies() {
490 assert_eq!(ok(Value::Null), r#"{"ok":true}"#);
491 assert_eq!(ok(json!(3)), r#"{"ok":true,"data":3}"#);
492 assert_eq!(error("x"), r#"{"ok":false,"error":"x"}"#);
493 assert_eq!(
494 event(Topic::Mode, json!("grid")),
495 r#"{"event":"mode","data":"grid"}"#
496 );
497 let mut out = Vec::new();
498 print_data(&mut out, &json!(["a", "b"])).unwrap();
499 print_data(&mut out, &json!({"k": 1})).unwrap();
500 assert_eq!(String::from_utf8(out).unwrap(), "a\nb\n{\"k\":1}\n");
501 }
502
503 #[test]
504 fn socket_names() {
505 assert_eq!(pid_of(Path::new("/run/remage/123.sock")), Some(123));
506 assert_eq!(pid_of(Path::new("/run/remage/x.sock")), None);
507 assert!(alive(std::process::id()));
508 }
509
510 #[test]
511 fn round_trip_over_a_socket() {
512 let dir = std::env::temp_dir().join(format!("remage-ipc-{}", std::process::id()));
513 std::fs::create_dir_all(&dir).unwrap();
514 let path = dir.join("t.sock");
515 let _ = std::fs::remove_file(&path);
516 let listener = UnixListener::bind(&path).unwrap();
517 let flight = Arc::new(AtomicUsize::new(0));
518 let f = flight.clone();
519 std::thread::spawn(move || {
520 let (s, _) = listener.accept().unwrap();
521 connection(
522 s,
523 &|inc: Incoming| {
524 let reply = match inc.request {
525 Request::Query(Query::Current) => ok(json!("/a.png")),
526 Request::Subscribe(_) => {
527 let ev = inc.events.unwrap();
528 let _ = ev.send(event(Topic::Mode, json!("grid")));
529 ok(Value::Null)
530 }
531 _ => error("unexpected"),
532 };
533 inc.reply.send(reply).is_ok()
534 },
535 &f,
536 );
537 });
538 let mut s = UnixStream::connect(&path).unwrap();
539 let mut r = BufReader::new(s.try_clone().unwrap());
540 let mut line = String::new();
541 writeln!(s, "{{\"query\":\"current\"}}\n\nbad").unwrap();
542 r.read_line(&mut line).unwrap();
543 assert_eq!(line.trim(), r#"{"ok":true,"data":"/a.png"}"#);
544 line.clear();
545 r.read_line(&mut line).unwrap();
546 assert!(line.contains("invalid JSON"), "{line}");
547 writeln!(s, "{{\"subscribe\":\"mode\"}}").unwrap();
548 line.clear();
549 r.read_line(&mut line).unwrap();
550 assert_eq!(line.trim(), r#"{"ok":true}"#);
551 line.clear();
552 r.read_line(&mut line).unwrap();
553 assert_eq!(line.trim(), r#"{"event":"mode","data":"grid"}"#);
554 assert_eq!(flight.load(Ordering::Acquire), 0);
555 std::fs::remove_dir_all(&dir).unwrap();
556 }
557}