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