keyforkd: extract serialization logic into middleware

This commit is contained in:
Ryan Heywood 2023-09-07 08:05:38 -05:00
parent 1a13acdfe3
commit a7feed1bcc
Signed by: ryan
GPG Key ID: 8E401478A3FBEF72
3 changed files with 125 additions and 21 deletions

View File

@ -3,6 +3,7 @@ use std::{collections::HashMap, path::PathBuf};
use keyfork_mnemonic_util::Mnemonic; use keyfork_mnemonic_util::Mnemonic;
use tokio::io::{self, AsyncBufReadExt, BufReader}; use tokio::io::{self, AsyncBufReadExt, BufReader};
use tower::ServiceBuilder;
#[cfg(feature = "tracing")] #[cfg(feature = "tracing")]
use tracing::debug; use tracing::debug;
@ -17,6 +18,7 @@ use tracing_subscriber::{
mod error; mod error;
mod server; mod server;
mod service; mod service;
mod middleware;
use error::KeyforkdError; use error::KeyforkdError;
use server::UnixServer; use server::UnixServer;
use service::Keyforkd; use service::Keyforkd;
@ -52,7 +54,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
debug!("reading mnemonic from standard input"); debug!("reading mnemonic from standard input");
let mnemonic = load_mnemonic().await?; let mnemonic = load_mnemonic().await?;
let service = Keyforkd::new(mnemonic); let service = ServiceBuilder::new()
.layer(middleware::SerdeLayer::new())
.service(Keyforkd::new(mnemonic));
let runtime_vars = std::env::vars() let runtime_vars = std::env::vars()
.filter(|(key, _)| ["XDG_RUNTIME_DIR", "KEYFORKD_SOCKET_PATH"].contains(&key.as_str())) .filter(|(key, _)| ["XDG_RUNTIME_DIR", "KEYFORKD_SOCKET_PATH"].contains(&key.as_str()))

View File

@ -0,0 +1,91 @@
use std::{future::Future, marker::PhantomData, pin::Pin};
use bincode::{deserialize, serialize};
use serde::{de::DeserializeOwned, Serialize};
use thiserror::Error;
use tower::{Layer, Service};
pub struct SerdeLayer<'a, Request> {
phantom: PhantomData<&'a ()>,
phantom_request: PhantomData<&'a Request>,
}
impl<'a, Request> SerdeLayer<'a, Request> {
pub fn new() -> Self {
Self {
phantom: PhantomData,
phantom_request: PhantomData,
}
}
}
impl<'a, S: 'a, Request> Layer<S> for SerdeLayer<'a, Request> {
type Service = SerdeService<S, Request>;
fn layer(&self, service: S) -> Self::Service {
SerdeService {
service,
phantom_request: PhantomData,
}
}
}
#[derive(Clone)]
pub struct SerdeService<S, Request> {
service: S,
phantom_request: PhantomData<Request>,
}
#[derive(Debug, Error)]
pub enum SerdeServiceError {
#[error("Error while polling: {0}")]
Poll(String),
#[error("Error while calling: {0}")]
Call(String),
#[error("Error while converting: {0}")]
Convert(String),
}
impl<S, Request> Service<Vec<u8>> for SerdeService<S, Request>
where
S: Service<Request> + Send + Sync,
Request: DeserializeOwned + Send,
<S as Service<Request>>::Error: std::error::Error + Send,
<S as Service<Request>>::Response: Serialize + Send,
<S as Service<Request>>::Future: Send + 'static,
{
type Response = Vec<u8>;
type Error = SerdeServiceError;
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
fn poll_ready(
&mut self,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
self.service
.poll_ready(cx)
.map_err(|e| SerdeServiceError::Poll(e.to_string()))
}
fn call(&mut self, req: Vec<u8>) -> Self::Future {
let request: Request = match deserialize(&req) {
Ok(r) => r,
Err(e) => {
return Box::pin(
async move { Err(SerdeServiceError::Convert(e.to_string())) },
)
}
};
let response = self.service.call(request);
Box::pin(async move {
let response = response
.await
.map_err(|e| SerdeServiceError::Call(e.to_string()))?;
serialize(&response).map_err(|e| SerdeServiceError::Convert(e.to_string()))
})
}
}

View File

@ -1,28 +1,14 @@
use crate::service::{DerivationError, Keyforkd};
use keyfork_derive_util::DerivationPath;
use keyfork_frame::asyncext::{try_decode_from, try_encode_to}; use keyfork_frame::asyncext::{try_decode_from, try_encode_to};
use std::{ use std::{
io::Error, io::Error,
path::{Path, PathBuf}, path::{Path, PathBuf},
}; };
use tokio::net::{UnixListener, UnixStream}; use tokio::net::UnixListener;
use tower::{Service, ServiceExt}; use tower::{Service, ServiceExt};
#[cfg(feature = "tracing")] #[cfg(feature = "tracing")]
use tracing::debug; use tracing::debug;
async fn read_path_from_socket(
socket: &mut UnixStream,
) -> Result<DerivationPath, Box<dyn std::error::Error + Send>> {
let data = try_decode_from(socket).await.unwrap();
let path: DerivationPath = bincode::deserialize(&data[..]).unwrap();
Ok(path)
}
async fn wait_and_run(app: &mut Keyforkd, path: DerivationPath) -> Result<Vec<u8>, DerivationError> {
app.ready().await?.call(path).await
}
#[allow(clippy::module_name_repetitions)] #[allow(clippy::module_name_repetitions)]
pub struct UnixServer { pub struct UnixServer {
listener: UnixListener, listener: UnixListener,
@ -54,7 +40,14 @@ impl UnixServer {
}) })
} }
pub async fn run(&mut self, app: Keyforkd) -> Result<(), Box<dyn std::error::Error>> { pub async fn run<S, R>(&mut self, app: S) -> Result<(), Box<dyn std::error::Error>>
where
S: Service<R> + Clone + Send + 'static,
R: From<Vec<u8>> + Send,
<S as Service<R>>::Error: std::error::Error + Send,
<S as Service<R>>::Response: std::convert::Into<Vec<u8>> + Send,
<S as Service<R>>::Future: Send,
{
#[cfg(feature = "tracing")] #[cfg(feature = "tracing")]
debug!("Listening for clients"); debug!("Listening for clients");
loop { loop {
@ -63,8 +56,8 @@ impl UnixServer {
#[cfg(feature = "tracing")] #[cfg(feature = "tracing")]
debug!("new socket connected"); debug!("new socket connected");
tokio::spawn(async move { tokio::spawn(async move {
let path = match read_path_from_socket(&mut socket).await { let bytes = match try_decode_from(&mut socket).await {
Ok(path) => path, Ok(bytes) => bytes,
Err(e) => { Err(e) => {
#[cfg(feature = "tracing")] #[cfg(feature = "tracing")]
debug!(%e, "Error reading DerivationPath from socket"); debug!(%e, "Error reading DerivationPath from socket");
@ -78,7 +71,22 @@ impl UnixServer {
} }
}; };
let response = match wait_and_run(&mut app, path).await { let app = match app.ready().await {
Ok(app) => app,
Err(e) => {
#[cfg(feature = "tracing")]
debug!(%e, "Could not poll ready");
let content = e.to_string().bytes().collect::<Vec<_>>();
let result = try_encode_to(&content[..], &mut socket).await;
#[cfg(feature = "tracing")]
if let Err(error) = result {
debug!(%error, "Error sending error to client");
}
return;
}
};
let response = match app.call(bytes.into()).await {
Ok(response) => response, Ok(response) => response,
Err(e) => { Err(e) => {
#[cfg(feature = "tracing")] #[cfg(feature = "tracing")]
@ -91,7 +99,8 @@ impl UnixServer {
} }
return; return;
} }
}; }
.into();
if let Err(e) = try_encode_to(&response[..], &mut socket).await { if let Err(e) = try_encode_to(&response[..], &mut socket).await {
#[cfg(feature = "tracing")] #[cfg(feature = "tracing")]