1use 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
54fn 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 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 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 _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;