mirror of https://github.com/chipsenkbeil/distant
Refactor code to have request/response format w/ client wrapper for easier processing
Chip Senkbeil
3 years ago
parent
f6fa3e606e
commit
f2cce4aa34
@ -1,75 +1,102 @@
|
||||
use crate::utils::Session;
|
||||
use codec::{DistantCodec, DistantCodecError};
|
||||
use derive_more::{Display, Error, From};
|
||||
use futures::SinkExt;
|
||||
use orion::{
|
||||
aead::{self, SecretKey},
|
||||
errors::UnknownCryptoError,
|
||||
mod transport;
|
||||
pub use transport::{Transport, TransportError};
|
||||
|
||||
use crate::{
|
||||
data::{Request, Response, ResponsePayload},
|
||||
utils::Session,
|
||||
};
|
||||
use log::*;
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
sync::{Arc, Mutex},
|
||||
};
|
||||
use serde::{de::DeserializeOwned, Serialize};
|
||||
use std::sync::Arc;
|
||||
use tokio::{io, net::TcpStream};
|
||||
use tokio_stream::StreamExt;
|
||||
use tokio_util::codec::Framed;
|
||||
use tokio::{
|
||||
io,
|
||||
sync::{oneshot, watch},
|
||||
};
|
||||
use tokio_stream::wrappers::WatchStream;
|
||||
|
||||
mod codec;
|
||||
type Callbacks = Arc<Mutex<HashMap<usize, oneshot::Sender<Response>>>>;
|
||||
|
||||
#[derive(Debug, Display, Error, From)]
|
||||
pub enum TransportError {
|
||||
CodecError(DistantCodecError),
|
||||
EncryptError(UnknownCryptoError),
|
||||
IoError(io::Error),
|
||||
SerializeError(serde_cbor::Error),
|
||||
}
|
||||
/// Represents a client that can make requests against a server
|
||||
pub struct Client {
|
||||
/// Underlying transport used by client
|
||||
transport: Arc<tokio::sync::Mutex<Transport>>,
|
||||
|
||||
/// Represents a transport of data across the network
|
||||
pub struct Transport {
|
||||
inner: Framed<TcpStream, DistantCodec>,
|
||||
key: Arc<SecretKey>,
|
||||
}
|
||||
/// Collection of callbacks to be invoked upon receiving a response to a request
|
||||
callbacks: Callbacks,
|
||||
|
||||
impl Transport {
|
||||
/// Wraps a `TcpStream` and associated credentials in a transport layer
|
||||
pub fn new(stream: TcpStream, key: Arc<SecretKey>) -> Self {
|
||||
Self {
|
||||
inner: Framed::new(stream, DistantCodec),
|
||||
key,
|
||||
}
|
||||
}
|
||||
/// Callback to trigger when a response is received without an origin or with an origin
|
||||
/// not found in the list of callbacks
|
||||
rx: watch::Receiver<Response>,
|
||||
}
|
||||
|
||||
impl Client {
|
||||
/// Establishes a connection using the provided session
|
||||
pub async fn connect(session: Session) -> io::Result<Self> {
|
||||
let stream = TcpStream::connect(session.to_socket_addr().await?).await?;
|
||||
Ok(Self::new(stream, Arc::new(session.key)))
|
||||
}
|
||||
let transport = Arc::new(tokio::sync::Mutex::new(Transport::connect(session).await?));
|
||||
let callbacks: Callbacks = Arc::new(Mutex::new(HashMap::new()));
|
||||
let (tx, rx) = watch::channel(Response::from(ResponsePayload::Error {
|
||||
description: String::from("Fake server response"),
|
||||
}));
|
||||
|
||||
// Start a task that continually checks for responses and triggers callbacks
|
||||
let transport_2 = Arc::clone(&transport);
|
||||
let callbacks_2 = Arc::clone(&callbacks);
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
match transport_2.lock().await.receive::<Response>().await {
|
||||
Ok(Some(res)) => {
|
||||
let maybe_callback = res
|
||||
.origin_id
|
||||
.as_ref()
|
||||
.and_then(|id| callbacks_2.lock().unwrap().remove(id));
|
||||
|
||||
// If there is an origin to this response, trigger the callback
|
||||
if let Some(tx) = maybe_callback {
|
||||
if let Err(res) = tx.send(res) {
|
||||
error!("Failed to trigger callback for response {}", res.id);
|
||||
}
|
||||
|
||||
/// Sends some data across the wire
|
||||
pub async fn send<T: Serialize>(&mut self, data: T) -> Result<(), TransportError> {
|
||||
// Serialize, encrypt, and then (TODO) sign
|
||||
// NOTE: Cannot used packed implementation for now due to issues with deserialization
|
||||
let data = serde_cbor::to_vec(&data)?;
|
||||
let data = aead::seal(&self.key, &data)?;
|
||||
// Otherwise, this goes into the junk draw of response handlers
|
||||
} else {
|
||||
if let Err(x) = tx.send(res) {
|
||||
error!("Failed to trigger watch: {}", x);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(None) => break,
|
||||
Err(x) => {
|
||||
error!("{}", x);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
self.inner
|
||||
.send(&data)
|
||||
.await
|
||||
.map_err(TransportError::CodecError)
|
||||
Ok(Self {
|
||||
transport,
|
||||
callbacks,
|
||||
rx,
|
||||
})
|
||||
}
|
||||
|
||||
/// Receives some data from out on the wire, waiting until it's available,
|
||||
/// returning none if the transport is now closed
|
||||
pub async fn receive<T: DeserializeOwned>(&mut self) -> Result<Option<T>, TransportError> {
|
||||
// If data is received, we process like usual
|
||||
if let Some(data) = self.inner.next().await {
|
||||
// Validate (TODO) signature, decrypt, and then deserialize
|
||||
let data = data?;
|
||||
let data = aead::open(&self.key, &data)?;
|
||||
let data = serde_cbor::from_slice(&data)?;
|
||||
Ok(Some(data))
|
||||
/// Sends a request and waits for a response
|
||||
pub async fn send(&self, req: Request) -> Result<Response, TransportError> {
|
||||
// First, add a callback that will trigger when we get the response for this request
|
||||
let (tx, rx) = oneshot::channel();
|
||||
self.callbacks.lock().unwrap().insert(req.id, tx);
|
||||
|
||||
// Second, send the request
|
||||
self.transport.lock().await.send(req).await?;
|
||||
|
||||
// Third, wait for the response
|
||||
rx.await
|
||||
.map_err(|x| TransportError::from(io::Error::new(io::ErrorKind::ConnectionAborted, x)))
|
||||
}
|
||||
|
||||
// Otherwise, if no data is received, this means that our socket has closed
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
/// Creates and returns a new stream of responses that are received with no originating request
|
||||
pub fn to_response_stream(&self) -> WatchStream<Response> {
|
||||
WatchStream::new(self.rx.clone())
|
||||
}
|
||||
}
|
||||
|
||||
@ -0,0 +1,75 @@
|
||||
use crate::utils::Session;
|
||||
use codec::{DistantCodec, DistantCodecError};
|
||||
use derive_more::{Display, Error, From};
|
||||
use futures::SinkExt;
|
||||
use orion::{
|
||||
aead::{self, SecretKey},
|
||||
errors::UnknownCryptoError,
|
||||
};
|
||||
use serde::{de::DeserializeOwned, Serialize};
|
||||
use std::sync::Arc;
|
||||
use tokio::{io, net::TcpStream};
|
||||
use tokio_stream::StreamExt;
|
||||
use tokio_util::codec::Framed;
|
||||
|
||||
mod codec;
|
||||
|
||||
#[derive(Debug, Display, Error, From)]
|
||||
pub enum TransportError {
|
||||
CodecError(DistantCodecError),
|
||||
EncryptError(UnknownCryptoError),
|
||||
IoError(io::Error),
|
||||
SerializeError(serde_cbor::Error),
|
||||
}
|
||||
|
||||
/// Represents a transport of data across the network
|
||||
pub struct Transport {
|
||||
inner: Framed<TcpStream, DistantCodec>,
|
||||
key: Arc<SecretKey>,
|
||||
}
|
||||
|
||||
impl Transport {
|
||||
/// Wraps a `TcpStream` and associated credentials in a transport layer
|
||||
pub fn new(stream: TcpStream, key: Arc<SecretKey>) -> Self {
|
||||
Self {
|
||||
inner: Framed::new(stream, DistantCodec),
|
||||
key,
|
||||
}
|
||||
}
|
||||
|
||||
/// Establishes a connection using the provided session
|
||||
pub async fn connect(session: Session) -> io::Result<Self> {
|
||||
let stream = TcpStream::connect(session.to_socket_addr().await?).await?;
|
||||
Ok(Self::new(stream, Arc::new(session.key)))
|
||||
}
|
||||
|
||||
/// Sends some data across the wire
|
||||
pub async fn send<T: Serialize>(&mut self, data: T) -> Result<(), TransportError> {
|
||||
// Serialize, encrypt, and then (TODO) sign
|
||||
// NOTE: Cannot used packed implementation for now due to issues with deserialization
|
||||
let data = serde_cbor::to_vec(&data)?;
|
||||
let data = aead::seal(&self.key, &data)?;
|
||||
|
||||
self.inner
|
||||
.send(&data)
|
||||
.await
|
||||
.map_err(TransportError::CodecError)
|
||||
}
|
||||
|
||||
/// Receives some data from out on the wire, waiting until it's available,
|
||||
/// returning none if the transport is now closed
|
||||
pub async fn receive<T: DeserializeOwned>(&mut self) -> Result<Option<T>, TransportError> {
|
||||
// If data is received, we process like usual
|
||||
if let Some(data) = self.inner.next().await {
|
||||
// Validate (TODO) signature, decrypt, and then deserialize
|
||||
let data = data?;
|
||||
let data = aead::open(&self.key, &data)?;
|
||||
let data = serde_cbor::from_slice(&data)?;
|
||||
Ok(Some(data))
|
||||
|
||||
// Otherwise, if no data is received, this means that our socket has closed
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue