use crate::framing::{read_frame, Frame, FLAG_FRAGMENTED, FLAG_MAC_PRESENT, FLAG_RAW_BINARY};
use prost::Message;
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use tokio::net::unix::{OwnedReadHalf, OwnedWriteHalf};
use tokio::net::UnixStream;
use veyron_wire::framing::{
parse_frag_header, serialize_header, write_frame_raw, FRAG_HEADER_SIZE, MAX_PAYLOAD_SIZE,
};
use veyron_wire::mac::{compute_tag, derive_session_key, verify_tag};
use veyron_wire::proto::veyron::{
envelope, ActionRequest, ActionRequestChunk, ActionResponse, ActionResponseChunk,
AudioStreamChunk, Envelope, EventAck, EventPublish, EventPublishAck, KernelCommand,
KernelCommandAck, Ping, PluginManifest, PluginRegister, PluginRegisterAck, SessionClose,
Subscribe, Unsubscribe,
};
use veyron_wire::WireError as VeyronError;
const MAX_REASSEMBLY_STREAMS: usize = 64;
const REASSEMBLY_TIMEOUT: Duration = Duration::from_secs(30);
const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
static ACTION_SEQ: AtomicU64 = AtomicU64::new(0);
struct ReassemblyBuf {
fragments: HashMap<u16, Vec<u8>>,
total: u16,
target: [u8; 32],
flags: u16,
first_seen: Instant,
buffered_bytes: usize,
}
impl ReassemblyBuf {
fn is_complete(&self) -> bool {
self.fragments.len() == self.total as usize
}
fn reassemble(mut self) -> Vec<u8> {
let mut out = Vec::with_capacity(self.buffered_bytes);
for seq in 0..self.total {
if let Some(chunk) = self.fragments.remove(&seq) {
out.extend_from_slice(&chunk);
}
}
out
}
}
pub struct VeyronClient {
read: OwnedReadHalf,
write: OwnedWriteHalf,
secret: Option<Vec<u8>>,
session_key: Option<[u8; 32]>,
reassembly: HashMap<u32, ReassemblyBuf>,
next_stream_id: u32,
}
impl VeyronClient {
pub async fn connect(socket_path: &str) -> Result<Self, VeyronError> {
Self::connect_inner(socket_path, None).await
}
pub async fn connect_with_secret(
socket_path: &str,
secret: &[u8],
) -> Result<Self, VeyronError> {
Self::connect_inner(socket_path, Some(secret.to_vec())).await
}
pub async fn connect_from_env() -> Result<Self, VeyronError> {
let socket_path = std::env::var("VEYRON_SOCKET_PATH")
.unwrap_or_else(|_| veyron_wire::socket::default_socket_path());
match std::env::var("VEYRON_JWT_SECRET") {
Ok(secret) if !secret.is_empty() => {
Self::connect_with_secret(&socket_path, secret.as_bytes()).await
}
_ => Self::connect(&socket_path).await,
}
}
pub fn from_stream(stream: UnixStream, secret: Option<Vec<u8>>) -> Self {
let (read, write) = stream.into_split();
Self {
read,
write,
secret,
session_key: None,
reassembly: HashMap::new(),
next_stream_id: 1,
}
}
async fn connect_inner(
socket_path: &str,
secret: Option<Vec<u8>>,
) -> Result<Self, VeyronError> {
let stream = UnixStream::connect(socket_path)
.await
.map_err(VeyronError::Io)?;
Ok(Self::from_stream(stream, secret))
}
pub fn is_secured(&self) -> bool {
self.session_key.is_some()
}
pub async fn register(
&mut self,
plugin_id: &str,
manifest: PluginManifest,
) -> Result<PluginRegisterAck, VeyronError> {
self.register_with_token(plugin_id, manifest, "").await
}
pub async fn register_with_token(
&mut self,
plugin_id: &str,
manifest: PluginManifest,
jwt_token: &str,
) -> Result<PluginRegisterAck, VeyronError> {
self.register_full(plugin_id, "1.0.0", manifest, jwt_token)
.await
}
pub async fn register_full(
&mut self,
plugin_id: &str,
version: &str,
manifest: PluginManifest,
jwt_token: &str,
) -> Result<PluginRegisterAck, VeyronError> {
let env = Envelope {
payload: Some(envelope::Payload::PluginRegister(PluginRegister {
plugin_id: plugin_id.to_string(),
version: version.to_string(),
manifest: Some(manifest),
jwt_token: jwt_token.to_string(),
..Default::default()
})),
..Default::default()
};
self.send("kernel", env).await?;
let response = self.recv().await?;
match response.payload {
Some(envelope::Payload::PluginRegisterAck(ack)) => {
if let Some(secret) = &self.secret {
if !ack.session_nonce.is_empty() {
self.session_key =
Some(derive_session_key(secret, &ack.session_nonce, plugin_id));
}
}
Ok(ack)
}
Some(envelope::Payload::Error(err)) => Err(VeyronError::Internal(format!(
"registration rejected: {} ({})",
err.message, err.details
))),
_ => Err(VeyronError::Internal("expected PluginRegisterAck".into())),
}
}
pub async fn send(&mut self, target: &str, envelope: Envelope) -> Result<(), VeyronError> {
let mut payload = Vec::new();
envelope
.encode(&mut payload)
.map_err(|_| VeyronError::Internal("encode failed".into()))?;
self.send_raw(target, payload).await
}
pub async fn send_raw(&mut self, target: &str, payload: Vec<u8>) -> Result<(), VeyronError> {
self.send_raw_with_flags(target, 0, payload).await
}
pub async fn send_raw_with_flags(
&mut self,
target: &str,
extra_flags: u16,
payload: Vec<u8>,
) -> Result<(), VeyronError> {
let base_flags = if self.session_key.is_some() {
FLAG_MAC_PRESENT
} else {
0
};
let mut frame = build_frame(target, base_flags | extra_flags, payload);
if let Some(key) = &self.session_key {
let header = serialize_header(&frame);
frame.mac = Some(compute_tag(key, &header, &frame.payload));
}
write_frame_raw(&mut self.write, &frame).await
}
pub async fn send_fragmented(
&mut self,
target: &str,
payload: &[u8],
chunk_size: usize,
) -> Result<(), VeyronError> {
if payload.len() > MAX_PAYLOAD_SIZE {
return Err(VeyronError::PayloadTooLarge(payload.len()));
}
if chunk_size == 0 || chunk_size + FRAG_HEADER_SIZE > MAX_PAYLOAD_SIZE {
return Err(VeyronError::Internal(format!(
"invalid fragment chunk_size: {chunk_size}"
)));
}
let total = payload.len().div_ceil(chunk_size).max(1);
if total > u16::MAX as usize {
return Err(VeyronError::Internal(format!(
"payload needs {total} fragments; max is {}",
u16::MAX
)));
}
let stream_id = self.next_stream_id;
self.next_stream_id = self.next_stream_id.wrapping_add(1).max(1);
let fragment_id = (stream_id & 0xFFFF) as u16;
for (seq, chunk) in payload.chunks(chunk_size).enumerate() {
let mut frag_payload = Vec::with_capacity(FRAG_HEADER_SIZE + chunk.len());
frag_payload.extend_from_slice(&fragment_id.to_be_bytes());
frag_payload.extend_from_slice(&(seq as u16).to_be_bytes());
frag_payload.extend_from_slice(&(total as u16).to_be_bytes());
frag_payload.extend_from_slice(&stream_id.to_be_bytes());
frag_payload.extend_from_slice(chunk);
self.send_raw_with_flags(target, FLAG_FRAGMENTED, frag_payload)
.await?;
}
Ok(())
}
pub async fn recv_frame(&mut self) -> Result<Frame, VeyronError> {
loop {
let frame = read_frame(&mut self.read).await?;
self.verify_frame_mac(&frame)?;
if frame.flags & FLAG_FRAGMENTED != 0 {
if let Some(complete) = self.absorb_fragment(frame)? {
return Ok(complete);
}
continue;
}
return Ok(frame);
}
}
pub async fn recv(&mut self) -> Result<Envelope, VeyronError> {
let frame = self.recv_frame().await?;
if frame.flags & FLAG_RAW_BINARY != 0 {
return Err(VeyronError::Internal(
"received raw-binary frame; use recv_frame() for audio".into(),
));
}
Envelope::decode(frame.payload.as_ref()).map_err(VeyronError::Proto)
}
pub async fn recv_timeout(&mut self, timeout: Duration) -> Result<Envelope, VeyronError> {
match tokio::time::timeout(timeout, self.recv()).await {
Ok(result) => result,
Err(_) => Err(VeyronError::Timeout),
}
}
fn verify_frame_mac(&self, frame: &Frame) -> Result<(), VeyronError> {
if let Some(key) = &self.session_key {
let valid = frame.flags & FLAG_MAC_PRESENT != 0
&& match &frame.mac {
Some(tag) => {
let header = serialize_header(frame);
verify_tag(key, &header, &frame.payload, tag)
}
None => false,
};
if !valid {
return Err(VeyronError::Internal(
"frame MAC verification failed".into(),
));
}
}
Ok(())
}
fn absorb_fragment(&mut self, frame: Frame) -> Result<Option<Frame>, VeyronError> {
self.reassembly
.retain(|_, buf| buf.first_seen.elapsed() < REASSEMBLY_TIMEOUT);
let hdr = parse_frag_header(&frame.payload)
.ok_or_else(|| VeyronError::Internal("fragment header too short".into()))?;
if hdr.total == 0 || hdr.sequence >= hdr.total {
return Err(VeyronError::Internal(format!(
"invalid fragment header: seq {} / total {}",
hdr.sequence, hdr.total
)));
}
if let Some(existing) = self.reassembly.get(&hdr.stream_id) {
if existing.total != hdr.total {
self.reassembly.remove(&hdr.stream_id);
return Err(VeyronError::Internal(
"fragment total mismatch within stream".into(),
));
}
} else if self.reassembly.len() >= MAX_REASSEMBLY_STREAMS {
return Err(VeyronError::Internal(
"too many concurrent fragment streams".into(),
));
}
let chunk = frame.payload[FRAG_HEADER_SIZE..].to_vec();
let entry = self
.reassembly
.entry(hdr.stream_id)
.or_insert_with(|| ReassemblyBuf {
fragments: HashMap::new(),
total: hdr.total,
target: frame.target,
flags: frame.flags & !(FLAG_FRAGMENTED | FLAG_MAC_PRESENT),
first_seen: Instant::now(),
buffered_bytes: 0,
});
let replaced_len = entry.fragments.get(&hdr.sequence).map_or(0, Vec::len);
let new_total = entry.buffered_bytes - replaced_len + chunk.len();
if new_total > MAX_PAYLOAD_SIZE {
self.reassembly.remove(&hdr.stream_id);
return Err(VeyronError::PayloadTooLarge(MAX_PAYLOAD_SIZE + 1));
}
entry.buffered_bytes = new_total;
entry.fragments.insert(hdr.sequence, chunk);
if entry.is_complete() {
let buf = self.reassembly.remove(&hdr.stream_id).unwrap();
let target = buf.target;
let flags = buf.flags;
let payload = buf.reassemble();
let crc32 = crc32fast::hash(&payload);
return Ok(Some(Frame {
magic: 0x5652,
flags,
length: payload.len() as u32,
target,
crc32,
payload: payload.into(),
mac: None,
}));
}
Ok(None)
}
pub async fn subscribe(&mut self, event_types: Vec<String>) -> Result<(), VeyronError> {
let env = Envelope {
payload: Some(envelope::Payload::Subscribe(Subscribe { event_types })),
..Default::default()
};
self.send("kernel", env).await
}
pub async fn unsubscribe(&mut self, event_types: Vec<String>) -> Result<(), VeyronError> {
let env = Envelope {
payload: Some(envelope::Payload::Unsubscribe(Unsubscribe { event_types })),
..Default::default()
};
self.send("kernel", env).await
}
pub async fn ack_event(&mut self, event_id: &str) -> Result<(), VeyronError> {
let env = Envelope {
payload: Some(envelope::Payload::EventAck(EventAck {
event_id: event_id.to_string(),
})),
..Default::default()
};
self.send("kernel", env).await
}
pub async fn publish_event(
&mut self,
event_type: &str,
payload_json: &[u8],
timeout_ms: u32,
) -> Result<EventPublishAck, VeyronError> {
let env = Envelope {
payload: Some(envelope::Payload::EventPublish(EventPublish {
event_type: event_type.to_string(),
payload_json: payload_json.to_vec(),
})),
..Default::default()
};
self.send("kernel", env).await?;
let timeout = if timeout_ms == 0 {
DEFAULT_REQUEST_TIMEOUT
} else {
Duration::from_millis(timeout_ms as u64)
};
let deadline = Instant::now() + timeout;
loop {
let remaining = deadline.saturating_duration_since(Instant::now());
if remaining.is_zero() {
return Err(VeyronError::Timeout);
}
let response = self.recv_timeout(remaining).await?;
match response.payload {
Some(envelope::Payload::EventPublishAck(ack)) => return Ok(ack),
Some(envelope::Payload::Error(err)) => {
return Err(VeyronError::Internal(format!(
"kernel error: {} ({})",
err.message, err.details
)));
}
_ => continue, }
}
}
pub async fn send_action(
&mut self,
action: &str,
params_json: &[u8],
timeout_ms: u32,
) -> Result<ActionResponse, VeyronError> {
let action_id = next_request_id("act");
let env = Envelope {
payload: Some(envelope::Payload::ActionRequest(ActionRequest {
action_id: action_id.clone(),
action: action.to_string(),
params_json: params_json.to_vec(),
timeout_ms,
streaming: false,
})),
..Default::default()
};
self.send("kernel", env).await?;
let timeout = if timeout_ms == 0 {
DEFAULT_REQUEST_TIMEOUT
} else {
Duration::from_millis(timeout_ms as u64)
};
let deadline = Instant::now() + timeout;
loop {
let remaining = deadline.saturating_duration_since(Instant::now());
if remaining.is_zero() {
return Err(VeyronError::Timeout);
}
let response = self.recv_timeout(remaining).await?;
match response.payload {
Some(envelope::Payload::ActionResponse(resp)) if resp.action_id == action_id => {
return Ok(resp);
}
Some(envelope::Payload::ActionStreamAbort(abort))
if abort.action_id == action_id =>
{
return Err(VeyronError::Internal(format!(
"stream aborted: {}",
abort.reason
)));
}
Some(envelope::Payload::Error(err)) => {
return Err(VeyronError::Internal(format!(
"kernel error: {} ({})",
err.message, err.details
)));
}
_ => continue, }
}
}
pub async fn send_action_streaming(
&mut self,
action: &str,
timeout_ms: u32,
) -> Result<String, VeyronError> {
let action_id = next_request_id("act");
let env = Envelope {
payload: Some(envelope::Payload::ActionRequest(ActionRequest {
action_id: action_id.clone(),
action: action.to_string(),
params_json: vec![],
timeout_ms,
streaming: true,
})),
..Default::default()
};
self.send("kernel", env).await?;
Ok(action_id)
}
pub async fn send_request_chunk(
&mut self,
action_id: &str,
seq: u32,
chunk: Vec<u8>,
is_final: bool,
) -> Result<(), VeyronError> {
let env = Envelope {
payload: Some(envelope::Payload::ActionRequestChunk(ActionRequestChunk {
action_id: action_id.to_string(),
seq,
chunk,
r#final: is_final,
})),
..Default::default()
};
self.send("kernel", env).await
}
pub async fn send_response_chunk(
&mut self,
action_id: &str,
seq: u32,
chunk: Vec<u8>,
) -> Result<(), VeyronError> {
let env = Envelope {
payload: Some(envelope::Payload::ActionResponseChunk(
ActionResponseChunk {
action_id: action_id.to_string(),
seq,
chunk,
},
)),
..Default::default()
};
self.send("kernel", env).await
}
pub async fn close_session(
&mut self,
action_id: &str,
reason: &str,
) -> Result<(), VeyronError> {
let env = Envelope {
payload: Some(envelope::Payload::SessionClose(SessionClose {
action_id: action_id.to_string(),
reason: reason.to_string(),
})),
..Default::default()
};
self.send("kernel", env).await
}
pub async fn send_command(
&mut self,
command_id: &str,
command: &str,
params_json: &[u8],
) -> Result<KernelCommandAck, VeyronError> {
let env = Envelope {
payload: Some(envelope::Payload::KernelCommand(KernelCommand {
command_id: command_id.to_string(),
command: command.to_string(),
params_json: params_json.to_vec(),
})),
..Default::default()
};
self.send("kernel", env).await?;
let response = self.recv().await?;
match response.payload {
Some(envelope::Payload::KernelCommandAck(ack)) => Ok(ack),
_ => Err(VeyronError::Internal("expected KernelCommandAck".into())),
}
}
pub async fn ping(&mut self) -> Result<Duration, VeyronError> {
let start = Instant::now();
let env = Envelope {
payload: Some(envelope::Payload::Ping(Ping {
timestamp: unix_millis(),
})),
..Default::default()
};
self.send("kernel", env).await?;
let response = self.recv().await?;
match response.payload {
Some(envelope::Payload::Pong(_)) => Ok(start.elapsed()),
_ => Err(VeyronError::Internal("expected Pong".into())),
}
}
pub async fn send_audio_chunk(
&mut self,
target: &str,
chunk: AudioStreamChunk,
) -> Result<(), VeyronError> {
let env = Envelope {
payload: Some(envelope::Payload::AudioStreamChunk(chunk)),
..Default::default()
};
self.send(target, env).await
}
pub async fn send_raw_audio(&mut self, target: &str, data: Vec<u8>) -> Result<(), VeyronError> {
self.send_raw_with_flags(target, FLAG_RAW_BINARY, data)
.await
}
}
fn build_frame(target: &str, flags: u16, payload: Vec<u8>) -> Frame {
let mut t = [0u8; 32];
let b = target.as_bytes();
let n = b.len().min(32);
t[..n].copy_from_slice(&b[..n]);
Frame {
magic: 0x5652,
flags,
length: payload.len() as u32,
target: t,
crc32: crc32fast::hash(&payload),
payload: payload.into(),
mac: None,
}
}
fn unix_millis() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_millis() as u64)
.unwrap_or(0)
}
fn next_request_id(prefix: &str) -> String {
let seq = ACTION_SEQ.fetch_add(1, Ordering::Relaxed);
format!("{prefix}-{}-{seq}", unix_millis())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn build_frame_truncates_long_target() {
let long = "x".repeat(64);
let frame = build_frame(&long, 0, vec![1, 2, 3]);
assert_eq!(frame.target, [b'x'; 32]);
assert_eq!(frame.length, 3);
assert_eq!(frame.crc32, crc32fast::hash(&[1, 2, 3]));
}
#[test]
fn request_ids_are_unique() {
let a = next_request_id("act");
let b = next_request_id("act");
assert_ne!(a, b);
}
}