返回 CodeWhale
daemon_socket.rs
根目录 / crates / app-server / tests / daemon_socket.rs
1 //! Desktop Phase 0 acceptance for the daemon socket: spawn the daemon, connect
2 //! over the unix socket, complete the attach/claim handshake, round-trip
3 //! requests through the same JSON-RPC dispatcher the stdio transport uses,
4 //! and shut down cleanly (socket file removed, listener gone).
5
6 #![cfg(unix)]
7
8 use std::os::unix::fs::PermissionsExt;
9 use std::path::{Path, PathBuf};
10 use std::sync::atomic::{AtomicU64, Ordering};
11 use std::time::{Duration, SystemTime, UNIX_EPOCH};
12
13 use codewhale_app_server::daemon_socket::{
14 DaemonSocketError, DaemonSocketOptions, bind_daemon_socket,
15 };
16 use serde_json::{Value, json};
17 use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
18 use tokio::net::UnixStream;
19 use tokio::net::unix::{OwnedReadHalf, OwnedWriteHalf};
20 use tokio::task::JoinHandle;
21
22 static NONCE: AtomicU64 = AtomicU64::new(0);
23
24 /// A short, unique socket path: unix socket paths are capped near 100 bytes,
25 /// so `std::env::temp_dir()` (deep under `/var/folders` on macOS) is too long.
26 /// `/tmp` is the same choice the hooks crate's socket test makes.
27 fn short_socket_root(label: &str) -> PathBuf {
28 let millis = SystemTime::now()
29 .duration_since(UNIX_EPOCH)
30 .expect("clock")
31 .as_millis()
32 % 1_000_000;
33 let nonce = NONCE.fetch_add(1, Ordering::Relaxed);
34 let pid = std::process::id();
35 let root = PathBuf::from("/tmp").join(format!("cw-ds-{label}-{pid}-{nonce}-{millis}"));
36 assert!(
37 root.as_os_str().len() < 60,
38 "socket root too long for a unix socket test: {}",
39 root.display()
40 );
41 root
42 }
43
44 struct Harness {
45 root: PathBuf,
46 socket_path: PathBuf,
47 _config_dir: tempfile::TempDir,
48 }
49
50 impl Harness {
51 fn new(label: &str) -> Self {
52 let root = short_socket_root(label);
53 let config_dir = tempfile::tempdir().expect("tempdir");
54 std::fs::write(config_dir.path().join("config.toml"), "").expect("config");
55 Self {
56 socket_path: root.join("run").join("daemon.sock"),
57 root,
58 _config_dir: config_dir,
59 }
60 }
61
62 fn options(&self) -> DaemonSocketOptions {
63 DaemonSocketOptions {
64 socket_path: Some(self.socket_path.clone()),
65 config_path: Some(self._config_dir.path().join("config.toml")),
66 }
67 }
68
69 /// Bind and serve on a background task; returns the serve join handle.
70 async fn spawn_daemon(&self) -> JoinHandle<Result<(), DaemonSocketError>> {
71 let daemon = bind_daemon_socket(self.options())
72 .await
73 .expect("bind daemon socket");
74 assert_eq!(daemon.local_path(), self.socket_path.as_path());
75 tokio::spawn(daemon.serve())
76 }
77 }
78
79 impl Drop for Harness {
80 fn drop(&mut self) {
81 let _ = std::fs::remove_dir_all(&self.root);
82 }
83 }
84
85 struct Client {
86 reader: BufReader<OwnedReadHalf>,
87 writer: OwnedWriteHalf,
88 }
89
90 impl Client {
91 async fn connect(path: &Path) -> Self {
92 let stream = tokio::time::timeout(Duration::from_secs(5), UnixStream::connect(path))
93 .await
94 .expect("connect timeout")
95 .expect("connect");
96 let (rx, writer) = stream.into_split();
97 Self {
98 reader: BufReader::new(rx),
99 writer,
100 }
101 }
102
103 async fn call(&mut self, id: u64, method: &str, params: Value) -> Value {
104 let line = serde_json::to_string(&json!({
105 "jsonrpc": "2.0",
106 "id": id,
107 "method": method,
108 "params": params,
109 }))
110 .expect("encode");
111 self.writer
112 .write_all(format!("{line}\n").as_bytes())
113 .await
114 .expect("write");
115 let mut response = String::new();
116 let read = tokio::time::timeout(
117 Duration::from_secs(10),
118 self.reader.read_line(&mut response),
119 )
120 .await
121 .expect("response timeout")
122 .expect("read");
123 assert!(
124 read > 0,
125 "daemon closed the connection before answering `{method}`"
126 );
127 let value: Value = serde_json::from_str(&response).expect("json response");
128 assert_eq!(value["id"], json!(id), "response id mismatch: {value}");
129 value
130 }
131
132 async fn attach(&mut self, id: u64, name: &str, mode: &str) -> Value {
133 self.call(
134 id,
135 "daemon/attach",
136 json!({ "client": { "name": name, "version": "0.0.0-test", "pid": std::process::id() }, "mode": mode }),
137 )
138 .await
139 }
140
141 /// Read until EOF; proves the daemon closed the socket.
142 async fn wait_for_close(mut self) {
143 let mut sink = String::new();
144 let read = tokio::time::timeout(Duration::from_secs(10), self.reader.read_line(&mut sink))
145 .await
146 .expect("close timeout")
147 .expect("read");
148 assert_eq!(read, 0, "expected EOF, got: {sink}");
149 }
150 }
151
152 async fn wait_for_socket_removed(path: &Path) {
153 tokio::time::timeout(Duration::from_secs(10), async {
154 while path.exists() {
155 tokio::time::sleep(Duration::from_millis(20)).await;
156 }
157 })
158 .await
159 .expect("socket file must be removed on shutdown");
160 }
161
162 #[tokio::test]
163 async fn owner_attaches_round_trips_and_shuts_down_cleanly() {
164 let harness = Harness::new("owner");
165 let server = harness.spawn_daemon().await;
166
167 let socket_mode = std::fs::metadata(&harness.socket_path)
168 .expect("socket metadata")
169 .permissions()
170 .mode()
171 & 0o777;
172 assert_eq!(socket_mode, 0o600, "socket must be private to the user");
173 let dir_mode = std::fs::metadata(harness.socket_path.parent().expect("parent"))
174 .expect("dir metadata")
175 .permissions()
176 .mode()
177 & 0o777;
178 assert_eq!(dir_mode, 0o700, "runtime dir must be private to the user");
179
180 let mut client = Client::connect(&harness.socket_path).await;
181
182 // Anything but healthz before attaching is refused with a typed error:
183 // a read-only probe, a thread/* read, and a prompt run alike.
184 for (id, method, params) in [
185 (1, "capabilities", json!({})),
186 (10, "thread/list", json!({})),
187 (11, "prompt/run", json!({ "prompt": "hi" })),
188 ] {
189 let early = client.call(id, method, params).await;
190 assert_eq!(early["error"]["code"], json!(-32010), "{method}: {early}");
191 assert_eq!(early["error"]["data"]["error"], json!("attach_required"));
192 assert_eq!(early["error"]["data"]["method"], json!(method));
193 }
194
195 // healthz is allowed pre-attach so a shell can probe liveness first.
196 let health = client.call(2, "healthz", json!({})).await;
197 assert_eq!(health["result"]["status"], json!("ok"), "{health}");
198 assert_eq!(health["result"]["transport"], json!("unix-socket"));
199
200 let attached = client.attach(3, "codewhale-desktop", "claim").await;
201 assert_eq!(attached["result"]["attached"], json!(true), "{attached}");
202 assert_eq!(attached["result"]["role"], json!("owner"));
203 assert_eq!(attached["result"]["transport"], json!("unix-socket"));
204 assert_eq!(
205 attached["result"]["daemon"]["pid"],
206 json!(std::process::id())
207 );
208 assert_eq!(
209 attached["result"]["daemon"]["version"],
210 json!(env!("CARGO_PKG_VERSION"))
211 );
212 assert_eq!(
213 attached["result"]["owner"]["name"],
214 json!("codewhale-desktop")
215 );
216 assert_eq!(attached["result"]["connections"], json!(1));
217
218 // Post-attach, the socket transport advertises its own handshake next to
219 // the stdio method set.
220 let advertised = client.call(8, "capabilities", json!({})).await;
221 let methods = advertised["result"]["methods"]
222 .as_array()
223 .expect("methods array");
224 assert_eq!(methods[0], json!("healthz"), "{advertised}");
225 assert_eq!(methods[1], json!("daemon/attach"), "{advertised}");
226 assert!(methods.contains(&json!("shutdown")));
227
228 // Round-trip JSON-RPC requests through the shared dispatcher: `app/*`
229 // methods in, their JSON results out — byte-for-byte the shapes the
230 // stdio transport emits. (No protocol-crate Op/EventMsg envelope is on
231 // this wire; the framing is the stdio transport's newline-delimited
232 // JSON-RPC.)
233 let caps = client.call(4, "app/capabilities", json!({})).await;
234 assert_eq!(caps["result"]["ok"], json!(true), "{caps}");
235 assert!(caps["result"]["data"]["routes"].is_array());
236 let config = client
237 .call(5, "app/config/get", json!({ "key": "model" }))
238 .await;
239 assert_eq!(config["result"]["ok"], json!(true), "{config}");
240 assert_eq!(config["result"]["data"]["key"], json!("model"));
241
242 // A second attach on an attached connection is a typed refusal, not
243 // method_not_found.
244 let again = client.attach(6, "codewhale-desktop", "attach").await;
245 assert_eq!(again["error"]["code"], json!(-32014), "{again}");
246
247 let stopped = client.call(7, "shutdown", json!({})).await;
248 assert_eq!(stopped["result"]["status"], json!("stopped"), "{stopped}");
249
250 let outcome = tokio::time::timeout(Duration::from_secs(10), server)
251 .await
252 .expect("daemon must exit after the owner's shutdown")
253 .expect("join");
254 outcome.expect("serve result");
255 wait_for_socket_removed(&harness.socket_path).await;
256 client.wait_for_close().await;
257 }
258
259 #[tokio::test]
260 async fn guests_share_the_daemon_but_cannot_stop_it() {
261 let harness = Harness::new("guest");
262 let server = harness.spawn_daemon().await;
263
264 let mut owner = Client::connect(&harness.socket_path).await;
265 let claimed = owner.attach(1, "desktop-window-1", "claim").await;
266 assert_eq!(claimed["result"]["role"], json!("owner"), "{claimed}");
267
268 let mut guest = Client::connect(&harness.socket_path).await;
269 let lost = guest.attach(1, "desktop-window-2", "claim").await;
270 assert_eq!(lost["error"]["code"], json!(-32011), "{lost}");
271 assert_eq!(
272 lost["error"]["data"]["owner"]["name"],
273 json!("desktop-window-1")
274 );
275
276 let attached = guest.attach(2, "desktop-window-2", "attach").await;
277 assert_eq!(attached["result"]["role"], json!("attached"), "{attached}");
278 assert_eq!(
279 attached["result"]["owner"]["name"],
280 json!("desktop-window-1")
281 );
282 assert_eq!(attached["result"]["connections"], json!(2));
283
284 let health = guest.call(3, "healthz", json!({})).await;
285 assert_eq!(health["result"]["status"], json!("ok"));
286
287 let refused = guest.call(4, "shutdown", json!({})).await;
288 assert_eq!(refused["error"]["code"], json!(-32012), "{refused}");
289 assert_eq!(refused["error"]["data"]["error"], json!("not_daemon_owner"));
290 assert!(
291 !server.is_finished(),
292 "a guest's shutdown must not stop the daemon"
293 );
294 assert!(harness.socket_path.exists());
295
296 // Once the owner leaves, the slot frees and a relaunched shell can claim.
297 drop(owner);
298 let mut relaunched = Client::connect(&harness.socket_path).await;
299 let reclaimed = tokio::time::timeout(Duration::from_secs(10), async {
300 loop {
301 let response = relaunched.attach(1, "desktop-relaunch", "claim").await;
302 if response.get("result").is_some() {
303 return response;
304 }
305 tokio::time::sleep(Duration::from_millis(20)).await;
306 }
307 })
308 .await
309 .expect("owner slot must free when the owner disconnects");
310 assert_eq!(reclaimed["result"]["role"], json!("owner"), "{reclaimed}");
311
312 // The guest is still attached and served while the new owner is in.
313 let health = guest.call(5, "healthz", json!({})).await;
314 assert_eq!(health["result"]["status"], json!("ok"));
315
316 let stopped = relaunched.call(2, "shutdown", json!({})).await;
317 assert_eq!(stopped["result"]["status"], json!("stopped"));
318 tokio::time::timeout(Duration::from_secs(10), server)
319 .await
320 .expect("daemon exits")
321 .expect("join")
322 .expect("serve result");
323 wait_for_socket_removed(&harness.socket_path).await;
324 // The owner's shutdown closes every other connection, not just its own.
325 guest.wait_for_close().await;
326 relaunched.wait_for_close().await;
327 }
328
329 #[tokio::test]
330 async fn version_skew_is_refused_at_attach() {
331 let harness = Harness::new("skew");
332 let server = harness.spawn_daemon().await;
333 let mut client = Client::connect(&harness.socket_path).await;
334 let refused = client
335 .call(
336 1,
337 "daemon/attach",
338 json!({ "client": { "name": "old-desktop" }, "expect_daemon_version": "0.0.1-other" }),
339 )
340 .await;
341 assert_eq!(refused["error"]["code"], json!(-32013), "{refused}");
342 assert_eq!(
343 refused["error"]["data"]["actual"],
344 json!(env!("CARGO_PKG_VERSION"))
345 );
346
347 let daemon = bind_daemon_socket(harness.options()).await;
348 // Meanwhile the original daemon is live, so a second bind must refuse.
349 match daemon {
350 Err(DaemonSocketError::AlreadyRunning { path }) => {
351 assert_eq!(path, harness.socket_path);
352 }
353 Err(other) => panic!("unexpected error: {other}"),
354 Ok(_) => panic!("second daemon must not replace a live socket"),
355 }
356 server.abort();
357 let _ = server.await;
358 }
359
360 #[tokio::test]
361 async fn stale_socket_is_cleaned_up_and_foreign_files_are_refused() {
362 let harness = Harness::new("stale");
363 std::fs::create_dir_all(harness.socket_path.parent().expect("parent")).expect("mkdir");
364
365 // A socket file whose listener is gone: bind must reclaim it.
366 {
367 let dead = tokio::net::UnixListener::bind(&harness.socket_path).expect("bind dead");
368 drop(dead);
369 }
370 assert!(
371 harness.socket_path.exists(),
372 "dropping a listener leaves the file"
373 );
374 let daemon = bind_daemon_socket(harness.options())
375 .await
376 .expect("stale socket must be reclaimed");
377 let handle = daemon.shutdown_handle();
378 let server = tokio::spawn(daemon.serve());
379 let mut client = Client::connect(&harness.socket_path).await;
380 let health = client.call(1, "healthz", json!({})).await;
381 assert_eq!(health["result"]["status"], json!("ok"));
382 handle.trigger();
383 tokio::time::timeout(Duration::from_secs(10), server)
384 .await
385 .expect("daemon exits on handle")
386 .expect("join")
387 .expect("serve result");
388 wait_for_socket_removed(&harness.socket_path).await;
389
390 // A regular file at the path is never deleted.
391 std::fs::write(&harness.socket_path, b"not a socket").expect("write file");
392 match bind_daemon_socket(harness.options()).await {
393 Err(DaemonSocketError::NotASocket { path }) => assert_eq!(path, harness.socket_path),
394 Err(other) => panic!("unexpected error: {other}"),
395 Ok(_) => panic!("must refuse to replace a non-socket"),
396 }
397 assert_eq!(
398 std::fs::read(&harness.socket_path).expect("file intact"),
399 b"not a socket"
400 );
401 }
402
402 lines RUST