Rust port of OmniVox, wire-compatible with the C++ original
Byte-exact cereal binary layout, RFC 8439 ChaCha20-Poly1305, RFC 7914 scrypt defaults; tokio transport preserving the original's observable behavior (contract B1-B3, B5 quirks included). Parity evidence: golden C++ vectors in tests/vectors.rs and loopback behavior tests; both cross-directions verified against the unmodified C++ binaries. Source: https://git.0xfeedc0de.com/michael/OmniVox
This commit is contained in:
+147
@@ -0,0 +1,147 @@
|
||||
//! Byte-exact codec for the cereal `BinaryOutputArchive` wire format.
|
||||
//!
|
||||
//! cereal writes raw little-endian fields with no header or magic:
|
||||
//! PODs as native LE bytes, fixed arrays as raw bytes, and
|
||||
//! `std::vector`/`std::string` as a `size_t` (u64) length followed by the
|
||||
//! bytes. Trailing bytes after the last field are ignored on read, exactly
|
||||
//! as cereal's stream reads ignore them.
|
||||
|
||||
use thiserror::Error;
|
||||
|
||||
/// Errors produced while decoding cereal-encoded bytes.
|
||||
#[derive(Debug, Error, PartialEq, Eq)]
|
||||
pub enum CodecError {
|
||||
/// Input ended before all fields were read.
|
||||
#[error("unexpected end of input at byte {offset} of {len}")]
|
||||
UnexpectedEof { offset: usize, len: usize },
|
||||
/// A vector length prefix exceeded the remaining input.
|
||||
#[error("vector length {declared} exceeds remaining input {available}")]
|
||||
LengthOverflow { declared: u64, available: usize },
|
||||
}
|
||||
|
||||
/// A type with a byte-exact cereal binary representation.
|
||||
///
|
||||
/// Implementations must write exactly the same bytes cereal's
|
||||
/// `BinaryOutputArchive` would write for the equivalent C++ `serialize`
|
||||
/// function, in field order.
|
||||
pub trait Cereal {
|
||||
/// Appends this value's cereal encoding to `out`.
|
||||
fn encode(&self, out: &mut Vec<u8>);
|
||||
|
||||
/// Decodes a value, ignoring any trailing bytes.
|
||||
fn decode(bytes: &[u8]) -> Result<Self, CodecError>
|
||||
where
|
||||
Self: Sized;
|
||||
}
|
||||
|
||||
/// Streaming reader over a byte slice, little-endian.
|
||||
pub struct Reader<'a> {
|
||||
bytes: &'a [u8],
|
||||
pos: usize,
|
||||
}
|
||||
|
||||
impl<'a> Reader<'a> {
|
||||
pub fn new(bytes: &'a [u8]) -> Self {
|
||||
Self { bytes, pos: 0 }
|
||||
}
|
||||
|
||||
pub fn take<const N: usize>(&mut self) -> Result<[u8; N], CodecError> {
|
||||
let end = self.pos + N;
|
||||
if end > self.bytes.len() {
|
||||
return Err(CodecError::UnexpectedEof {
|
||||
offset: self.pos,
|
||||
len: self.bytes.len(),
|
||||
});
|
||||
}
|
||||
let mut out = [0u8; N];
|
||||
out.copy_from_slice(&self.bytes[self.pos..end]);
|
||||
self.pos = end;
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
pub fn u16(&mut self) -> Result<u16, CodecError> {
|
||||
Ok(u16::from_le_bytes(self.take()?))
|
||||
}
|
||||
|
||||
pub fn u32(&mut self) -> Result<u32, CodecError> {
|
||||
Ok(u32::from_le_bytes(self.take()?))
|
||||
}
|
||||
|
||||
pub fn u64(&mut self) -> Result<u64, CodecError> {
|
||||
Ok(u64::from_le_bytes(self.take()?))
|
||||
}
|
||||
|
||||
/// Reads a cereal `std::vector<uint8_t>`: u64 length, then bytes.
|
||||
pub fn byte_vec(&mut self) -> Result<Vec<u8>, CodecError> {
|
||||
let declared = self.u64()?;
|
||||
let available = self.bytes.len() - self.pos;
|
||||
let len = usize::try_from(declared).map_err(|_| CodecError::LengthOverflow {
|
||||
declared,
|
||||
available,
|
||||
})?;
|
||||
if len > available {
|
||||
return Err(CodecError::LengthOverflow {
|
||||
declared,
|
||||
available,
|
||||
});
|
||||
}
|
||||
let out = self.bytes[self.pos..self.pos + len].to_vec();
|
||||
self.pos += len;
|
||||
Ok(out)
|
||||
}
|
||||
}
|
||||
|
||||
/// Pushes a cereal `std::vector<uint8_t>`: u64 LE length, then bytes.
|
||||
pub fn write_byte_vec(out: &mut Vec<u8>, bytes: &[u8]) {
|
||||
out.extend_from_slice(&(bytes.len() as u64).to_le_bytes());
|
||||
out.extend_from_slice(bytes);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn reader_matches_cereal_field_layout() {
|
||||
// 3412 8877665544332211 0500000000000000 48656c6c6f
|
||||
// (golden: cereal_datagram from tests/vectors.txt)
|
||||
let bytes = [
|
||||
0x34, 0x12, 0x88, 0x77, 0x66, 0x55, 0x44, 0x33, 0x22, 0x11, 0x05, 0x00, 0x00, 0x00,
|
||||
0x00, 0x00, 0x00, 0x00, 0x48, 0x65, 0x6c, 0x6c, 0x6f,
|
||||
];
|
||||
let mut r = Reader::new(&bytes);
|
||||
assert_eq!(r.u16().unwrap(), 0x1234);
|
||||
assert_eq!(r.u64().unwrap(), 0x1122334455667788);
|
||||
assert_eq!(r.byte_vec().unwrap(), b"Hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trailing_bytes_are_ignored() {
|
||||
let mut bytes = vec![0x01, 0x00];
|
||||
bytes.extend_from_slice(b"junk");
|
||||
let mut r = Reader::new(&bytes);
|
||||
assert_eq!(r.u16().unwrap(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncated_input_errors() {
|
||||
let mut r = Reader::new(&[0x01]);
|
||||
assert!(matches!(
|
||||
r.u16(),
|
||||
Err(CodecError::UnexpectedEof { offset: 0, len: 1 })
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lying_length_prefix_errors_instead_of_allocating() {
|
||||
let mut bytes = Vec::new();
|
||||
write_byte_vec(&mut bytes, b"hi");
|
||||
let declared = u64::MAX;
|
||||
bytes[..8].copy_from_slice(&declared.to_le_bytes());
|
||||
let mut r = Reader::new(&bytes);
|
||||
assert!(matches!(
|
||||
r.byte_vec(),
|
||||
Err(CodecError::LengthOverflow { .. })
|
||||
));
|
||||
}
|
||||
}
|
||||
+138
@@ -0,0 +1,138 @@
|
||||
//! Authenticated encryption recovered from the C++ CryptoPP usage:
|
||||
//! scrypt key derivation (RFC 7914) and ChaCha20-Poly1305 (RFC 8439),
|
||||
//! byte-compatible with CryptoPP's `Scrypt` and `ChaCha20Poly1305`.
|
||||
|
||||
use chacha20poly1305::{
|
||||
AeadInPlace, ChaCha20Poly1305, Key, KeyInit, aead::generic_array::GenericArray,
|
||||
};
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::codec::{Cereal, CodecError, Reader, write_byte_vec};
|
||||
|
||||
/// 256-bit ChaCha20-Poly1305 key.
|
||||
pub const KEY_LEN: usize = 32;
|
||||
/// 96-bit IETF nonce.
|
||||
pub const NONCE_LEN: usize = 12;
|
||||
/// 128-bit Poly1305 tag.
|
||||
pub const TAG_LEN: usize = 16;
|
||||
|
||||
/// scrypt cost parameter used by the C++ default (N = 2^17).
|
||||
pub const SCRYPT_LOG_N: u8 = 17;
|
||||
/// scrypt block size used by the C++ default.
|
||||
pub const SCRYPT_R: u32 = 8;
|
||||
/// scrypt parallelism used by the C++ default.
|
||||
pub const SCRYPT_P: u32 = 1;
|
||||
|
||||
/// Errors from key derivation or AEAD operations.
|
||||
#[derive(Debug, Error, PartialEq, Eq)]
|
||||
pub enum CryptoError {
|
||||
/// C++ threw `std::invalid_argument`; carried as an error instead.
|
||||
#[error("Passphrase and salt must be non-empty")]
|
||||
EmptyInput,
|
||||
/// The C++ port of `DecryptAndVerify` returning false.
|
||||
#[error("Authentication failed: message has been tampered with")]
|
||||
AuthenticationFailed,
|
||||
#[error("scrypt parameters rejected: {0}")]
|
||||
ScryptParams(String),
|
||||
}
|
||||
|
||||
/// Derives a ChaCha20 key from a passphrase and salt with scrypt.
|
||||
///
|
||||
/// Matches CryptoPP `Scrypt::DeriveKey(key, 32, passphrase, salt, N, r, p)`
|
||||
/// with the C++ defaults `N = 2^17, r = 8, p = 1`.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// [`CryptoError::EmptyInput`] if either input is empty, or
|
||||
/// [`CryptoError::ScryptParams`] if the parameters are rejected.
|
||||
pub fn derive_encrypt_key(
|
||||
passphrase: &[u8],
|
||||
salt: &[u8],
|
||||
log_n: u8,
|
||||
r: u32,
|
||||
p: u32,
|
||||
) -> Result<[u8; KEY_LEN], CryptoError> {
|
||||
if passphrase.is_empty() || salt.is_empty() {
|
||||
return Err(CryptoError::EmptyInput);
|
||||
}
|
||||
let params = scrypt::Params::new(log_n, r, p, KEY_LEN)
|
||||
.map_err(|e| CryptoError::ScryptParams(e.to_string()))?;
|
||||
let mut key = [0u8; KEY_LEN];
|
||||
scrypt::scrypt(passphrase, salt, ¶ms, &mut key)
|
||||
.map_err(|e| CryptoError::ScryptParams(e.to_string()))?;
|
||||
Ok(key)
|
||||
}
|
||||
|
||||
/// An encrypted message on the wire: nonce, detached tag, ciphertext.
|
||||
///
|
||||
/// Serializes in cereal byte order: `nonce[12] || tag[16] || u64 len || payload`.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AeadPacket {
|
||||
pub nonce: [u8; NONCE_LEN],
|
||||
pub tag: [u8; TAG_LEN],
|
||||
pub payload: Vec<u8>,
|
||||
}
|
||||
|
||||
impl Cereal for AeadPacket {
|
||||
fn encode(&self, out: &mut Vec<u8>) {
|
||||
out.extend_from_slice(&self.nonce);
|
||||
out.extend_from_slice(&self.tag);
|
||||
write_byte_vec(out, &self.payload);
|
||||
}
|
||||
|
||||
fn decode(bytes: &[u8]) -> Result<Self, CodecError> {
|
||||
let mut r = Reader::new(bytes);
|
||||
Ok(Self {
|
||||
nonce: r.take()?,
|
||||
tag: r.take()?,
|
||||
payload: r.byte_vec()?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Encrypts `plaintext` with a fresh random 96-bit nonce and no AAD,
|
||||
/// matching the C++ `encrypt` (CryptoPP `EncryptAndAuthenticate`).
|
||||
pub fn encrypt(key: &[u8; KEY_LEN], plaintext: &[u8]) -> AeadPacket {
|
||||
let mut nonce = [0u8; NONCE_LEN];
|
||||
getrandom::fill(&mut nonce).expect("OS RNG failure");
|
||||
encrypt_with_nonce(key, &nonce, plaintext)
|
||||
}
|
||||
|
||||
/// Deterministic variant for tests; production traffic uses [`encrypt`].
|
||||
pub fn encrypt_with_nonce(
|
||||
key: &[u8; KEY_LEN],
|
||||
nonce: &[u8; NONCE_LEN],
|
||||
plaintext: &[u8],
|
||||
) -> AeadPacket {
|
||||
let cipher = ChaCha20Poly1305::new(Key::from_slice(key));
|
||||
let mut payload = plaintext.to_vec();
|
||||
let tag = cipher
|
||||
.encrypt_in_place_detached(GenericArray::from_slice(nonce), b"", &mut payload)
|
||||
.expect("ChaCha20Poly1305 detached encryption is infallible");
|
||||
let mut packet_tag = [0u8; TAG_LEN];
|
||||
packet_tag.copy_from_slice(tag.as_slice());
|
||||
AeadPacket {
|
||||
nonce: *nonce,
|
||||
tag: packet_tag,
|
||||
payload,
|
||||
}
|
||||
}
|
||||
|
||||
/// Decrypts and verifies a packet.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// [`CryptoError::AuthenticationFailed`] if the tag does not verify.
|
||||
pub fn decrypt(key: &[u8; KEY_LEN], packet: &AeadPacket) -> Result<Vec<u8>, CryptoError> {
|
||||
let cipher = ChaCha20Poly1305::new(Key::from_slice(key));
|
||||
let mut payload = packet.payload.clone();
|
||||
cipher
|
||||
.decrypt_in_place_detached(
|
||||
GenericArray::from_slice(&packet.nonce),
|
||||
b"",
|
||||
&mut payload,
|
||||
GenericArray::from_slice(&packet.tag),
|
||||
)
|
||||
.map_err(|_| CryptoError::AuthenticationFailed)?;
|
||||
Ok(payload)
|
||||
}
|
||||
+618
@@ -0,0 +1,618 @@
|
||||
//! OmniVox — encrypted async TCP client/server messaging.
|
||||
//!
|
||||
//! Rust port of the OmniVox C++ library, wire-compatible byte-for-byte with
|
||||
//! the original (framing, cereal field layout, scrypt + ChaCha20-Poly1305).
|
||||
//! See `PARITY-CONTRACT.md` at the workspace root for the normative spec,
|
||||
//! including the quirks this port deliberately preserves (B1, B2, B5).
|
||||
//!
|
||||
//! # Example
|
||||
//!
|
||||
//! ```no_run
|
||||
//! use omnivox::{Cereal, CodecError, ConnectionContext, Decision, OmniVox, OmniVoxError};
|
||||
//! use std::sync::Arc;
|
||||
//!
|
||||
//! #[derive(Default)]
|
||||
//! struct Ping;
|
||||
//!
|
||||
//! impl Cereal for Ping {
|
||||
//! fn encode(&self, out: &mut Vec<u8>) { out.push(1); }
|
||||
//! fn decode(_: &[u8]) -> Result<Self, CodecError> { Ok(Ping) }
|
||||
//! }
|
||||
//!
|
||||
//! # async fn run() -> Result<(), OmniVoxError> {
|
||||
//! let server = Arc::new(OmniVox::<Ping>::new());
|
||||
//! server.set_encrypt_key([7u8; 32]);
|
||||
//! server
|
||||
//! .start_server(
|
||||
//! 8080,
|
||||
//! 2,
|
||||
//! 15000,
|
||||
//! Arc::new(|_data: &Ping, _ctx: &ConnectionContext| {}),
|
||||
//! Arc::new(|_ctx| {}),
|
||||
//! Arc::new(|_ctx| Decision::Allow),
|
||||
//! )
|
||||
//! .await?;
|
||||
//! # Ok(()) }
|
||||
//! ```
|
||||
|
||||
mod codec;
|
||||
mod crypto;
|
||||
|
||||
pub use codec::{Cereal, CodecError, Reader, write_byte_vec};
|
||||
pub use crypto::{
|
||||
AeadPacket, CryptoError, KEY_LEN, NONCE_LEN, SCRYPT_LOG_N, SCRYPT_P, SCRYPT_R, TAG_LEN,
|
||||
decrypt, derive_encrypt_key, encrypt, encrypt_with_nonce,
|
||||
};
|
||||
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
net::SocketAddr,
|
||||
sync::{
|
||||
Arc, Mutex, OnceLock,
|
||||
atomic::{AtomicBool, AtomicU32, Ordering},
|
||||
},
|
||||
time::{SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use tokio::{
|
||||
io::{AsyncReadExt, AsyncWriteExt},
|
||||
net::{
|
||||
TcpListener, TcpStream,
|
||||
tcp::{OwnedReadHalf, OwnedWriteHalf},
|
||||
},
|
||||
sync::Mutex as AsyncMutex,
|
||||
task::JoinHandle,
|
||||
time::{Duration, timeout},
|
||||
};
|
||||
use tracing::{error, info};
|
||||
|
||||
/// 5 MiB per-message allocation cap, matching `MAX_MESSAGE_SIZE`.
|
||||
pub const MAX_MESSAGE_SIZE: u32 = 1024 * 1024 * 5;
|
||||
/// Fixed body-read stall cap, matching `MAX_TIMEOUT_STALL_MS`.
|
||||
pub const MAX_TIMEOUT_STALL_MS: u64 = 5000;
|
||||
|
||||
/// State of a connection, matching the C++ `ConnectionStatus`.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ConnectionStatus {
|
||||
NotConnected,
|
||||
Connected,
|
||||
ConnectionFailed,
|
||||
Rejected,
|
||||
}
|
||||
|
||||
/// Per-connection metadata handed to every callback.
|
||||
///
|
||||
/// `total_sent` mirrors the C++ `iTotalSent`, which the original never
|
||||
/// increments; it is preserved as always-zero for parity.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ConnectionContext {
|
||||
pub connected: ConnectionStatus,
|
||||
pub id: u32,
|
||||
pub total_sent: usize,
|
||||
pub creation_time: u64,
|
||||
pub ip: String,
|
||||
pub port: u16,
|
||||
}
|
||||
|
||||
impl ConnectionContext {
|
||||
fn unknown(host: &str, port: u16) -> Self {
|
||||
Self {
|
||||
connected: ConnectionStatus::ConnectionFailed,
|
||||
id: 0,
|
||||
total_sent: 0,
|
||||
creation_time: 0,
|
||||
ip: host.to_owned(),
|
||||
port,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Server-side accept decision, matching `ConnectionDecision`.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum Decision {
|
||||
Allow,
|
||||
Deny,
|
||||
}
|
||||
|
||||
pub type DataHandler<T> = Arc<dyn Fn(&T, &ConnectionContext) + Send + Sync + 'static>;
|
||||
pub type ConnHandler = Arc<dyn Fn(&ConnectionContext) + Send + Sync + 'static>;
|
||||
pub type DecisionHandler = Arc<dyn Fn(&ConnectionContext) -> Decision + Send + Sync + 'static>;
|
||||
|
||||
/// Library errors surfaced from setup calls. Runtime message failures are
|
||||
/// logged and swallowed, matching the C++ exception swallowing (contract B2).
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum OmniVoxError {
|
||||
#[error(transparent)]
|
||||
Io(#[from] std::io::Error),
|
||||
#[error(transparent)]
|
||||
Crypto(#[from] CryptoError),
|
||||
#[error("already started")]
|
||||
AlreadyStarted,
|
||||
}
|
||||
|
||||
fn now_micros() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_micros() as u64)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
struct ClientHandle {
|
||||
/// Held across the full frame write, mirroring the C++ clients-mutex
|
||||
/// hold across the blocking `asio::write` (contract B3). `Option` so a
|
||||
/// disconnect can close the socket even while the read task still holds
|
||||
/// the handle, matching the C++ `handleDisconnect` socket close.
|
||||
writer: AsyncMutex<Option<OwnedWriteHalf>>,
|
||||
context: ConnectionContext,
|
||||
}
|
||||
|
||||
struct Inner<T> {
|
||||
key: Mutex<Option<[u8; KEY_LEN]>>,
|
||||
clients: Mutex<HashMap<u32, Arc<ClientHandle>>>,
|
||||
next_id: AtomicU32,
|
||||
stopping: AtomicBool,
|
||||
/// Client-mode connection state; id 0 means "never connected", so the
|
||||
/// C++ `ctx.id != serverContext.id` timer-gating reproduces exactly.
|
||||
server_ctx: Mutex<Option<ConnectionContext>>,
|
||||
server_writer: AsyncMutex<Option<OwnedWriteHalf>>,
|
||||
received: OnceLock<DataHandler<T>>,
|
||||
disconnected: OnceLock<ConnHandler>,
|
||||
connect_decision: OnceLock<DecisionHandler>,
|
||||
server_connected: OnceLock<ConnHandler>,
|
||||
tasks: Mutex<Vec<JoinHandle<()>>>,
|
||||
}
|
||||
|
||||
/// Encrypted TCP endpoint usable in either client or server role.
|
||||
///
|
||||
/// Cheap to hand out: state is shared behind an `Arc`; clone or borrow the
|
||||
/// same instance from callbacks.
|
||||
pub struct OmniVox<T> {
|
||||
inner: Arc<Inner<T>>,
|
||||
}
|
||||
|
||||
impl<T: Cereal + Send + Sync + 'static> Default for OmniVox<T> {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Cereal + Send + Sync + 'static> OmniVox<T> {
|
||||
/// Creates an endpoint with no key set (traffic then uses the all-zero
|
||||
/// key, reproducing the C++ default-constructed `SecByteBlock`).
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
inner: Arc::new(Inner {
|
||||
key: Mutex::new(None),
|
||||
clients: Mutex::new(HashMap::new()),
|
||||
next_id: AtomicU32::new(1),
|
||||
stopping: AtomicBool::new(false),
|
||||
server_ctx: Mutex::new(None),
|
||||
server_writer: AsyncMutex::const_new(None),
|
||||
received: OnceLock::new(),
|
||||
disconnected: OnceLock::new(),
|
||||
connect_decision: OnceLock::new(),
|
||||
server_connected: OnceLock::new(),
|
||||
tasks: Mutex::new(Vec::new()),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Sets the ChaCha20-Poly1305 key. Until set, traffic uses the all-zero
|
||||
/// key, which is what the C++ default-constructed `SecByteBlock` did.
|
||||
pub fn set_encrypt_key(&self, key: [u8; KEY_LEN]) {
|
||||
*self.inner.key.lock().unwrap() = Some(key);
|
||||
}
|
||||
|
||||
fn key_bytes(&self) -> [u8; KEY_LEN] {
|
||||
self.inner.key.lock().unwrap().unwrap_or([0u8; KEY_LEN])
|
||||
}
|
||||
|
||||
/// Starts the server on `0.0.0.0:port`.
|
||||
///
|
||||
/// `num_threads` is accepted for API parity with the C++ signature; the
|
||||
/// tokio runtime's worker pool already provides the parallelism the C++
|
||||
/// `ThreadPool_` existed to add, so it configures nothing (contract:
|
||||
/// thread structure is an accepted difference).
|
||||
///
|
||||
/// `timeout_ms` is the per-client idle cap for waiting on a message
|
||||
/// header. Errors only if already started; runtime failures are logged.
|
||||
pub async fn start_server(
|
||||
&self,
|
||||
port: u16,
|
||||
num_threads: usize,
|
||||
timeout_ms: u64,
|
||||
data: DataHandler<T>,
|
||||
disconnect: ConnHandler,
|
||||
connect: DecisionHandler,
|
||||
) -> Result<(), OmniVoxError> {
|
||||
let _ = num_threads;
|
||||
let _ = self.inner.received.set(data);
|
||||
let _ = self.inner.disconnected.set(disconnect);
|
||||
let _ = self.inner.connect_decision.set(connect);
|
||||
|
||||
let listener = TcpListener::bind(("0.0.0.0", port)).await?;
|
||||
let accept_inner = Arc::clone(&self.inner);
|
||||
let accept_task = tokio::spawn(async move {
|
||||
loop {
|
||||
let (stream, peer) = match listener.accept().await {
|
||||
Ok(v) => v,
|
||||
Err(_) if accept_inner.stopping.load(Ordering::Relaxed) => break,
|
||||
Err(e) => {
|
||||
error!("[OMNIVOX] accept error: {e}");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let conn_inner = Arc::clone(&accept_inner);
|
||||
let task = tokio::spawn(async move {
|
||||
Self::serve_connection(conn_inner, stream, peer, timeout_ms).await;
|
||||
});
|
||||
accept_inner.tasks.lock().unwrap().push(task);
|
||||
}
|
||||
});
|
||||
self.inner.tasks.lock().unwrap().push(accept_task);
|
||||
|
||||
info!(
|
||||
"[OMNIVOX] Server started on port {port} with {num_threads} threads and {timeout_ms} ms timeout"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn serve_connection(
|
||||
inner: Arc<Inner<T>>,
|
||||
mut stream: TcpStream,
|
||||
peer: SocketAddr,
|
||||
timeout_ms: u64,
|
||||
) {
|
||||
let id = inner.next_id.fetch_add(1, Ordering::Relaxed);
|
||||
let mut ctx = ConnectionContext {
|
||||
connected: ConnectionStatus::Connected,
|
||||
id,
|
||||
total_sent: 0,
|
||||
creation_time: now_micros(),
|
||||
ip: peer.ip().to_string(),
|
||||
port: peer.port(),
|
||||
};
|
||||
|
||||
let decision = inner
|
||||
.connect_decision
|
||||
.get()
|
||||
.map(|f| f(&ctx))
|
||||
.unwrap_or(Decision::Allow);
|
||||
|
||||
if decision == Decision::Deny {
|
||||
let _ = stream.shutdown().await;
|
||||
ctx.connected = ConnectionStatus::Rejected;
|
||||
if let Some(f) = inner.disconnected.get() {
|
||||
f(&ctx);
|
||||
}
|
||||
info!(
|
||||
"[OMNIVOX] Client connection from {}:{} rejected",
|
||||
ctx.ip, ctx.port
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
let (mut reader, writer) = stream.into_split();
|
||||
let handle = Arc::new(ClientHandle {
|
||||
writer: AsyncMutex::new(Some(writer)),
|
||||
context: ctx.clone(),
|
||||
});
|
||||
inner
|
||||
.clients
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert(id, Arc::clone(&handle));
|
||||
info!(
|
||||
"[OMNIVOX] Client {id} connected from {}:{}",
|
||||
ctx.ip, ctx.port
|
||||
);
|
||||
|
||||
Self::read_loop(inner, handle, &mut reader, timeout_ms).await;
|
||||
}
|
||||
|
||||
async fn read_loop(
|
||||
inner: Arc<Inner<T>>,
|
||||
handle: Arc<ClientHandle>,
|
||||
reader: &mut OwnedReadHalf,
|
||||
timeout_ms: u64,
|
||||
) {
|
||||
let ctx = &handle.context;
|
||||
loop {
|
||||
// Header: u32 LE message length, bounded by the client timeout.
|
||||
let mut header = [0u8; 4];
|
||||
match timeout(
|
||||
Duration::from_millis(timeout_ms),
|
||||
reader.read_exact(&mut header),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Err(_elapsed) => {
|
||||
info!(
|
||||
"[OMNIVOX] Client {} {}:{} Timed out",
|
||||
ctx.id, ctx.ip, ctx.port
|
||||
);
|
||||
Self::handle_client_disconnect(&inner, ctx).await;
|
||||
return;
|
||||
}
|
||||
Ok(Err(_)) => {
|
||||
Self::handle_client_disconnect(&inner, ctx).await;
|
||||
return;
|
||||
}
|
||||
Ok(Ok(_)) => {}
|
||||
}
|
||||
|
||||
let size = u32::from_le_bytes(header);
|
||||
// Contract B1, preserved: on an unreasonable size the C++ code
|
||||
// disconnects but does NOT return — it still attempts the body
|
||||
// read (a 0-byte read succeeds and its empty packet then fails
|
||||
// to decrypt, which is logged, not fatal).
|
||||
if size == 0 || size > MAX_MESSAGE_SIZE {
|
||||
info!(
|
||||
"[OMNIVOX] Client {} {}:{} has an unreasonable header size!",
|
||||
ctx.id, ctx.ip, ctx.port
|
||||
);
|
||||
Self::handle_client_disconnect(&inner, ctx).await;
|
||||
}
|
||||
|
||||
// Body: fixed 5 s stall cap against slow-body attacks.
|
||||
let mut body = vec![0u8; size as usize];
|
||||
match timeout(
|
||||
Duration::from_millis(MAX_TIMEOUT_STALL_MS),
|
||||
reader.read_exact(&mut body),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Err(_elapsed) => {
|
||||
info!(
|
||||
"[OMNIVOX] Client {} {}:{} had a stall during body read!",
|
||||
ctx.id, ctx.ip, ctx.port
|
||||
);
|
||||
Self::handle_client_disconnect(&inner, ctx).await;
|
||||
return;
|
||||
}
|
||||
Ok(Err(_)) => {
|
||||
Self::handle_client_disconnect(&inner, ctx).await;
|
||||
return;
|
||||
}
|
||||
Ok(Ok(_)) => {}
|
||||
}
|
||||
|
||||
// Contract B2, preserved: bad packets are logged and skipped;
|
||||
// the connection stays open.
|
||||
if let Err(e) = Self::process_message(&inner, &body, ctx) {
|
||||
error!("[OMNIVOX] Error processing message: {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn process_message(
|
||||
inner: &Inner<T>,
|
||||
body: &[u8],
|
||||
ctx: &ConnectionContext,
|
||||
) -> Result<(), ProcessError> {
|
||||
let packet = AeadPacket::decode(body)?;
|
||||
let plaintext = decrypt(&inner_key(inner), &packet)?;
|
||||
let payload = T::decode(&plaintext)?;
|
||||
if let Some(f) = inner.received.get() {
|
||||
f(&payload, ctx);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_client_disconnect(inner: &Inner<T>, ctx: &ConnectionContext) {
|
||||
let removed = {
|
||||
let mut clients = inner.clients.lock().unwrap();
|
||||
clients.remove(&ctx.id)
|
||||
};
|
||||
if let Some(handle) = removed {
|
||||
if let Some(mut writer) = handle.writer.lock().await.take() {
|
||||
let _ = writer.shutdown().await;
|
||||
}
|
||||
if let Some(f) = inner.disconnected.get() {
|
||||
f(ctx);
|
||||
}
|
||||
info!("[OMNIVOX] Client {} disconnected", ctx.id);
|
||||
}
|
||||
}
|
||||
|
||||
/// Connects to `host:port`, registering callbacks. The read loop runs
|
||||
/// only when `data` is provided — the C++ gate on `receivedCallback`,
|
||||
/// including contract B5 (client mode never arms stall timers).
|
||||
///
|
||||
/// The `connection` callback fires exactly once with the outcome; data
|
||||
/// arriving later reaches `data` on a background task.
|
||||
pub async fn connect_to_server(
|
||||
&self,
|
||||
host: &str,
|
||||
port: u16,
|
||||
data: Option<DataHandler<T>>,
|
||||
connection: ConnHandler,
|
||||
) -> Result<(), OmniVoxError> {
|
||||
if let Some(d) = data {
|
||||
let _ = self.inner.received.set(d);
|
||||
}
|
||||
let _ = self.inner.server_connected.set(connection);
|
||||
|
||||
let ctx = match Self::dial(host, port).await {
|
||||
Err(_) => ConnectionContext::unknown(host, port),
|
||||
Ok(stream) => {
|
||||
let ctx = ConnectionContext {
|
||||
connected: ConnectionStatus::Connected,
|
||||
id: self.inner.next_id.fetch_add(1, Ordering::Relaxed),
|
||||
total_sent: 0,
|
||||
creation_time: now_micros(),
|
||||
ip: host.to_owned(),
|
||||
port,
|
||||
};
|
||||
*self.inner.server_ctx.lock().unwrap() = Some(ctx.clone());
|
||||
let (reader, writer) = stream.into_split();
|
||||
*self.inner.server_writer.lock().await = Some(writer);
|
||||
if self.inner.received.get().is_some() {
|
||||
let inner = Arc::clone(&self.inner);
|
||||
let read_ctx = ctx.clone();
|
||||
self.inner
|
||||
.tasks
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(tokio::spawn(async move {
|
||||
Self::client_read_loop(inner, reader, read_ctx).await;
|
||||
}));
|
||||
}
|
||||
ctx
|
||||
}
|
||||
};
|
||||
|
||||
if let Some(f) = self.inner.server_connected.get() {
|
||||
f(&ctx);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn dial(host: &str, port: u16) -> Result<TcpStream, std::io::Error> {
|
||||
// asio::async_connect tries the resolved list in order and stops at
|
||||
// the first success.
|
||||
let addrs = tokio::net::lookup_host((host, port)).await?;
|
||||
let mut last = Err(std::io::Error::other("no addresses resolved"));
|
||||
for addr in addrs {
|
||||
match TcpStream::connect(addr).await {
|
||||
Ok(s) => return Ok(s),
|
||||
Err(e) => last = Err(e),
|
||||
}
|
||||
}
|
||||
last
|
||||
}
|
||||
|
||||
async fn client_read_loop(
|
||||
inner: Arc<Inner<T>>,
|
||||
mut reader: OwnedReadHalf,
|
||||
ctx: ConnectionContext,
|
||||
) {
|
||||
loop {
|
||||
let mut header = [0u8; 4];
|
||||
if reader.read_exact(&mut header).await.is_err() {
|
||||
break;
|
||||
}
|
||||
let size = u32::from_le_bytes(header);
|
||||
if size == 0 || size > MAX_MESSAGE_SIZE {
|
||||
info!(
|
||||
"[OMNIVOX] Client {} {}:{} has an unreasonable header size!",
|
||||
ctx.id, ctx.ip, ctx.port
|
||||
);
|
||||
}
|
||||
let mut body = vec![0u8; size as usize];
|
||||
if reader.read_exact(&mut body).await.is_err() {
|
||||
break;
|
||||
}
|
||||
if let Err(e) = Self::process_message(&inner, &body, &ctx) {
|
||||
error!("[OMNIVOX] Error processing message: {e}");
|
||||
}
|
||||
}
|
||||
*inner.server_writer.lock().await = None;
|
||||
if let Some(mut sc) = inner.server_ctx.lock().unwrap().take() {
|
||||
sc.connected = ConnectionStatus::NotConnected;
|
||||
if let Some(f) = inner.server_connected.get() {
|
||||
f(&sc);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Sends to the connected server. `false` if not connected — the C++
|
||||
/// checks `is_open() && status == Connected` and returns, not throws.
|
||||
pub async fn send_to_server(&self, data: &T) -> bool {
|
||||
let connected = self
|
||||
.inner
|
||||
.server_ctx
|
||||
.lock()
|
||||
.unwrap()
|
||||
.as_ref()
|
||||
.is_some_and(|c| c.connected == ConnectionStatus::Connected);
|
||||
if !connected {
|
||||
return false;
|
||||
}
|
||||
let frame = match Self::sealed_frame(&self.key_bytes(), data) {
|
||||
Ok(f) => f,
|
||||
Err(()) => return false,
|
||||
};
|
||||
let mut guard = self.inner.server_writer.lock().await;
|
||||
match guard.as_mut() {
|
||||
Some(writer) => writer.write_all(&frame).await.is_ok(),
|
||||
None => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Sends to client `id`. `false` if unknown — contract B3 preserved: the
|
||||
/// write lock is held across the full frame write, as the C++ held the
|
||||
/// clients mutex across the blocking `asio::write`.
|
||||
pub async fn send_to_client(&self, data: &T, id: u32) -> bool {
|
||||
let client = self.inner.clients.lock().unwrap().get(&id).cloned();
|
||||
let Some(client) = client else { return false };
|
||||
let Ok(frame) = Self::sealed_frame(&self.key_bytes(), data) else {
|
||||
return false;
|
||||
};
|
||||
let mut guard = client.writer.lock().await;
|
||||
match guard.as_mut() {
|
||||
Some(writer) => writer.write_all(&frame).await.is_ok(),
|
||||
None => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn sealed_frame(key: &[u8; KEY_LEN], data: &T) -> Result<Vec<u8>, ()> {
|
||||
let mut serialized = Vec::new();
|
||||
data.encode(&mut serialized);
|
||||
let packet = encrypt(key, &serialized);
|
||||
let mut packet_bytes = Vec::new();
|
||||
packet.encode(&mut packet_bytes);
|
||||
let mut frame = Vec::with_capacity(4 + packet_bytes.len());
|
||||
// C++ memcpy'd a size_t into a u32 slot — truncation kept.
|
||||
frame.extend_from_slice(&(packet_bytes.len() as u32).to_le_bytes());
|
||||
frame.extend_from_slice(&packet_bytes);
|
||||
Ok(frame)
|
||||
}
|
||||
|
||||
/// Stops everything: accept loop, read loops, all sockets. Idempotent.
|
||||
pub async fn stop(&self) {
|
||||
self.inner.stopping.store(true, Ordering::Relaxed);
|
||||
let tasks: Vec<_> = self.inner.tasks.lock().unwrap().drain(..).collect();
|
||||
for t in tasks {
|
||||
t.abort();
|
||||
}
|
||||
let clients: Vec<_> = self
|
||||
.inner
|
||||
.clients
|
||||
.lock()
|
||||
.unwrap()
|
||||
.drain()
|
||||
.map(|(_, c)| c)
|
||||
.collect();
|
||||
drop(clients);
|
||||
let server = self.inner.server_writer.lock().await.take();
|
||||
drop(server);
|
||||
info!("[OMNIVOX] OmniVox stopped");
|
||||
}
|
||||
|
||||
/// Server-side view of a client, for callbacks that want metadata.
|
||||
pub fn client_context(&self, id: u32) -> Option<ConnectionContext> {
|
||||
self.inner
|
||||
.clients
|
||||
.lock()
|
||||
.unwrap()
|
||||
.get(&id)
|
||||
.map(|c| c.context.clone())
|
||||
}
|
||||
|
||||
pub fn server_context(&self) -> Option<ConnectionContext> {
|
||||
self.inner.server_ctx.lock().unwrap().clone()
|
||||
}
|
||||
}
|
||||
|
||||
fn inner_key<T>(inner: &Inner<T>) -> [u8; KEY_LEN] {
|
||||
inner.key.lock().unwrap().unwrap_or([0u8; KEY_LEN])
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
enum ProcessError {
|
||||
#[error(transparent)]
|
||||
Codec(#[from] CodecError),
|
||||
#[error(transparent)]
|
||||
Crypto(#[from] CryptoError),
|
||||
}
|
||||
Reference in New Issue
Block a user