Add more sophisticated "axum-like" routing

Lots of futzing with generics to allow handler functions of many
different signatures, whose arguments cause the API to automatically
extract and/or parse parts of the request, and whose return values cause
the API to automatically encode responses.

All of this still needs documentation and cleaning up/refactoring into
different modules.

xxdk-rs/Cargo.lock

@@ -566,6 +566,12 @@ dependencies = [
  "either",
 ]
 
+[[package]]
+name = "itoa"
+version = "1.0.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b"
+
 [[package]]
 name = "js-sys"
 version = "0.3.69"
@@ -898,6 +904,12 @@ dependencies = [
  "windows-sys 0.52.0",
 ]
 
+[[package]]
+name = "ryu"
+version = "1.0.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f3cb5ba0dc43242ce17de99c180e96db90b235b8a9fdc9543c96d2209116bd9f"
+
 [[package]]
 name = "serde"
 version = "1.0.202"
@@ -918,6 +930,17 @@ dependencies = [
  "syn 2.0.53",
 ]
 
+[[package]]
+name = "serde_json"
+version = "1.0.120"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4e0d21c9a8cae1235ad58a00c11cb40d4b1e5c784f1ef2c537876ed6ffd8b7c5"
+dependencies = [
+ "itoa",
+ "ryu",
+ "serde",
+]
+
 [[package]]
 name = "sharded-slab"
 version = "0.1.7"
@@ -1431,6 +1454,7 @@ dependencies = [
  "lazy_static",
  "libc",
  "serde",
+ "serde_json",
  "tokio",
  "tower",
  "tracing",

xxdk-rs/rpc-example/src/lib.rs

@@ -5,7 +5,7 @@ use std::sync::Arc;
 use base64::prelude::*;
 use structopt::StructOpt;
 use xxdk::base::*;
-use xxdk::service::*;
+use xxdk::service::{CMixServerConfig, SenderId, Utf8Lossy};
 
 const SECRET: &str = "Hello";
 const REGISTRATION_CODE: &str = "";
@@ -56,17 +56,16 @@ pub async fn run() -> Result<(), String> {
         private_key: String::from(""),
     };
 
-    let xx_router = xxdk::service::Router::new(Arc::new(cmix)).route("demo", xx_rpc_handler);
-    CMixServer::serve(xx_router, cmix_config).await
+    let xx_router = xxdk::service::Router::with_state(Arc::new(cmix)).route("demo", xx_rpc_handler);
+    xxdk::service::serve(xx_router, cmix_config).await
 }
 
-pub async fn xx_rpc_handler(_: Arc<CMix>, request: IncomingRequest) -> Result<Vec<u8>, String> {
-    let sender: String = request.sender_id.iter().fold(String::new(), |mut s, b| {
+pub async fn xx_rpc_handler(id: SenderId, req: Utf8Lossy) -> String {
+    let sender: String = id.0.iter().fold(String::new(), |mut s, b| {
         write!(s, "{b:02x}").unwrap();
         s
     });
     tracing::info!(sender, "Received message via cMix",);
-    let text = String::from_utf8_lossy(&request.request);
-
-    Ok(format!("Hi from rust rpc example! Echoed message: {text}").into_bytes())
+    let text = req.0;
+    format!("Hi from rust rpc example! Echoed message: {text}")
 }

xxdk-rs/xxdk/Cargo.toml

@@ -10,6 +10,7 @@ base64 = "0.22.1"
 lazy_static = "1.4.0"
 libc = "0.2.153"
 serde = { version = "1.0.202", features = ["derive"] }
+serde_json = "1.0.120"
 tokio = { version = "1.37.0", features = ["rt", "fs", "sync", "time"] }
 tower = "0.4.13"
 tracing = "0.1.40"

xxdk-rs/xxdk/src/service.rs

-use base64::prelude::*;
-use serde::Deserialize;
+use std::borrow::Cow;
 use std::collections::HashMap;
-use std::future::{self, poll_fn, Future};
+use std::future::Future;
 use std::pin::Pin;
 use std::sync::Arc;
-use std::task::Poll;
+use std::task::{Context, Poll};
 use std::time::Duration;
-use tokio::sync::mpsc;
+
+use base64::prelude::*;
+use serde::de::DeserializeOwned;
+use serde::{Deserialize, Serialize};
+use serde_json as json;
 use tower::Service;
 
-use crate::base;
-use crate::rpc;
+use crate::{base, rpc};
+
+pub type PinnedFuture<T> = Pin<Box<dyn Future<Output = T> + Send + 'static>>;
 
 #[derive(Debug, Clone)]
 pub struct IncomingRequest {
-    pub sender_id: Vec<u8>,
-    pub request: Vec<u8>,
+    sender_id: Vec<u8>,
+    request: Vec<u8>,
+    separator_idx: usize,
 }
-#[derive(Debug, Clone)]
-struct Response {
-    pub text: Vec<u8>,
+
+impl IncomingRequest {
+    fn new(sender_id: Vec<u8>, request: Vec<u8>) -> Result<Self, String> {
+        let separator_idx = request
+            .iter()
+            .position(|b| *b == b',')
+            .ok_or_else(|| "no endpoint in request".to_string())?;
+
+        std::str::from_utf8(&request[..separator_idx])
+            .map_err(|e| format!("non-UTF-8 endpoint: {e}"))?;
+
+        Ok(Self {
+            sender_id,
+            request,
+            separator_idx,
+        })
+    }
+
+    pub fn sender_id(&self) -> &[u8] {
+        &self.sender_id
+    }
+
+    pub fn endpoint(&self) -> &str {
+        unsafe { std::str::from_utf8_unchecked(&self.request[..self.separator_idx]) }
     }
 
-type HandlerFnInner<T> = dyn Fn(
-        Arc<T>,
-        IncomingRequest,
-    ) -> Pin<Box<dyn Future<Output = Result<Vec<u8>, String>> + Send + 'static>>
-    + Send
-    + Sync
-    + 'static;
+    pub fn request(&self) -> &[u8] {
+        &self.request[self.separator_idx + 1..]
+    }
+}
 
-pub trait HandlerFn<T>:
-    Fn(Arc<T>, IncomingRequest) -> Self::Future + Send + Sync + 'static
+// TODO If we're a bit more careful about it, we can probably get rid of the Sync bound here
+pub trait Handler<T, S, Res>: Clone + Send + Sync + Sized + 'static {
+    fn call(self, req: IncomingRequest, state: S) -> PinnedFuture<Result<Vec<u8>, String>>;
+}
+
+macro_rules! impl_handler {
+    ($($ty:ident),*) => {
+        impl<F, Fut, S, Res, $($ty),*> Handler<($($ty,)*), S, Res> for F
+        where
+            F: FnOnce($($ty),*) -> Fut + Clone + Send + Sync + Sized + 'static,
+            Fut: Future<Output = Res> + Send + 'static,
+            S: Send + 'static,
+            Res: IntoResponse,
+            $(
+                $ty: FromRequest<S>,
+            )*
         {
-    type Future: Future<Output = Result<Vec<u8>, String>> + Send + 'static;
+            #[allow(non_snake_case, unused_variables)]
+            fn call(self, req: IncomingRequest, state: S) -> PinnedFuture<Result<Vec<u8>, String>> {
+                Box::pin(async move {
+                    $(
+                        let $ty = $ty::extract(&req, &state)?;
+                    )*
+                    self($($ty),*).await.into_response()
+                })
+            }
+        }
+    };
+}
+
+macro_rules! tuples {
+    ($name:ident) => {
+        $name!();
+        $name!(T1);
+        $name!(T1, T2);
+        $name!(T1, T2, T3);
+        $name!(T1, T2, T3, T4);
+        $name!(T1, T2, T3, T4, T5);
+        $name!(T1, T2, T3, T4, T5, T6);
+        $name!(T1, T2, T3, T4, T5, T6, T7);
+        $name!(T1, T2, T3, T4, T5, T6, T7, T8);
+        $name!(T1, T2, T3, T4, T5, T6, T7, T8, T9);
+        $name!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10);
+        $name!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11);
+        $name!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12);
+        $name!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13);
+        $name!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14);
+        $name!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15);
+        $name!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16);
+    };
+}
+
+tuples!(impl_handler);
+
+trait ErasedHandler<S>: Send + Sync + 'static {
+    fn call(&self, req: IncomingRequest, state: S) -> PinnedFuture<Result<Vec<u8>, String>>;
 }
 
-impl<T, F, Fut> HandlerFn<T> for F
+struct MakeErasedHandler<H, S> {
+    handler: H,
+    #[allow(clippy::type_complexity)]
+    call: fn(H, IncomingRequest, S) -> PinnedFuture<Result<Vec<u8>, String>>,
+}
+
+impl<H, S> ErasedHandler<S> for MakeErasedHandler<H, S>
 where
-    F: Fn(Arc<T>, IncomingRequest) -> Fut + Send + Sync + 'static,
-    Fut: Future<Output = Result<Vec<u8>, String>> + Send + 'static,
+    H: Clone + Send + Sync + 'static,
+    S: 'static,
 {
-    type Future = Fut;
+    fn call(&self, req: IncomingRequest, state: S) -> PinnedFuture<Result<Vec<u8>, String>> {
+        let h = self.handler.clone();
+        (self.call)(h, req, state)
+    }
 }
 
-pub struct Router<T> {
-    handlers: HashMap<String, Arc<HandlerFnInner<T>>>,
-    state: Arc<T>,
+impl<H, S> MakeErasedHandler<H, S> {
+    fn make<T, Res>(handler: H) -> Self
+    where
+        H: Handler<T, S, Res>,
+    {
+        let call = |h: H, req, state| h.call(req, state);
+        Self { handler, call }
+    }
 }
 
-// Manual implementation to avoid derive putting a Clone bound on T
-impl<T> Clone for Router<T> {
-    fn clone(&self) -> Self {
-        Self {
-            handlers: self.handlers.clone(),
-            state: self.state.clone(),
+type BoxedErasedHandler<S> = Arc<dyn ErasedHandler<S>>;
+
+#[derive(Clone)]
+pub struct Router<S> {
+    inner: Arc<RouterInner<S>>,
 }
+
+#[derive(Clone)]
+struct RouterInner<S> {
+    handlers: HashMap<String, BoxedErasedHandler<S>>,
+    state: S,
+}
+
+impl Router<()> {
+    pub fn without_state() -> Self {
+        Self::with_state(())
     }
 }
 
-impl<T> Router<T> {
-    pub fn new(state: Arc<T>) -> Self {
-        Self {
-            handlers: HashMap::new(),
-            state,
+impl<S> Router<S>
+where
+    S: Send + Clone + 'static,
+{
+    pub fn with_state(state: S) -> Self {
+        let handlers = HashMap::new();
+        let inner = Arc::new(RouterInner { handlers, state });
+        Self { inner }
     }
+
+    pub fn route<H, T, Res>(self, endpoint: &str, handler: H) -> Self
+    where
+        H: Handler<T, S, Res>,
+    {
+        let handler = Arc::new(MakeErasedHandler::make(handler));
+        self.with_inner(|inner| {
+            inner.handlers.insert(String::from(endpoint), handler);
+        })
     }
 
-    pub fn route<F>(mut self, endpoint: &str, handler: F) -> Self
+    fn with_inner<F>(self, f: F) -> Self
     where
-        F: HandlerFn<T>,
+        F: FnOnce(&mut RouterInner<S>),
     {
-        self.handlers.insert(
-            endpoint.to_string(),
-            Arc::new(move |state, req| Box::pin(handler(state, req))),
-        );
-        self
+        let mut inner = self.into_inner();
+        f(&mut inner);
+        Self {
+            inner: Arc::new(inner),
         }
     }
 
-impl<T> Service<IncomingRequest> for Router<T>
+    fn into_inner(self) -> RouterInner<S> {
+        match Arc::try_unwrap(self.inner) {
+            Ok(inner) => inner,
+            Err(arc) => RouterInner::clone(&*arc),
+        }
+    }
+}
+
+impl<S> Service<IncomingRequest> for Router<S>
 where
-    T: Send + Sync + 'static,
+    S: Clone + Send + 'static,
 {
     type Response = Vec<u8>;
 
     type Error = String;
 
-    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
+    type Future = PinnedFuture<Result<Vec<u8>, String>>;
 
-    fn poll_ready(&mut self, _cx: &mut std::task::Context<'_>) -> Poll<Result<(), Self::Error>> {
+    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         Poll::Ready(Ok(()))
     }
 
     fn call(&mut self, req: IncomingRequest) -> Self::Future {
-        // TODO All this manual matching on Results is gross, see if we can find a cleaner way to
-        // do this
-        let separator_idx = match req
-            .request
-            .iter()
-            .position(|b| *b == b',')
-            .ok_or_else(|| "no endpoint in request".to_string())
-        {
-            Ok(idx) => idx,
-            Err(e) => return Box::pin(future::ready(Err(e))),
+        let endpoint = req.endpoint();
+        let handler = match self.inner.handlers.get(req.endpoint()) {
+            Some(h) => h,
+            None => {
+                return Box::pin(std::future::ready(Err(format!(
+                    "unrecognized endpoint `{endpoint}`"
+                ))))
+            }
         };
-        let (endpoint, request) = req.request.split_at(separator_idx);
 
-        let endpoint =
-            match std::str::from_utf8(endpoint).map_err(|e| format!("non-UTF-8 endpoint ({e})")) {
-                Ok(endpoint) => endpoint,
-                Err(e) => return Box::pin(future::ready(Err(e))),
-            };
+        let state = self.inner.state.clone();
+        handler.call(req, state)
+    }
+}
 
-        if let Some(handler) = self.handlers.get(endpoint) {
-            let (_, request) = request.split_first().unwrap();
-            let req = IncomingRequest {
-                request: request.into(),
-                sender_id: req.sender_id,
-            };
+// TODO We can put a lifetime parameter on this to allow borrowing directly from the request buffer
+pub trait FromRequest<S>: Sized {
+    fn extract(req: &IncomingRequest, state: &S) -> Result<Self, String>;
+}
 
-            handler(self.state.clone(), req)
-        } else {
-            Box::pin(future::ready(Err(format!(
-                "unrecognized endpoint `{endpoint}`"
-            ))))
+pub trait IntoResponse {
+    fn into_response(self) -> Result<Vec<u8>, String>;
+}
+
+impl<R> IntoResponse for Result<R, String>
+where
+    R: IntoResponse,
+{
+    fn into_response(self) -> Result<Vec<u8>, String> {
+        self.and_then(|r| r.into_response())
+    }
+}
+
+impl IntoResponse for () {
+    fn into_response(self) -> Result<Vec<u8>, String> {
+        Ok(Vec::new())
+    }
+}
+
+impl IntoResponse for Vec<u8> {
+    fn into_response(self) -> Result<Vec<u8>, String> {
+        Ok(self)
+    }
 }
+
+impl IntoResponse for &[u8] {
+    fn into_response(self) -> Result<Vec<u8>, String> {
+        Ok(Vec::from(self))
+    }
+}
+
+impl<const N: usize> IntoResponse for [u8; N] {
+    fn into_response(self) -> Result<Vec<u8>, String> {
+        Ok(Vec::from(&self))
+    }
+}
+
+impl IntoResponse for Cow<'_, [u8]> {
+    fn into_response(self) -> Result<Vec<u8>, String> {
+        Ok(self.into_owned())
+    }
+}
+
+impl IntoResponse for String {
+    fn into_response(self) -> Result<Vec<u8>, String> {
+        Ok(self.into_bytes())
+    }
+}
+
+impl IntoResponse for &str {
+    fn into_response(self) -> Result<Vec<u8>, String> {
+        Ok(Vec::from(self.as_bytes()))
+    }
+}
+
+impl IntoResponse for Cow<'_, str> {
+    fn into_response(self) -> Result<Vec<u8>, String> {
+        Ok(self.into_owned().into_bytes())
+    }
+}
+
+#[derive(Debug, Clone, Copy)]
+pub struct Json<T>(pub T);
+
+impl<T, S> FromRequest<S> for Json<T>
+where
+    T: DeserializeOwned,
+{
+    fn extract(req: &IncomingRequest, _state: &S) -> Result<Self, String> {
+        Ok(Self(
+            json::from_slice(req.request()).map_err(|e| e.to_string())?,
+        ))
+    }
+}
+
+impl<T> IntoResponse for Json<T>
+where
+    T: Serialize,
+{
+    fn into_response(self) -> Result<Vec<u8>, String> {
+        json::to_vec(&self.0).map_err(|e| e.to_string())
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct SenderId(pub Vec<u8>);
+
+impl<S> FromRequest<S> for SenderId {
+    fn extract(req: &IncomingRequest, _state: &S) -> Result<Self, String> {
+        Ok(Self(req.sender_id.clone()))
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct RawRequest(pub Vec<u8>);
+
+impl<S> FromRequest<S> for RawRequest {
+    fn extract(req: &IncomingRequest, _state: &S) -> Result<Self, String> {
+        Ok(Self(Vec::from(req.request())))
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct Utf8(pub String);
+
+impl<S> FromRequest<S> for Utf8 {
+    fn extract(req: &IncomingRequest, _state: &S) -> Result<Self, String> {
+        Ok(Self(String::from(
+            std::str::from_utf8(req.request()).map_err(|e| e.to_string())?,
+        )))
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct Utf8Lossy(pub String);
+
+impl<S> FromRequest<S> for Utf8Lossy {
+    fn extract(req: &IncomingRequest, _state: &S) -> Result<Self, String> {
+        Ok(Self(String::from_utf8_lossy(req.request()).into_owned()))
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct State<S>(pub S);
+
+impl<S> FromRequest<S> for State<S>
+where
+    S: Clone,
+{
+    fn extract(_req: &IncomingRequest, state: &S) -> Result<Self, String> {
+        Ok(Self(state.clone()))
     }
 }
 
@@ -138,13 +374,9 @@ pub struct CMixServerConfig {
     pub private_key: String,
 }
 
-#[derive(Debug)]
-pub struct CMixServer;
-
-impl CMixServer {
-    pub async fn serve<T>(router: Router<T>, config: CMixServerConfig) -> Result<(), String>
+pub async fn serve<S>(service: S, config: CMixServerConfig) -> Result<(), String>
 where
-        T: Send + Sync + 'static,
+    S: Service<IncomingRequest, Response = Vec<u8>, Error = String> + Clone + Send + 'static,
 {
     tracing::info!("Starting cMix server");
     let ndf_contents = tokio::fs::read_to_string(&config.ndf_path)
@@ -231,16 +463,7 @@ impl CMixServer {
     cmix.ekv_set("rpc_server_private_key", &private_key)?;
 
     let runtime = tokio::runtime::Handle::current();
-        let (sender, mut response_queue) = mpsc::channel(256);
-        let cbs = CMixServerCallbacks {
-            router,
-            runtime,
-            response_queue: sender,
-        };
-
-        tracing::info!("Spawning RPC server");
-        base::callbacks::set_rpc_callbacks();
-        let rpc_server = rpc::new_server(&cmix, cbs, reception_id, private_key)?;
+    let cbs = CMixServerCallback { service, runtime };
 
     let cmix = Arc::new(cmix);
     tokio::task::spawn_blocking({
@@ -262,8 +485,10 @@ impl CMixServer {
         tokio::time::sleep(Duration::from_secs(1)).await;
     }
 
+    tracing::info!("Spawning RPC server");
+    base::callbacks::set_rpc_callbacks();
+    let rpc_server = rpc::new_server(&cmix, cbs, reception_id, private_key)?;
     rpc_server.start();
-
     tracing::info!(
         "RPC Server CB PTR: {:#x}",
         rpc_server.cb as *const _ as *const libc::c_void as usize
@@ -272,61 +497,50 @@ impl CMixServer {
     tracing::info!("RPC Public Key: {public_key_b64}");
     tracing::info!("RPC Reception ID: {reception_id_b64}");
 
-        while let Some(resp) = response_queue.recv().await {
-            tokio::spawn(async move {
-                tracing::debug!("request received, sending response");
-                tracing::debug!("{}", String::from_utf8_lossy(&resp.text));
-            });
-        }
+    // TODO We need a better way to shut down the server. This never actually completes or gets
+    // past this line, it just runs until the process gets a kill signal.
+    std::future::pending::<()>().await;
 
-        // rpc_server.stop();
+    rpc_server.stop();
     cmix.stop_network_follower()
 }
-}
-
-unsafe impl Send for CMixServer {}
 
-struct CMixServerCallbacks<T> {
-    router: Router<T>,
+struct CMixServerCallback<S> {
+    service: S,
     runtime: tokio::runtime::Handle,
-    response_queue: mpsc::Sender<Response>,
 }
 
-impl<T> rpc::ServerCallback for CMixServerCallbacks<T>
+impl<S> rpc::ServerCallback for CMixServerCallback<S>
 where
-    T: Send + Sync + 'static,
+    S: Service<IncomingRequest, Response = Vec<u8>, Error = String> + Clone + Send + 'static,
 {
-    #[tracing::instrument(level = "debug", skip(self))]
     fn serve_req(&self, sender_id: Vec<u8>, request: Vec<u8>) -> Vec<u8> {
-        let mut router = self.router.clone();
-        let response_queue = self.response_queue.clone();
-        let req = IncomingRequest { sender_id, request };
-        let r = self.runtime.block_on(async {
-            let ret: Vec<u8>;
-            tracing::debug!("Evaluating router on request");
-            if poll_fn(|cx| router.poll_ready(cx)).await.is_ok() {
-                match router.call(req).await {
-                    Err(e) => {
-                        tracing::warn!(error = e, "Error in servicing request");
-                        ret = e.into_bytes();
-                    }
-                    Ok(resp) => {
-                        if response_queue
-                            .send(Response { text: resp.clone() })
+        let mut service = self.service.clone();
+        let res: Result<Vec<u8>, String> = self.runtime.block_on(async move {
+            tracing::debug!("evaluating service on request");
+            if std::future::poll_fn(|cx| service.poll_ready(cx))
                 .await
-                            .is_err()
+                .is_ok()
             {
-                            tracing::warn!("couldn't send to queue");
-                        };
-                        ret = resp;
-                    }
-                }
+                let req = IncomingRequest::new(sender_id, request)?;
+                service.call(req).await
             } else {
-                ret = String::from("error unable to service request").into_bytes();
+                Err("unable to service request".to_string())
             }
-            ret
         });
-        tracing::warn!("sending response: {}", String::from_utf8_lossy(&r));
-        return r;
+
+        let res = match res {
+            Ok(bytes) => bytes,
+            Err(text) => {
+                tracing::warn!(error = text, "error servicing request");
+                text.into_bytes()
+            }
+        };
+
+        tracing::info!(
+            res = String::from_utf8_lossy(&res).as_ref(),
+            "sending response"
+        );
+        res
     }
 }