a3s_box_runtime/pool/
client.rs

1//! Socket protocol and client helpers for the warm-pool daemon.
2
3use a3s_box_core::error::{BoxError, Result};
4use serde::{Deserialize, Serialize};
5
6/// Maximum encoded payload accepted by the pool socket protocol.
7///
8/// Pool responses can contain command output, so the limit is intentionally
9/// larger than a typical request while still preventing an untrusted peer from
10/// forcing an arbitrary allocation from the length prefix.
11pub const MAX_FRAME_SIZE: usize = 16 * 1024 * 1024;
12
13/// Bound how long a pool socket read may wait for a complete frame. VM boot and
14/// command execution happen after the request frame has been read, so a peer
15/// that connects and sends only a partial frame cannot pin a task indefinitely.
16pub const FRAME_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(15);
17
18/// Wire protocol for the `pool` Unix socket.
19///
20/// Client→daemon request: run a one-shot command, query status, stop the
21/// daemon, or manage a short-lived leased VM session.
22#[derive(Serialize, Deserialize)]
23#[serde(tag = "op", rename_all = "snake_case")]
24pub enum PoolRequest {
25    Run(PoolRunRequest),
26    Status,
27    Stop,
28    Lease(PoolLeaseRequest),
29    Exec(PoolLeaseExecRequest),
30    Release(PoolLeaseReleaseRequest),
31}
32
33#[derive(Serialize, Deserialize)]
34pub struct PoolRunRequest {
35    /// Image to run in; `None` means use the daemon's default image.
36    #[serde(default)]
37    pub image: Option<String>,
38    /// User to run as (`uid[:gid]` or name); `None` runs as the image default.
39    #[serde(default)]
40    pub user: Option<String>,
41    /// Working directory inside the sandbox.
42    #[serde(default)]
43    pub workdir: Option<String>,
44    /// Optional guest-visible rootfs to chroot into before executing.
45    #[serde(default)]
46    pub rootfs: Option<String>,
47    /// Extra KEY=VALUE environment entries.
48    #[serde(default)]
49    pub env: Vec<String>,
50    /// Boot-time volume specs for this sandbox pool.
51    #[serde(default)]
52    pub volumes: Vec<String>,
53    /// Boot-time vCPU count for lazily-created pools.
54    #[serde(default)]
55    pub vcpus: Option<u32>,
56    /// Boot-time memory size for lazily-created pools.
57    #[serde(default)]
58    pub memory_mb: Option<u32>,
59    /// Force exec mode for this request.
60    #[serde(default)]
61    pub exec: bool,
62    /// Guest-side execution timeout in nanoseconds.
63    #[serde(default)]
64    pub timeout_ns: Option<u64>,
65    pub cmd: Vec<String>,
66}
67
68#[derive(Serialize, Deserialize)]
69pub struct PoolRunResponse {
70    pub stdout: Vec<u8>,
71    pub stderr: Vec<u8>,
72    pub exit_code: i32,
73    pub error: Option<String>,
74}
75
76#[derive(Serialize, Deserialize)]
77pub struct PoolLeaseRequest {
78    /// Image for the helper VM; `None` means use the daemon's default image.
79    #[serde(default)]
80    pub image: Option<String>,
81    /// Boot-time volume specs for this leased VM.
82    #[serde(default)]
83    pub volumes: Vec<String>,
84    /// Boot-time vCPU count.
85    #[serde(default)]
86    pub vcpus: Option<u32>,
87    /// Boot-time memory size.
88    #[serde(default)]
89    pub memory_mb: Option<u32>,
90}
91
92#[derive(Serialize, Deserialize)]
93pub struct PoolLeaseResponse {
94    pub lease_id: Option<String>,
95    pub error: Option<String>,
96}
97
98#[derive(Serialize, Deserialize)]
99pub struct PoolLeaseExecRequest {
100    pub lease_id: String,
101    pub cmd: Vec<String>,
102    #[serde(default)]
103    pub timeout_ns: Option<u64>,
104    #[serde(default)]
105    pub env: Vec<String>,
106    #[serde(default)]
107    pub working_dir: Option<String>,
108    #[serde(default)]
109    pub rootfs: Option<String>,
110    #[serde(default)]
111    pub stdin: Option<Vec<u8>>,
112    #[serde(default)]
113    pub user: Option<String>,
114}
115
116#[derive(Serialize, Deserialize)]
117pub struct PoolLeaseReleaseRequest {
118    pub lease_id: String,
119}
120
121#[derive(Serialize, Deserialize)]
122pub struct PoolLeaseReleaseResponse {
123    pub error: Option<String>,
124}
125
126/// Live stats for one image's warm pool.
127#[derive(Serialize, Deserialize)]
128pub struct PoolImageStat {
129    pub image: String,
130    pub pool: String,
131    /// Maximum concurrent sandboxes for this pool key.
132    #[serde(default)]
133    pub max: usize,
134    pub idle: usize,
135    /// Sandboxes currently checked out by one-shot runs or leases.
136    #[serde(default)]
137    pub active: usize,
138    /// Active sandboxes held by lease clients.
139    #[serde(default)]
140    pub leased: usize,
141    pub total_created: u64,
142    pub total_acquired: u64,
143    pub total_evicted: u64,
144}
145
146#[derive(Serialize, Deserialize)]
147pub struct PoolStatusResponse {
148    pub images: Vec<PoolImageStat>,
149}
150
151#[derive(Serialize, Deserialize)]
152pub struct PoolStopResponse {
153    pub error: Option<String>,
154}
155
156pub struct PoolClientRun {
157    pub socket: String,
158    pub image: Option<String>,
159    pub user: Option<String>,
160    pub workdir: Option<String>,
161    pub rootfs: Option<String>,
162    pub env: Vec<String>,
163    pub volumes: Vec<String>,
164    pub vcpus: u32,
165    pub memory_mb: u32,
166    pub exec: bool,
167    pub timeout_ns: Option<u64>,
168    pub cmd: Vec<String>,
169}
170
171pub struct PoolClientOutput {
172    pub stdout: Vec<u8>,
173    pub stderr: Vec<u8>,
174    pub exit_code: i32,
175}
176
177pub struct PoolLeaseClient {
178    socket: String,
179    lease_id: String,
180    released: bool,
181}
182
183impl PoolLeaseClient {
184    pub fn lease_id(&self) -> &str {
185        &self.lease_id
186    }
187
188    pub async fn acquire(req: PoolClientLease) -> Result<Self> {
189        let response = lease_client(&req).await?;
190        let lease_id = response.lease_id.ok_or_else(|| {
191            BoxError::PoolError("pool lease response did not include a lease id".to_string())
192        })?;
193        Ok(Self {
194            socket: req.socket,
195            lease_id,
196            released: false,
197        })
198    }
199
200    pub async fn exec(&self, req: PoolLeaseExec) -> Result<PoolClientOutput> {
201        lease_exec_client(
202            &self.socket,
203            PoolLeaseExecRequest {
204                lease_id: self.lease_id.clone(),
205                cmd: req.cmd,
206                timeout_ns: req.timeout_ns,
207                env: req.env,
208                working_dir: req.working_dir,
209                rootfs: req.rootfs,
210                stdin: req.stdin,
211                user: req.user,
212            },
213        )
214        .await
215    }
216
217    pub async fn release(mut self) -> Result<()> {
218        let result = release_client(&self.socket, &self.lease_id).await;
219        if result.is_ok() {
220            self.released = true;
221        }
222        result
223    }
224}
225
226impl Drop for PoolLeaseClient {
227    fn drop(&mut self) {
228        #[cfg(not(windows))]
229        if !self.released {
230            release_client_blocking_best_effort(&self.socket, &self.lease_id);
231        }
232    }
233}
234
235pub struct PoolClientLease {
236    pub socket: String,
237    pub image: Option<String>,
238    pub volumes: Vec<String>,
239    pub vcpus: u32,
240    pub memory_mb: u32,
241}
242
243pub struct PoolLeaseExec {
244    pub cmd: Vec<String>,
245    pub timeout_ns: Option<u64>,
246    pub env: Vec<String>,
247    pub working_dir: Option<String>,
248    pub rootfs: Option<String>,
249    pub stdin: Option<Vec<u8>>,
250    pub user: Option<String>,
251}
252
253#[cfg(not(windows))]
254pub async fn run_client(req: PoolClientRun) -> Result<PoolClientOutput> {
255    use tokio::net::UnixStream;
256
257    let mut stream = UnixStream::connect(&req.socket).await.map_err(|e| {
258        BoxError::PoolError(format!(
259            "Failed to connect to pool daemon at {} ({}). Is `a3s-box pool start` running?",
260            req.socket, e
261        ))
262    })?;
263
264    write_frame(
265        &mut stream,
266        &serde_json::to_vec(&PoolRequest::Run(PoolRunRequest {
267            image: req.image,
268            user: req.user,
269            workdir: req.workdir,
270            rootfs: req.rootfs,
271            env: req.env,
272            volumes: req.volumes,
273            vcpus: Some(req.vcpus),
274            memory_mb: Some(req.memory_mb),
275            exec: req.exec,
276            timeout_ns: req.timeout_ns,
277            cmd: req.cmd,
278        }))?,
279    )
280    .await?;
281    let resp: PoolRunResponse = serde_json::from_slice(&read_frame(&mut stream).await?)?;
282
283    if let Some(err) = resp.error {
284        return Err(BoxError::PoolError(err));
285    }
286
287    Ok(PoolClientOutput {
288        stdout: resp.stdout,
289        stderr: resp.stderr,
290        exit_code: resp.exit_code,
291    })
292}
293
294#[cfg(windows)]
295pub async fn run_client(_req: PoolClientRun) -> Result<PoolClientOutput> {
296    Err(BoxError::PoolError(
297        "`pool run` is not supported on Windows".to_string(),
298    ))
299}
300
301#[cfg(not(windows))]
302pub async fn status_client(socket: &str) -> Result<PoolStatusResponse> {
303    use tokio::net::UnixStream;
304
305    let mut stream = UnixStream::connect(socket).await.map_err(|e| {
306        BoxError::PoolError(format!("Failed to connect to pool daemon at {socket}: {e}"))
307    })?;
308    write_frame(&mut stream, &serde_json::to_vec(&PoolRequest::Status)?).await?;
309    Ok(serde_json::from_slice(&read_frame(&mut stream).await?)?)
310}
311
312#[cfg(not(windows))]
313pub async fn stop_client(socket: &str) -> Result<()> {
314    use tokio::net::UnixStream;
315
316    let mut stream = UnixStream::connect(socket).await.map_err(|e| {
317        BoxError::PoolError(format!("Failed to connect to pool daemon at {socket}: {e}"))
318    })?;
319    write_frame(&mut stream, &serde_json::to_vec(&PoolRequest::Stop)?).await?;
320    let resp: PoolStopResponse = serde_json::from_slice(&read_frame(&mut stream).await?)?;
321    if let Some(error) = resp.error {
322        return Err(BoxError::PoolError(error));
323    }
324    Ok(())
325}
326
327#[cfg(windows)]
328pub async fn stop_client(_socket: &str) -> Result<()> {
329    Err(BoxError::PoolError(
330        "`pool stop` is not supported on Windows".to_string(),
331    ))
332}
333
334#[cfg(not(windows))]
335async fn lease_client(req: &PoolClientLease) -> Result<PoolLeaseResponse> {
336    use tokio::net::UnixStream;
337
338    let mut stream = UnixStream::connect(&req.socket).await.map_err(|e| {
339        BoxError::PoolError(format!(
340            "Failed to connect to pool daemon at {} ({}). Is `a3s-box pool start` running?",
341            req.socket, e
342        ))
343    })?;
344    write_frame(
345        &mut stream,
346        &serde_json::to_vec(&PoolRequest::Lease(PoolLeaseRequest {
347            image: req.image.clone(),
348            volumes: req.volumes.clone(),
349            vcpus: Some(req.vcpus),
350            memory_mb: Some(req.memory_mb),
351        }))?,
352    )
353    .await?;
354    let resp: PoolLeaseResponse = serde_json::from_slice(&read_frame(&mut stream).await?)?;
355    if let Some(error) = resp.error.as_ref() {
356        return Err(BoxError::PoolError(error.clone()));
357    }
358    Ok(resp)
359}
360
361#[cfg(windows)]
362async fn lease_client(_req: &PoolClientLease) -> Result<PoolLeaseResponse> {
363    Err(BoxError::PoolError(
364        "warm-pool leases are not supported on Windows".to_string(),
365    ))
366}
367
368#[cfg(not(windows))]
369async fn lease_exec_client(socket: &str, req: PoolLeaseExecRequest) -> Result<PoolClientOutput> {
370    use tokio::net::UnixStream;
371
372    let mut stream = UnixStream::connect(socket).await.map_err(|e| {
373        BoxError::PoolError(format!("Failed to connect to pool daemon at {socket}: {e}"))
374    })?;
375    write_frame(&mut stream, &serde_json::to_vec(&PoolRequest::Exec(req))?).await?;
376    let resp: PoolRunResponse = serde_json::from_slice(&read_frame(&mut stream).await?)?;
377    if let Some(error) = resp.error {
378        return Err(BoxError::PoolError(error));
379    }
380    Ok(PoolClientOutput {
381        stdout: resp.stdout,
382        stderr: resp.stderr,
383        exit_code: resp.exit_code,
384    })
385}
386
387#[cfg(windows)]
388async fn lease_exec_client(_socket: &str, _req: PoolLeaseExecRequest) -> Result<PoolClientOutput> {
389    Err(BoxError::PoolError(
390        "warm-pool leases are not supported on Windows".to_string(),
391    ))
392}
393
394#[cfg(not(windows))]
395async fn release_client(socket: &str, lease_id: &str) -> Result<()> {
396    use tokio::net::UnixStream;
397
398    let mut stream = UnixStream::connect(socket).await.map_err(|e| {
399        BoxError::PoolError(format!("Failed to connect to pool daemon at {socket}: {e}"))
400    })?;
401    write_frame(
402        &mut stream,
403        &serde_json::to_vec(&PoolRequest::Release(PoolLeaseReleaseRequest {
404            lease_id: lease_id.to_string(),
405        }))?,
406    )
407    .await?;
408    let resp: PoolLeaseReleaseResponse = serde_json::from_slice(&read_frame(&mut stream).await?)?;
409    if let Some(error) = resp.error {
410        return Err(BoxError::PoolError(error));
411    }
412    Ok(())
413}
414
415#[cfg(windows)]
416async fn release_client(_socket: &str, _lease_id: &str) -> Result<()> {
417    Err(BoxError::PoolError(
418        "warm-pool leases are not supported on Windows".to_string(),
419    ))
420}
421
422#[cfg(not(windows))]
423fn release_client_blocking_best_effort(socket: &str, lease_id: &str) {
424    use std::io::Write;
425    use std::os::unix::net::UnixStream;
426    use std::time::Duration;
427
428    let Ok(mut stream) = UnixStream::connect(socket) else {
429        return;
430    };
431    let timeout = Some(Duration::from_millis(500));
432    let _ = stream.set_read_timeout(timeout);
433    let _ = stream.set_write_timeout(timeout);
434
435    let Ok(payload) = serde_json::to_vec(&PoolRequest::Release(PoolLeaseReleaseRequest {
436        lease_id: lease_id.to_string(),
437    })) else {
438        return;
439    };
440    let _ = stream
441        .write_all(&(payload.len() as u32).to_le_bytes())
442        .and_then(|_| stream.write_all(&payload))
443        .and_then(|_| stream.flush());
444}
445
446/// Length-prefixed (u32 LE) framing for the pool Unix-socket protocol.
447#[cfg(not(windows))]
448pub async fn write_frame<W>(w: &mut W, data: &[u8]) -> std::io::Result<()>
449where
450    W: tokio::io::AsyncWrite + Unpin,
451{
452    use tokio::io::AsyncWriteExt;
453
454    if data.len() > MAX_FRAME_SIZE {
455        return Err(std::io::Error::new(
456            std::io::ErrorKind::InvalidInput,
457            format!(
458                "pool frame is {} bytes; maximum is {} bytes",
459                data.len(),
460                MAX_FRAME_SIZE
461            ),
462        ));
463    }
464    w.write_all(&(data.len() as u32).to_le_bytes()).await?;
465    w.write_all(data).await?;
466    w.flush().await
467}
468
469#[cfg(not(windows))]
470pub async fn read_frame<R>(r: &mut R) -> std::io::Result<Vec<u8>>
471where
472    R: tokio::io::AsyncRead + Unpin,
473{
474    use tokio::io::AsyncReadExt;
475
476    let mut len = [0u8; 4];
477    r.read_exact(&mut len).await?;
478    let frame_len = u32::from_le_bytes(len) as usize;
479    if frame_len > MAX_FRAME_SIZE {
480        return Err(std::io::Error::new(
481            std::io::ErrorKind::InvalidData,
482            format!(
483                "pool frame is {} bytes; maximum is {} bytes",
484                frame_len, MAX_FRAME_SIZE
485            ),
486        ));
487    }
488    let mut buf = vec![0u8; frame_len];
489    r.read_exact(&mut buf).await?;
490    Ok(buf)
491}
492
493/// Read one pool frame with a bounded wait for the length prefix and payload.
494#[cfg(not(windows))]
495pub async fn read_frame_with_timeout<R>(
496    r: &mut R,
497    timeout: std::time::Duration,
498) -> std::io::Result<Vec<u8>>
499where
500    R: tokio::io::AsyncRead + Unpin,
501{
502    tokio::time::timeout(timeout, read_frame(r))
503        .await
504        .map_err(|_| {
505            std::io::Error::new(
506                std::io::ErrorKind::TimedOut,
507                "timed out waiting for pool frame",
508            )
509        })?
510}
511
512#[cfg(test)]
513mod tests {
514    #[cfg(not(windows))]
515    use super::{read_frame, read_frame_with_timeout, write_frame, MAX_FRAME_SIZE};
516
517    #[cfg(not(windows))]
518    #[tokio::test]
519    async fn oversized_frame_is_rejected_before_allocation() {
520        use tokio::io::AsyncWriteExt;
521
522        let (mut writer, mut reader) = tokio::io::duplex(64);
523        writer
524            .write_all(&((MAX_FRAME_SIZE as u32) + 1).to_le_bytes())
525            .await
526            .unwrap();
527
528        let error = read_frame(&mut reader).await.unwrap_err();
529        assert_eq!(error.kind(), std::io::ErrorKind::InvalidData);
530        assert!(error.to_string().contains("maximum"));
531    }
532
533    #[cfg(not(windows))]
534    #[tokio::test]
535    async fn oversized_frame_is_rejected_before_write() {
536        let (mut writer, _) = tokio::io::duplex(64);
537        let payload = vec![0u8; MAX_FRAME_SIZE + 1];
538
539        let error = write_frame(&mut writer, &payload).await.unwrap_err();
540        assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput);
541    }
542
543    #[cfg(not(windows))]
544    #[tokio::test]
545    async fn frame_read_timeout_returns_a_timed_out_error() {
546        let (_writer, mut reader) = tokio::io::duplex(64);
547
548        let error = read_frame_with_timeout(&mut reader, std::time::Duration::from_millis(10))
549            .await
550            .unwrap_err();
551        assert_eq!(error.kind(), std::io::ErrorKind::TimedOut);
552    }
553
554    #[cfg(not(windows))]
555    #[test]
556    fn lease_drop_releases_synchronously() {
557        use super::*;
558        use std::io::Read;
559        use std::os::unix::net::UnixListener;
560
561        let tmp = tempfile::TempDir::new().unwrap();
562        let socket = tmp.path().join("pool.sock");
563        let listener = UnixListener::bind(&socket).unwrap();
564        let socket_arg = socket.to_string_lossy().to_string();
565
566        let server = std::thread::spawn(move || {
567            let (mut stream, _) = listener.accept().unwrap();
568            let mut len = [0_u8; 4];
569            stream.read_exact(&mut len).unwrap();
570            let mut request = vec![0_u8; u32::from_le_bytes(len) as usize];
571            stream.read_exact(&mut request).unwrap();
572            let request: PoolRequest = serde_json::from_slice(&request).unwrap();
573            match request {
574                PoolRequest::Release(req) => assert_eq!(req.lease_id, "lease-drop"),
575                _ => panic!("drop should send release request"),
576            }
577        });
578
579        let lease = PoolLeaseClient {
580            socket: socket_arg,
581            lease_id: "lease-drop".to_string(),
582            released: false,
583        };
584        drop(lease);
585
586        server.join().unwrap();
587    }
588}
a3s_box_runtime/pool/client.rs

a3s_box_runtime/pool/
client.rs
