Skip to main content

cherenkov/
control.rs

1//! Local control: one bounded JSON exchange per connection, authenticated by UID.
2
3use crate::units::BYTES_PER_KIB;
4use anyhow::{Context, Result, ensure};
5use serde::{Deserialize, Serialize};
6use serde_json::{Value, json};
7use std::fs::{File, OpenOptions};
8use std::io::{Read, Write};
9use std::os::fd::AsRawFd;
10use std::os::unix::fs::{DirBuilderExt, FileTypeExt, MetadataExt, OpenOptionsExt, PermissionsExt};
11use std::os::unix::net::{UnixListener, UnixStream};
12use std::path::{Path, PathBuf};
13use std::sync::{
14    Arc,
15    atomic::{AtomicBool, Ordering},
16};
17use std::thread::JoinHandle;
18use std::time::{Duration, Instant};
19
20mod activity;
21pub mod state;
22pub mod stats;
23pub use state::State;
24
25#[derive(Clone, Debug, Serialize, Deserialize)]
26#[serde(tag = "op", rename_all = "snake_case", deny_unknown_fields)]
27pub enum Command {
28    Status,
29    StatsSummary,
30    StatsLayers {
31        #[serde(default)]
32        offset: usize,
33        #[serde(default = "page_limit")]
34        limit: usize,
35    },
36    StatsExperts {
37        layer: usize,
38        #[serde(default)]
39        offset: usize,
40        #[serde(default = "page_limit")]
41        limit: usize,
42    },
43    ConfigShow,
44    ConfigReload,
45}
46
47fn page_limit() -> usize {
48    64
49}
50
51const MAX_FRAME: usize = 64 * BYTES_PER_KIB;
52const TIMEOUT: Duration = Duration::from_secs(2);
53
54/// The deadline applies to the whole frame, even if a client dribbles bytes.
55fn read_frame(stream: &mut UnixStream) -> Result<Vec<u8>> {
56    let deadline = Instant::now() + TIMEOUT;
57    let mut bytes = Vec::new();
58
59    loop {
60        wait_readable(stream, deadline)?;
61
62        let mut chunk = [0; BYTES_PER_KIB];
63        let n = stream.read(&mut chunk).context("reading control frame")?;
64
65        ensure!(n > 0, "incomplete control frame");
66
67        let end = chunk[..n].iter().position(|&b| b == b'\n');
68
69        bytes.extend_from_slice(&chunk[..end.unwrap_or(n)]);
70        ensure!(bytes.len() <= MAX_FRAME, "control frame exceeds 64 KiB");
71
72        if end.is_some() {
73            return Ok(bytes);
74        }
75    }
76}
77
78fn wait_readable(stream: &UnixStream, deadline: Instant) -> Result<()> {
79    // macOS rejects timeout changes after peer close, even with unread data.
80    // Poll also bounds the entire frame. This connection has only one reader.
81    let mut fd = libc::pollfd {
82        fd: stream.as_raw_fd(),
83        events: libc::POLLIN,
84        revents: 0,
85    };
86
87    loop {
88        let remaining = deadline
89            .checked_duration_since(Instant::now())
90            .context("control request timed out")?;
91        // TIMEOUT is two seconds; round up to poll's millisecond precision.
92        let ready = unsafe { libc::poll(&mut fd, 1, remaining.as_millis() as i32 + 1) };
93
94        if ready > 0 {
95            return Ok(());
96        }
97
98        ensure!(ready != 0, "control request timed out");
99
100        let error = std::io::Error::last_os_error();
101
102        if error.kind() == std::io::ErrorKind::Interrupted {
103            continue;
104        }
105
106        return Err(error).context("polling control socket");
107    }
108}
109
110fn write_frame(stream: &mut UnixStream, value: &Value) -> Result<()> {
111    let mut bytes = serde_json::to_vec(value)?;
112
113    ensure!(bytes.len() <= MAX_FRAME, "control response exceeds 64 KiB");
114    bytes.push(b'\n');
115    stream
116        .set_write_timeout(Some(TIMEOUT))
117        .context("setting control write timeout")?;
118    stream.write_all(&bytes).context("writing control frame")?;
119
120    Ok(())
121}
122
123fn same_user(stream: &UnixStream) -> Result<()> {
124    let (mut uid, mut gid) = (0, 0);
125    let result = unsafe { libc::getpeereid(stream.as_raw_fd(), &mut uid, &mut gid) };
126
127    ensure!(result == 0, "cannot determine control peer identity");
128    ensure!(
129        uid == unsafe { libc::geteuid() },
130        "control peer must have the server's UID"
131    );
132
133    Ok(())
134}
135
136pub fn query(socket: &Path, command: Command) -> Result<Value> {
137    let mut stream = UnixStream::connect(socket)
138        .with_context(|| format!("connecting to {}", socket.display()))?;
139
140    same_user(&stream)?;
141    write_frame(&mut stream, &serde_json::to_value(command)?)?;
142
143    let response: Value = serde_json::from_slice(&read_frame(&mut stream)?)?;
144
145    ensure!(
146        response["ok"] == true,
147        "{}",
148        response["error"]
149            .as_str()
150            .unwrap_or("invalid control response")
151    );
152
153    Ok(response["data"].clone())
154}
155
156fn handle(mut stream: UnixStream, state: &State) -> Result<()> {
157    same_user(&stream)?;
158
159    let result = (|| -> Result<Value> {
160        Ok(
161            match serde_json::from_slice::<Command>(&read_frame(&mut stream)?)? {
162                Command::Status => state.status(),
163                Command::StatsSummary => state.summary()?,
164                Command::StatsLayers { offset, limit } => state.layers(offset, limit)?,
165                Command::StatsExperts {
166                    layer,
167                    offset,
168                    limit,
169                } => state.experts(layer, offset, limit)?,
170                Command::ConfigShow => state.show_config(),
171                Command::ConfigReload => state.reload()?,
172            },
173        )
174    })();
175    let (data, error) = match result {
176        Ok(data) => (Some(data), None),
177        Err(error) => (None, Some(format!("{error:#}"))),
178    };
179
180    write_frame(
181        &mut stream,
182        &json!({"ok": error.is_none(), "data": data, "error": error}),
183    )
184}
185
186pub struct Listener {
187    socket: PathBuf,
188    identity: (u64, u64),
189    // Held until the socket is removed; serializes startup and stale cleanup.
190    _lock: File,
191    stopped: Arc<AtomicBool>,
192    thread: Option<JoinHandle<()>>,
193}
194
195impl Listener {
196    pub fn start(socket: &Path, state: Arc<State>) -> Result<Self> {
197        let parent = socket.parent().context("socket needs a parent directory")?;
198
199        match std::fs::DirBuilder::new().mode(0o700).create(parent) {
200            Ok(()) => {}
201            Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => {}
202            Err(e) => return Err(e).context("creating control directory"),
203        }
204
205        let uid = unsafe { libc::geteuid() };
206        let dir = std::fs::symlink_metadata(parent)?;
207
208        ensure!(
209            dir.is_dir() && dir.uid() == uid && dir.mode() & 0o777 == 0o700,
210            "control directory must be owned by this user, not a symlink, and mode 0700: {}",
211            parent.display()
212        );
213
214        let lock = OpenOptions::new()
215            .read(true)
216            .write(true)
217            .create(true)
218            .truncate(false)
219            .mode(0o600)
220            .custom_flags(libc::O_NOFOLLOW)
221            .open(socket.with_extension("lock"))?;
222        let meta = lock.metadata()?;
223
224        ensure!(
225            meta.is_file() && meta.uid() == uid && meta.mode() & 0o777 == 0o600,
226            "unsafe control lock file"
227        );
228        ensure!(
229            unsafe { libc::flock(lock.as_raw_fd(), libc::LOCK_EX | libc::LOCK_NB) } == 0,
230            "another server owns control socket {}",
231            socket.display()
232        );
233
234        match std::fs::symlink_metadata(socket) {
235            Ok(meta) => {
236                ensure!(
237                    meta.file_type().is_socket() && meta.uid() == uid,
238                    "refusing to replace non-socket control path"
239                );
240
241                match UnixStream::connect(socket) {
242                    Err(e) if e.kind() == std::io::ErrorKind::ConnectionRefused => {
243                        std::fs::remove_file(socket)?
244                    }
245                    _ => anyhow::bail!("control socket already in use: {}", socket.display()),
246                }
247            }
248            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
249            Err(e) => return Err(e.into()),
250        }
251
252        let listener = UnixListener::bind(socket).context("binding control socket")?;
253
254        std::fs::set_permissions(socket, std::fs::Permissions::from_mode(0o600))?;
255
256        let meta = std::fs::symlink_metadata(socket)?;
257        let stopped = Arc::new(AtomicBool::new(false));
258        let stop = stopped.clone();
259        let thread = std::thread::spawn(move || {
260            for stream in listener.incoming() {
261                if stop.load(Ordering::Acquire) {
262                    break;
263                }
264
265                if let Ok(stream) = stream {
266                    let _ = handle(stream, &state);
267                }
268            }
269        });
270
271        Ok(Self {
272            socket: socket.to_owned(),
273            identity: (meta.dev(), meta.ino()),
274            _lock: lock,
275            stopped,
276            thread: Some(thread),
277        })
278    }
279}
280
281impl Drop for Listener {
282    fn drop(&mut self) {
283        self.stopped.store(true, Ordering::Release);
284
285        let _ = UnixStream::connect(&self.socket);
286
287        if let Some(thread) = self.thread.take() {
288            let _ = thread.join();
289        }
290
291        if let Ok(meta) = std::fs::symlink_metadata(&self.socket)
292            && (meta.dev(), meta.ino()) == self.identity
293        {
294            let _ = std::fs::remove_file(&self.socket);
295        }
296    }
297}
298
299#[cfg(test)]
300#[path = "../tests/unit/control.rs"]
301mod tests;