Skip to main content

aster/
client.rs

1use std::fs::OpenOptions;
2use std::io::{BufRead, BufReader, Read, Write};
3use std::os::unix::net::UnixStream;
4use std::os::unix::process::CommandExt;
5use std::process::{Command, Stdio};
6use std::thread;
7use std::time::{Duration, Instant};
8
9use anyhow::{Context, Result, bail};
10use fs2::FileExt;
11
12use crate::config::Paths;
13use crate::protocol::{PROTOCOL_VERSION, Request, RequestEnvelope, Response};
14
15const CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
16const IO_TIMEOUT: Duration = Duration::from_secs(1);
17const MAX_RESPONSE_BYTES: u64 = 1024 * 1024;
18
19pub fn request(paths: &Paths, request: Request) -> Result<Response> {
20    match request_once(paths, Request::Ping) {
21        Ok(Response::Pong { .. }) => {}
22        Ok(Response::Error { message }) if message.contains("unsupported protocol version") => {
23            replace_previous_daemon(paths)?;
24        }
25        _ => {
26            start_daemon(paths)?;
27            wait_for_daemon(paths)?;
28        }
29    }
30    request_once(paths, request)
31}
32
33pub fn request_idempotent(paths: &Paths, request: Request) -> Result<Response> {
34    match request_once(paths, request.clone()) {
35        Ok(Response::Error { message }) if message.contains("unsupported protocol version") => {
36            replace_previous_daemon(paths)?;
37            request_once(paths, request)
38        }
39        Ok(response) => Ok(response),
40        Err(_) => {
41            if request_once(paths, Request::Ping).is_err() {
42                start_daemon(paths)?;
43                wait_for_daemon(paths)?;
44            }
45            request_once(paths, request)
46        }
47    }
48}
49
50pub fn request_once(paths: &Paths, request: Request) -> Result<Response> {
51    request_once_version(paths, request, PROTOCOL_VERSION)
52}
53
54fn request_once_version(paths: &Paths, request: Request, version: u32) -> Result<Response> {
55    let mut stream = UnixStream::connect(&paths.socket_file)
56        .with_context(|| format!("failed to connect to {}", paths.socket_file.display()))?;
57    stream.set_read_timeout(Some(IO_TIMEOUT))?;
58    stream.set_write_timeout(Some(IO_TIMEOUT))?;
59
60    let envelope = RequestEnvelope { version, request };
61    serde_json::to_writer(&mut stream, &envelope)?;
62    stream.write_all(b"\n")?;
63    stream.flush()?;
64
65    let mut line = String::new();
66    BufReader::new(stream)
67        .take(MAX_RESPONSE_BYTES + 1)
68        .read_line(&mut line)?;
69    if line.len() as u64 > MAX_RESPONSE_BYTES {
70        bail!("daemon response exceeds 1 MiB");
71    }
72    if line.is_empty() {
73        bail!("daemon closed the connection without a response");
74    }
75    serde_json::from_str(&line).context("daemon returned an invalid response")
76}
77
78fn replace_previous_daemon(paths: &Paths) -> Result<()> {
79    for version in (0..PROTOCOL_VERSION).rev() {
80        if matches!(
81            request_once_version(paths, Request::Shutdown, version),
82            Ok(Response::ShuttingDown)
83        ) {
84            break;
85        }
86    }
87    let deadline = Instant::now() + Duration::from_secs(10);
88    while Instant::now() < deadline {
89        let socket_closed = UnixStream::connect(&paths.socket_file).is_err();
90        if socket_closed && daemon_lock_is_available(paths) {
91            start_daemon(paths)?;
92            return wait_for_daemon(paths);
93        }
94        thread::sleep(Duration::from_millis(20));
95    }
96    bail!("previous Aster daemon did not release its socket and lock")
97}
98
99fn daemon_lock_is_available(paths: &Paths) -> bool {
100    let Ok(file) = OpenOptions::new()
101        .create(true)
102        .read(true)
103        .write(true)
104        .truncate(false)
105        .open(&paths.daemon_lock_file)
106    else {
107        return false;
108    };
109    if file.try_lock_exclusive().is_err() {
110        return false;
111    }
112    let _ = FileExt::unlock(&file);
113    true
114}
115
116fn start_daemon(paths: &Paths) -> Result<()> {
117    paths.ensure_directories()?;
118    let executable = std::env::current_exe().context("failed to locate aster executable")?;
119    let log = OpenOptions::new()
120        .create(true)
121        .append(true)
122        .open(paths.state_dir.join("daemon.log"))?;
123    let error_log = log.try_clone()?;
124
125    let mut command = Command::new(executable);
126    command
127        .arg("daemon")
128        .current_dir(&paths.state_dir)
129        .env("ASTER_CONFIG", &paths.config_file)
130        .env("ASTER_STATE_DIR", &paths.state_dir)
131        .env("ASTER_SOCKET", &paths.socket_file)
132        .stdin(Stdio::null())
133        .stdout(Stdio::from(log))
134        .stderr(Stdio::from(error_log));
135
136    // A new session keeps the shared daemon alive when the shell, tmux client,
137    // or SSH connection that first requested it exits.
138    unsafe {
139        command.pre_exec(|| {
140            if libc::setsid() == -1 {
141                return Err(std::io::Error::last_os_error());
142            }
143            Ok(())
144        });
145    }
146    command.spawn().context("failed to start aster daemon")?;
147    Ok(())
148}
149
150fn wait_for_daemon(paths: &Paths) -> Result<()> {
151    let deadline = Instant::now() + CONNECT_TIMEOUT;
152    while Instant::now() < deadline {
153        if let Ok(Response::Pong { .. }) = request_once(paths, Request::Ping) {
154            return Ok(());
155        }
156        thread::sleep(Duration::from_millis(20));
157    }
158    bail!(
159        "aster daemon did not become ready at {}",
160        paths.socket_file.display()
161    )
162}