use std::time::Duration;
use arcbox_vm_driver::Vsock;
use serde::Serialize;
use tokio::io::AsyncReadExt as _;
use tokio::net::UnixStream;
use super::{MAX_FRAME_SIZE, connect_to_port, read_frame, write_frame};
use crate::agent::{DirWatch, FsEvents};
use crate::error::{ComputerError, Result};
const FILE_IO_TIMEOUT: Duration = Duration::from_mins(1);
pub use arcbox_vm_proto::file as proto;
pub use proto::FILE_PORT;
use proto::{
FILE_ACK, FILE_DATA, FILE_DONE, FILE_ERR, FILE_EVENT, FILE_LIST, FILE_LIST_REQ, FILE_MKDIR_REQ,
FILE_MOVE_REQ, FILE_READ_REQ, FILE_REMOVE_REQ, FILE_STAT, FILE_STAT_REQ, FILE_WATCH_REQ,
FILE_WRITE_REQ, FileStatDto, FsEventDto, ListDirReq, MAX_FILE_SIZE, MakeDirReq, MoveReq,
RemoveReq, StatReq, WatchReq,
};
#[derive(Serialize)]
struct WriteReq<'a> {
path: &'a str,
mode: u32,
}
#[derive(Serialize)]
struct ReadReq<'a> {
path: &'a str,
}
pub(super) async fn write_file(
vsock: &dyn Vsock,
path: &str,
mode: u32,
data: &[u8],
) -> Result<()> {
if data.len() > MAX_FILE_SIZE {
return Err(ComputerError::Vsock(format!(
"file too large ({} bytes, max {MAX_FILE_SIZE})",
data.len()
)));
}
tokio::time::timeout(FILE_IO_TIMEOUT, write_file_inner(vsock, path, mode, data))
.await
.map_err(|_| ComputerError::Vsock("file write: timed out".into()))?
}
async fn write_file_inner(vsock: &dyn Vsock, path: &str, mode: u32, data: &[u8]) -> Result<()> {
let mut stream = connect_to_port(vsock, FILE_PORT).await?;
let req = serde_json::to_vec(&WriteReq { path, mode })
.map_err(|e| ComputerError::Vsock(format!("serialize WriteReq: {e}")))?;
write_frame(&mut stream, FILE_WRITE_REQ, &req)
.await
.map_err(|e| ComputerError::Vsock(format!("write FILE_WRITE_REQ: {e}")))?;
for chunk in data.chunks(MAX_FRAME_SIZE) {
write_frame(&mut stream, FILE_DATA, chunk)
.await
.map_err(|e| ComputerError::Vsock(format!("write FILE_DATA: {e}")))?;
}
write_frame(&mut stream, FILE_DONE, &[])
.await
.map_err(|e| ComputerError::Vsock(format!("write FILE_DONE: {e}")))?;
let (resp_type, payload) = read_frame(&mut stream)
.await
.map_err(|e| ComputerError::Vsock(format!("read write response: {e}")))?;
match resp_type {
FILE_ACK => Ok(()),
FILE_ERR => Err(decode_file_err(&payload)),
other => Err(ComputerError::Vsock(format!(
"file write: unexpected response type 0x{other:02x}"
))),
}
}
pub(super) async fn read_file(vsock: &dyn Vsock, path: &str) -> Result<Vec<u8>> {
tokio::time::timeout(FILE_IO_TIMEOUT, read_file_inner(vsock, path))
.await
.map_err(|_| ComputerError::Vsock("file read: timed out".into()))?
}
async fn read_file_inner(vsock: &dyn Vsock, path: &str) -> Result<Vec<u8>> {
let mut stream = connect_to_port(vsock, FILE_PORT).await?;
let req = serde_json::to_vec(&ReadReq { path })
.map_err(|e| ComputerError::Vsock(format!("serialize ReadReq: {e}")))?;
write_frame(&mut stream, FILE_READ_REQ, &req)
.await
.map_err(|e| ComputerError::Vsock(format!("write FILE_READ_REQ: {e}")))?;
let mut buf = Vec::new();
loop {
let (frame_type, payload) = read_frame(&mut stream)
.await
.map_err(|e| ComputerError::Vsock(format!("read file data: {e}")))?;
match frame_type {
FILE_DATA => {
buf.extend_from_slice(&payload);
if buf.len() > MAX_FILE_SIZE {
return Err(ComputerError::Vsock(format!(
"file too large (>{MAX_FILE_SIZE} bytes)"
)));
}
}
FILE_DONE => return Ok(buf),
FILE_ERR => return Err(decode_file_err(&payload)),
other => {
return Err(ComputerError::Vsock(format!(
"file read: unexpected frame type 0x{other:02x}"
)));
}
}
}
}
fn decode_file_err(payload: &[u8]) -> ComputerError {
let text = String::from_utf8_lossy(payload).into_owned();
if let Some(path) = text.strip_prefix(proto::ERR_NOT_FOUND) {
ComputerError::PathNotFound(path.to_owned())
} else if let Some(path) = text.strip_prefix(proto::ERR_NOT_A_DIRECTORY) {
ComputerError::NotADirectory(path.to_owned())
} else if let Some(path) = text.strip_prefix(proto::ERR_NOT_EMPTY) {
ComputerError::DirectoryNotEmpty(path.to_owned())
} else {
ComputerError::Vsock(text)
}
}
async fn unary_file_op(
vsock: &dyn Vsock,
req_type: u8,
req: &(impl Serialize + Sync),
ok_type: u8,
) -> Result<Vec<u8>> {
let payload = serde_json::to_vec(req)
.map_err(|e| ComputerError::Vsock(format!("serialize request: {e}")))?;
let op = async {
let mut stream = connect_to_port(vsock, FILE_PORT).await?;
write_frame(&mut stream, req_type, &payload)
.await
.map_err(|e| ComputerError::Vsock(format!("write request 0x{req_type:02x}: {e}")))?;
let (resp_type, resp) = read_frame(&mut stream)
.await
.map_err(|e| ComputerError::Vsock(format!("read response: {e}")))?;
match resp_type {
t if t == ok_type => Ok(resp),
FILE_ERR => Err(decode_file_err(&resp)),
other => Err(ComputerError::Vsock(format!(
"unexpected response type 0x{other:02x}"
))),
}
};
tokio::time::timeout(FILE_IO_TIMEOUT, op)
.await
.map_err(|_| ComputerError::Vsock("file operation timed out".into()))?
}
fn parse_json<T: serde::de::DeserializeOwned>(payload: &[u8]) -> Result<T> {
serde_json::from_slice(payload)
.map_err(|e| ComputerError::Vsock(format!("malformed agent response: {e}")))
}
pub(super) async fn stat_file(vsock: &dyn Vsock, path: &str) -> Result<FileStatDto> {
let req = StatReq {
path: path.to_owned(),
};
let payload = unary_file_op(vsock, FILE_STAT_REQ, &req, FILE_STAT).await?;
parse_json(&payload)
}
pub(super) async fn list_dir(vsock: &dyn Vsock, path: &str) -> Result<Vec<FileStatDto>> {
let req = ListDirReq {
path: path.to_owned(),
};
let payload = unary_file_op(vsock, FILE_LIST_REQ, &req, FILE_LIST).await?;
parse_json(&payload)
}
pub(super) async fn make_dir(vsock: &dyn Vsock, path: &str, mode: u32) -> Result<()> {
let req = MakeDirReq {
path: path.to_owned(),
mode,
};
unary_file_op(vsock, FILE_MKDIR_REQ, &req, FILE_ACK).await?;
Ok(())
}
pub(super) async fn remove_entry(vsock: &dyn Vsock, path: &str, recursive: bool) -> Result<()> {
let req = RemoveReq {
path: path.to_owned(),
recursive,
};
unary_file_op(vsock, FILE_REMOVE_REQ, &req, FILE_ACK).await?;
Ok(())
}
pub(super) async fn move_entry(vsock: &dyn Vsock, from: &str, to: &str) -> Result<()> {
let req = MoveReq {
from: from.to_owned(),
to: to.to_owned(),
};
unary_file_op(vsock, FILE_MOVE_REQ, &req, FILE_ACK).await?;
Ok(())
}
struct VsockWatch {
stream: UnixStream,
}
#[async_trait::async_trait]
impl FsEvents for VsockWatch {
async fn next_event(&mut self) -> Result<Option<FsEventDto>> {
let mut ty = [0u8; 1];
let n = self
.stream
.read(&mut ty)
.await
.map_err(|e| ComputerError::Vsock(format!("watch read: {e}")))?;
if n == 0 {
return Ok(None);
}
let len = self
.stream
.read_u32_le()
.await
.map_err(|e| ComputerError::Vsock(format!("watch read: {e}")))?
as usize;
if len > MAX_FRAME_SIZE {
return Err(ComputerError::Vsock(format!(
"watch frame too large: {len}"
)));
}
let mut payload = vec![0u8; len];
if len > 0 {
self.stream
.read_exact(&mut payload)
.await
.map_err(|e| ComputerError::Vsock(format!("watch read: {e}")))?;
}
match ty[0] {
FILE_EVENT => Ok(Some(parse_json(&payload)?)),
FILE_ERR => Err(decode_file_err(&payload)),
other => Err(ComputerError::Vsock(format!(
"unexpected watch frame type 0x{other:02x}"
))),
}
}
}
pub(super) async fn watch_dir(vsock: &dyn Vsock, path: &str, recursive: bool) -> Result<DirWatch> {
let req = WatchReq {
path: path.to_owned(),
recursive,
};
let payload = serde_json::to_vec(&req)
.map_err(|e| ComputerError::Vsock(format!("serialize WatchReq: {e}")))?;
let setup = async {
let mut stream = connect_to_port(vsock, FILE_PORT).await?;
write_frame(&mut stream, FILE_WATCH_REQ, &payload)
.await
.map_err(|e| ComputerError::Vsock(format!("write FILE_WATCH_REQ: {e}")))?;
let (resp_type, resp) = read_frame(&mut stream)
.await
.map_err(|e| ComputerError::Vsock(format!("read watch ack: {e}")))?;
match resp_type {
FILE_ACK => Ok(DirWatch::new(Box::new(VsockWatch { stream }))),
FILE_ERR => Err(decode_file_err(&resp)),
other => Err(ComputerError::Vsock(format!(
"unexpected watch ack type 0x{other:02x}"
))),
}
};
tokio::time::timeout(FILE_IO_TIMEOUT, setup)
.await
.map_err(|_| ComputerError::Vsock("watch setup timed out".into()))?
}
#[cfg(test)]
mod tests {
use super::*;
use std::path::PathBuf;
use arcbox_vm_driver::{IoMode, Vsock, VsockConn};
use async_trait::async_trait;
use super::super::{read_frame as async_read_frame, write_frame as async_write_frame};
struct UnixVsock(PathBuf);
#[async_trait]
impl Vsock for UnixVsock {
async fn dial(&self, _port: u32) -> arcbox_vm_driver::Result<VsockConn> {
let stream = UnixStream::connect(&self.0).await?;
Ok(VsockConn {
fd: stream.into_std()?.into(),
mode: IoMode::Async,
})
}
}
#[test]
fn test_write_req_serializes() {
let req = WriteReq {
path: "/tmp/test.txt",
mode: 0o644,
};
let json = serde_json::to_string(&req).unwrap();
assert!(json.contains("/tmp/test.txt"));
assert!(json.contains("420")); }
#[test]
fn test_read_req_serializes() {
let req = ReadReq { path: "/etc/hosts" };
let json = serde_json::to_string(&req).unwrap();
assert!(json.contains("/etc/hosts"));
}
#[tokio::test]
async fn test_write_file_protocol_success() {
let (mut agent, host) = tokio::io::duplex(8192);
let agent_handle = tokio::spawn(async move {
let (ty, payload) = async_read_frame(&mut agent).await.unwrap();
assert_eq!(ty, FILE_WRITE_REQ);
let parsed: serde_json::Value = serde_json::from_slice(&payload).unwrap();
assert_eq!(parsed["path"], "/tmp/hello.txt");
let mut data = Vec::new();
loop {
let (ty, chunk) = async_read_frame(&mut agent).await.unwrap();
match ty {
FILE_DATA => data.extend_from_slice(&chunk),
FILE_DONE => break,
_ => panic!("unexpected frame type 0x{ty:02x}"),
}
}
assert_eq!(data, b"hello world");
async_write_frame(&mut agent, FILE_ACK, &[]).await.unwrap();
});
let mut stream = host;
let req = serde_json::to_vec(&WriteReq {
path: "/tmp/hello.txt",
mode: 0o644,
})
.unwrap();
async_write_frame(&mut stream, FILE_WRITE_REQ, &req)
.await
.unwrap();
for chunk in b"hello world".chunks(MAX_FRAME_SIZE) {
async_write_frame(&mut stream, FILE_DATA, chunk)
.await
.unwrap();
}
async_write_frame(&mut stream, FILE_DONE, &[])
.await
.unwrap();
let (resp_type, _) = async_read_frame(&mut stream).await.unwrap();
assert_eq!(resp_type, FILE_ACK);
agent_handle.await.unwrap();
}
#[tokio::test]
async fn test_write_file_protocol_error() {
let (mut agent, host) = tokio::io::duplex(8192);
let agent_handle = tokio::spawn(async move {
let _ = async_read_frame(&mut agent).await.unwrap();
loop {
let (ty, _) = async_read_frame(&mut agent).await.unwrap();
if ty == FILE_DONE {
break;
}
}
async_write_frame(&mut agent, FILE_ERR, b"permission denied")
.await
.unwrap();
});
let mut stream = host;
let req = serde_json::to_vec(&WriteReq {
path: "/root/secret",
mode: 0o600,
})
.unwrap();
async_write_frame(&mut stream, FILE_WRITE_REQ, &req)
.await
.unwrap();
async_write_frame(&mut stream, FILE_DONE, &[])
.await
.unwrap();
let (resp_type, payload) = async_read_frame(&mut stream).await.unwrap();
assert_eq!(resp_type, FILE_ERR);
assert_eq!(std::str::from_utf8(&payload).unwrap(), "permission denied");
agent_handle.await.unwrap();
}
#[tokio::test]
async fn test_read_file_protocol_success() {
let (mut agent, host) = tokio::io::duplex(8192);
let agent_handle = tokio::spawn(async move {
let (ty, _payload) = async_read_frame(&mut agent).await.unwrap();
assert_eq!(ty, FILE_READ_REQ);
async_write_frame(&mut agent, FILE_DATA, b"part1")
.await
.unwrap();
async_write_frame(&mut agent, FILE_DATA, b"part2")
.await
.unwrap();
async_write_frame(&mut agent, FILE_DONE, &[]).await.unwrap();
});
let mut stream = host;
let req = serde_json::to_vec(&ReadReq {
path: "/tmp/test.txt",
})
.unwrap();
async_write_frame(&mut stream, FILE_READ_REQ, &req)
.await
.unwrap();
let mut buf = Vec::new();
loop {
let (ty, payload) = async_read_frame(&mut stream).await.unwrap();
match ty {
FILE_DATA => buf.extend_from_slice(&payload),
FILE_DONE => break,
_ => panic!("unexpected frame type 0x{ty:02x}"),
}
}
assert_eq!(buf, b"part1part2");
agent_handle.await.unwrap();
}
#[tokio::test]
async fn test_read_file_protocol_error() {
let (mut agent, host) = tokio::io::duplex(8192);
let agent_handle = tokio::spawn(async move {
let _ = async_read_frame(&mut agent).await.unwrap();
async_write_frame(&mut agent, FILE_ERR, b"no such file")
.await
.unwrap();
});
let mut stream = host;
let req = serde_json::to_vec(&ReadReq {
path: "/nonexistent",
})
.unwrap();
async_write_frame(&mut stream, FILE_READ_REQ, &req)
.await
.unwrap();
let (ty, payload) = async_read_frame(&mut stream).await.unwrap();
assert_eq!(ty, FILE_ERR);
assert_eq!(std::str::from_utf8(&payload).unwrap(), "no such file");
agent_handle.await.unwrap();
}
#[test]
fn file_err_prefixes_decode_to_typed_errors() {
assert!(matches!(
decode_file_err(b"ENOENT: /a/b"),
ComputerError::PathNotFound(p) if p == "/a/b"
));
assert!(matches!(
decode_file_err(b"ENOTDIR: /a/file"),
ComputerError::NotADirectory(p) if p == "/a/file"
));
assert!(matches!(
decode_file_err(b"ENOTEMPTY: /a/dir"),
ComputerError::DirectoryNotEmpty(p) if p == "/a/dir"
));
assert!(matches!(
decode_file_err(b"read file: boom"),
ComputerError::Vsock(m) if m == "read file: boom"
));
}
async fn mock_vsock_server<F, Fut>(script: F) -> (tempfile::TempDir, UnixVsock)
where
F: FnOnce(UnixStream) -> Fut + Send + 'static,
Fut: Future<Output = ()> + Send,
{
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("v.sock");
let listener = tokio::net::UnixListener::bind(&path).unwrap();
tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
script(stream).await;
});
(dir, UnixVsock(path))
}
#[tokio::test]
async fn stat_file_round_trips_the_dto() {
let dto = FileStatDto {
name: "b".into(),
kind: proto::KIND_FILE.into(),
size: 42,
mode: 0o644,
mtime_secs: 1_700_000_000,
mtime_nanos: 5,
uid: 0,
gid: 0,
symlink_target: String::new(),
};
let expected = dto.clone();
let (_dir, vsock) = mock_vsock_server(move |mut stream| async move {
let (ty, payload) = async_read_frame(&mut stream).await.unwrap();
assert_eq!(ty, FILE_STAT_REQ);
let req: StatReq = serde_json::from_slice(&payload).unwrap();
assert_eq!(req.path, "/a/b");
let body = serde_json::to_vec(&dto).unwrap();
async_write_frame(&mut stream, FILE_STAT, &body)
.await
.unwrap();
})
.await;
let got = stat_file(&vsock, "/a/b").await.unwrap();
assert_eq!(got, expected);
}
#[tokio::test]
async fn remove_entry_surfaces_directory_not_empty() {
let (_dir, vsock) = mock_vsock_server(|mut stream| async move {
let (ty, _) = async_read_frame(&mut stream).await.unwrap();
assert_eq!(ty, FILE_REMOVE_REQ);
async_write_frame(&mut stream, FILE_ERR, b"ENOTEMPTY: /full")
.await
.unwrap();
})
.await;
let err = remove_entry(&vsock, "/full", false).await.unwrap_err();
assert!(matches!(err, ComputerError::DirectoryNotEmpty(p) if p == "/full"));
}
#[tokio::test]
async fn read_file_error_is_classified() {
let (_dir, vsock) = mock_vsock_server(|mut stream| async move {
let _ = async_read_frame(&mut stream).await.unwrap();
async_write_frame(&mut stream, FILE_ERR, b"ENOENT: /missing")
.await
.unwrap();
})
.await;
let err = read_file(&vsock, "/missing").await.unwrap_err();
assert!(matches!(err, ComputerError::PathNotFound(p) if p == "/missing"));
}
#[tokio::test]
async fn write_file_error_is_classified() {
let (_dir, vsock) = mock_vsock_server(|mut stream| async move {
let _ = async_read_frame(&mut stream).await.unwrap();
loop {
let (ty, _) = async_read_frame(&mut stream).await.unwrap();
if ty == FILE_DONE {
break;
}
}
async_write_frame(&mut stream, FILE_ERR, b"ENOTDIR: /plain.txt/sub")
.await
.unwrap();
})
.await;
let err = write_file(&vsock, "/plain.txt/sub", 0, b"x")
.await
.unwrap_err();
assert!(matches!(err, ComputerError::NotADirectory(p) if p == "/plain.txt/sub"));
}
#[tokio::test]
async fn watch_dir_streams_events_until_clean_eof() {
let (_dir, vsock) = mock_vsock_server(|mut stream| async move {
let (ty, payload) = async_read_frame(&mut stream).await.unwrap();
assert_eq!(ty, FILE_WATCH_REQ);
let req: WatchReq = serde_json::from_slice(&payload).unwrap();
assert!(req.recursive);
async_write_frame(&mut stream, FILE_ACK, &[]).await.unwrap();
let event = FsEventDto {
kind: proto::EVENT_CREATED.into(),
path: "/w/new".into(),
renamed_to: String::new(),
};
let body = serde_json::to_vec(&event).unwrap();
async_write_frame(&mut stream, FILE_EVENT, &body)
.await
.unwrap();
})
.await;
let mut watch = watch_dir(&vsock, "/w", true).await.unwrap();
let event = watch.next_event().await.unwrap().unwrap();
assert_eq!(event.kind, proto::EVENT_CREATED);
assert_eq!(event.path, "/w/new");
assert!(watch.next_event().await.unwrap().is_none());
}
#[tokio::test]
async fn watch_dir_setup_error_is_typed() {
let (_dir, vsock) = mock_vsock_server(|mut stream| async move {
let _ = async_read_frame(&mut stream).await.unwrap();
async_write_frame(&mut stream, FILE_ERR, b"ENOENT: /missing")
.await
.unwrap();
})
.await;
let err = watch_dir(&vsock, "/missing", false).await.unwrap_err();
assert!(matches!(err, ComputerError::PathNotFound(p) if p == "/missing"));
}
}