use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use serde::Serialize;
use serde::de::DeserializeOwned;
use crate::codec::{Decode, Encode};
use crate::error::Error;
type BoxedHandler =
Box<dyn Fn(Vec<u8>, &dyn std::any::Any) -> Result<Vec<u8>, Error> + Send + Sync>;
pub struct Router {
handlers: RwLock<HashMap<String, Arc<BoxedHandler>>>,
}
impl Router {
pub fn new() -> Self {
Self {
handlers: RwLock::new(HashMap::new()),
}
}
pub fn register<F>(&self, name: impl Into<String>, handler: F)
where
F: Fn(Vec<u8>) -> Vec<u8> + Send + Sync + 'static,
{
let boxed: BoxedHandler = Box::new(move |payload, _ctx| Ok(handler(payload)));
crate::write_or_recover(&self.handlers).insert(name.into(), Arc::new(boxed));
}
pub fn register_simple<F>(&self, name: impl Into<String>, handler: F)
where
F: Fn() -> Vec<u8> + Send + Sync + 'static,
{
let boxed: BoxedHandler = Box::new(move |_payload, _ctx| Ok(handler()));
crate::write_or_recover(&self.handlers).insert(name.into(), Arc::new(boxed));
}
pub fn register_json<F, A, R>(&self, name: impl Into<String>, handler: F)
where
F: Fn(A) -> R + Send + Sync + 'static,
A: DeserializeOwned + 'static,
R: Serialize + 'static,
{
let boxed: BoxedHandler = Box::new(move |payload, _ctx| {
let arg: A = sonic_rs::from_slice(&payload).map_err(Error::from)?;
let result = handler(arg);
sonic_rs::to_vec(&result).map_err(Error::from)
});
crate::write_or_recover(&self.handlers).insert(name.into(), Arc::new(boxed));
}
pub fn register_json_result<F, A, R, E>(&self, name: impl Into<String>, handler: F)
where
F: Fn(A) -> Result<R, E> + Send + Sync + 'static,
A: DeserializeOwned + 'static,
R: Serialize + 'static,
E: std::fmt::Display + 'static,
{
let boxed: BoxedHandler = Box::new(move |payload, _ctx| {
let arg: A = sonic_rs::from_slice(&payload).map_err(Error::from)?;
let result = handler(arg).map_err(|e| Error::Handler(e.to_string()))?;
sonic_rs::to_vec(&result).map_err(Error::from)
});
crate::write_or_recover(&self.handlers).insert(name.into(), Arc::new(boxed));
}
pub fn register_binary<F, A, R>(&self, name: impl Into<String>, handler: F)
where
F: Fn(A) -> R + Send + Sync + 'static,
A: Decode + 'static,
R: Encode + 'static,
{
let boxed: BoxedHandler = Box::new(move |payload, _ctx| {
let (arg, _consumed) = A::decode(&payload).ok_or(Error::DecodeFailed)?;
let result = handler(arg);
let mut buf = Vec::with_capacity(result.encode_size());
result.encode(&mut buf);
Ok(buf)
});
crate::write_or_recover(&self.handlers).insert(name.into(), Arc::new(boxed));
}
pub fn register_with_context<F>(&self, name: impl Into<String>, handler: F)
where
F: Fn(Vec<u8>, &dyn std::any::Any) -> Result<Vec<u8>, Error> + Send + Sync + 'static,
{
let boxed: BoxedHandler = Box::new(handler);
crate::write_or_recover(&self.handlers).insert(name.into(), Arc::new(boxed));
}
pub fn call_with_context(
&self,
name: &str,
payload: Vec<u8>,
ctx: &dyn std::any::Any,
) -> Result<Vec<u8>, Error> {
let handler = {
let handlers = crate::read_or_recover(&self.handlers);
handlers.get(name).cloned()
};
match handler {
Some(h) => h(payload, ctx),
None => Err(Error::UnknownCommand(name.to_string())),
}
}
#[must_use]
pub fn call_or_error_bytes_with_context(
&self,
name: &str,
payload: Vec<u8>,
ctx: &dyn std::any::Any,
) -> Vec<u8> {
match self.call_with_context(name, payload, ctx) {
Ok(bytes) => bytes,
Err(e) => e.to_string().into_bytes(),
}
}
pub fn call(&self, name: &str, payload: Vec<u8>) -> Result<Vec<u8>, Error> {
self.call_with_context(name, payload, &())
}
#[must_use]
pub fn call_or_error_bytes(&self, name: &str, payload: Vec<u8>) -> Vec<u8> {
self.call_or_error_bytes_with_context(name, payload, &())
}
#[must_use]
pub fn has(&self, name: &str) -> bool {
crate::read_or_recover(&self.handlers).contains_key(name)
}
}
impl std::fmt::Debug for Router {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let count = crate::read_or_recover(&self.handlers).len();
f.debug_struct("Router")
.field("handler_count", &count)
.finish()
}
}
impl Default for Router {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn register_and_dispatch() {
let table = Router::new();
table.register("echo", |payload: Vec<u8>| payload);
let resp = table.call("echo", b"hello".to_vec()).unwrap();
assert_eq!(resp, b"hello");
}
#[test]
fn unknown_command() {
let table = Router::new();
let err = table.call("nope", vec![]).unwrap_err();
assert!(matches!(err, Error::UnknownCommand(ref name) if name == "nope"));
assert_eq!(err.to_string(), "unknown command: nope");
}
#[test]
fn has_command() {
let table = Router::new();
assert!(!table.has("ping"));
table.register("ping", |_payload: Vec<u8>| b"pong".to_vec());
assert!(table.has("ping"));
}
#[test]
fn register_simple_test() {
let table = Router::new();
table.register_simple("version", || b"1.0".to_vec());
let resp = table.call("version", vec![0xFF]).unwrap();
assert_eq!(resp, b"1.0");
}
#[test]
fn call_or_error_bytes_success() {
let table = Router::new();
table.register("echo", |payload: Vec<u8>| payload);
let resp = table.call_or_error_bytes("echo", b"hello".to_vec());
assert_eq!(resp, b"hello");
}
#[test]
fn call_or_error_bytes_unknown() {
let table = Router::new();
let resp = table.call_or_error_bytes("nope", vec![]);
assert_eq!(resp, b"unknown command: nope");
}
#[test]
fn register_json_roundtrip() {
let table = Router::new();
table.register_json("add", |args: (i32, i32)| args.0 + args.1);
let payload = sonic_rs::to_vec(&(3, 4)).unwrap();
let resp = table.call("add", payload).unwrap();
let result: i32 = sonic_rs::from_slice(&resp).unwrap();
assert_eq!(result, 7);
}
#[test]
fn register_json_bad_input() {
let table = Router::new();
table.register_json("add", |args: (i32, i32)| args.0 + args.1);
let err = table.call("add", b"not json!".to_vec()).unwrap_err();
assert!(matches!(err, Error::Serialize(_)));
}
#[test]
fn register_json_result_ok() {
let table = Router::new();
table.register_json_result("divide", |args: (f64, f64)| -> Result<f64, String> {
if args.1 == 0.0 {
Err("division by zero".into())
} else {
Ok(args.0 / args.1)
}
});
let payload = sonic_rs::to_vec(&(10.0_f64, 2.0_f64)).unwrap();
let resp = table.call("divide", payload).unwrap();
let result: f64 = sonic_rs::from_slice(&resp).unwrap();
assert!((result - 5.0).abs() < f64::EPSILON);
}
#[test]
fn register_json_result_err() {
let table = Router::new();
table.register_json_result("divide", |args: (f64, f64)| -> Result<f64, String> {
if args.1 == 0.0 {
Err("division by zero".into())
} else {
Ok(args.0 / args.1)
}
});
let payload = sonic_rs::to_vec(&(10.0_f64, 0.0_f64)).unwrap();
let err = table.call("divide", payload).unwrap_err();
assert!(matches!(err, Error::Handler(ref msg) if msg == "division by zero"));
}
#[test]
fn register_json_result_bad_input() {
let table = Router::new();
table.register_json_result("divide", |args: (f64, f64)| -> Result<f64, String> {
Ok(args.0 / args.1)
});
let err = table.call("divide", b"garbage".to_vec()).unwrap_err();
assert!(matches!(err, Error::Serialize(_)));
}
#[derive(Debug, PartialEq)]
struct Pair(u32, u32);
impl crate::codec::Encode for Pair {
fn encode(&self, buf: &mut Vec<u8>) {
self.0.encode(buf);
self.1.encode(buf);
}
fn encode_size(&self) -> usize {
8
}
}
impl crate::codec::Decode for Pair {
fn decode(data: &[u8]) -> Option<(Self, usize)> {
let (a, ca) = u32::decode(data)?;
let (b, cb) = u32::decode(&data[ca..])?;
Some((Pair(a, b), ca + cb))
}
}
#[test]
fn register_binary_roundtrip() {
let table = Router::new();
table.register_binary("sum", |p: Pair| p.0 + p.1);
let mut payload = Vec::new();
Pair(10, 20).encode(&mut payload);
let resp = table.call("sum", payload).unwrap();
let (result, _) = u32::decode(&resp).unwrap();
assert_eq!(result, 30);
}
#[test]
fn register_binary_bad_input() {
let table = Router::new();
table.register_binary("sum", |p: Pair| p.0 + p.1);
let err = table.call("sum", vec![1, 2, 3]).unwrap_err();
assert!(matches!(err, Error::DecodeFailed));
}
#[test]
fn register_with_context_basic() {
let table = Router::new();
table.register_with_context("echo_ctx", |payload: Vec<u8>, _ctx: &dyn std::any::Any| {
Ok(payload)
});
let resp = table.call("echo_ctx", b"hello".to_vec()).unwrap();
assert_eq!(resp, b"hello");
}
#[test]
fn call_with_context_passes_through() {
let table = Router::new();
table.register_with_context("check_ctx", |_payload: Vec<u8>, ctx: &dyn std::any::Any| {
if ctx.downcast_ref::<String>().is_some() {
Ok(b"got string".to_vec())
} else {
Ok(b"no string".to_vec())
}
});
let ctx = String::from("hello");
let resp = table.call_with_context("check_ctx", vec![], &ctx).unwrap();
assert_eq!(resp, b"got string");
let resp = table.call("check_ctx", vec![]).unwrap();
assert_eq!(resp, b"no string");
}
#[test]
fn call_or_error_bytes_with_context_success() {
let table = Router::new();
table.register("echo", |payload: Vec<u8>| payload);
let ctx = String::from("unused");
let resp = table.call_or_error_bytes_with_context("echo", b"hello".to_vec(), &ctx);
assert_eq!(resp, b"hello");
}
#[test]
fn call_or_error_bytes_with_context_unknown() {
let table = Router::new();
let resp = table.call_or_error_bytes_with_context("nope", vec![], &());
assert_eq!(resp, b"unknown command: nope");
}
#[test]
fn register_replaces_handler() {
let table = Router::new();
table.register("cmd", |_payload: Vec<u8>| b"first".to_vec());
table.register("cmd", |_payload: Vec<u8>| b"second".to_vec());
let resp = table.call("cmd", vec![]).unwrap();
assert_eq!(resp, b"second");
}
#[test]
fn register_with_context_error_propagation() {
let table = Router::new();
table.register_with_context("fail", |_payload: Vec<u8>, _ctx: &dyn std::any::Any| {
Err(Error::Handler("context handler failed".into()))
});
let err = table.call("fail", vec![]).unwrap_err();
assert!(matches!(err, Error::Handler(ref msg) if msg == "context handler failed"));
let bytes = table.call_or_error_bytes("fail", vec![]);
assert_eq!(bytes, b"handler error: context handler failed");
}
}