use super::route::{DirectPath, Route, RouteError, RouteResult};
use crate::client::WsSender;
use crate::talker::{Talker, TalkerCompatible, TermReason, WsIncoming};
use anyhow::Error;
use async_trait::async_trait;
use futures::channel::mpsc;
use headers::HeaderMapExt;
use hyper::upgrade::Upgraded;
use hyper::{Body, Request, Response, StatusCode};
use meio::prelude::{Action, ActionHandler, Actor, Address, LiteTask, StopReceiver};
use meio_protocol::Protocol;
use std::net::SocketAddr;
use tokio_tungstenite::WebSocketStream;
use tungstenite::protocol::Role;
pub trait WsFromRequest: Sized + Send + Sync + 'static {
type Output: Send;
type Protocol: Protocol;
fn from_request(&self, request: &Request<Body>) -> Option<Result<Self::Output, Error>>;
}
impl<T> WsFromRequest for T
where
T: DirectPath,
T::Parameter: Protocol,
{
type Output = <T as DirectPath>::Output;
type Protocol = T::Parameter;
fn from_request(&self, request: &Request<Body>) -> Option<Result<Self::Output, Error>> {
let uri = request.uri();
let path = uri.path();
if Self::paths().iter().any(|p| p == &path) {
let query = uri.query().unwrap_or("");
let output =
serde_qs::from_str(query).map_err(|err| RouteError::new(path, query, err).into());
Some(output)
} else {
None
}
}
}
pub struct WsReq<T: WsFromRequest> {
pub request: T::Output,
pub stream: WsHandler<T::Protocol>,
}
impl<T: WsFromRequest> Action for WsReq<T> {}
pub struct WsRoute<E, A>
where
A: Actor,
{
extractor: E,
address: Address<A>,
}
impl<E, A> WsRoute<E, A>
where
A: Actor,
{
pub fn new(extractor: E, address: Address<A>) -> Self {
Self { extractor, address }
}
}
impl<E, A> Route for WsRoute<E, A>
where
E: WsFromRequest,
A: Actor + ActionHandler<WsReq<E>>,
{
fn try_route(&self, addr: &SocketAddr, mut request: Request<Body>) -> RouteResult {
match self.extractor.from_request(&request) {
Some(Ok(value)) => {
let mut res = Response::new(Body::empty());
let address = self.address.clone();
if request.headers().typed_get::<headers::Upgrade>().is_none() {
*res.status_mut() = StatusCode::BAD_REQUEST;
}
let ws_key = request.headers().typed_get::<headers::SecWebsocketKey>();
let addr = *addr;
tokio::task::spawn(async move {
match hyper::upgrade::on(&mut request).await {
Ok(upgraded) => {
let websocket = tokio_tungstenite::WebSocketStream::from_raw_socket(
upgraded,
Role::Server,
None,
)
.await;
let stream = WsHandler::new(addr, websocket);
let msg = WsReq {
request: value,
stream,
};
address.act(msg)?;
Ok(())
}
Err(err) => {
log::error!("upgrade error: {}", err);
Err(Error::from(err))
}
}
});
*res.status_mut() = StatusCode::SWITCHING_PROTOCOLS;
res.headers_mut()
.typed_insert(headers::Connection::upgrade());
res.headers_mut()
.typed_insert(headers::Upgrade::websocket());
if let Some(value) = ws_key {
res.headers_mut()
.typed_insert(headers::SecWebsocketAccept::from(value));
}
let fut = futures::future::ready(Ok(res));
Ok(Box::pin(fut))
}
None => Err(request),
Some(Err(err)) => {
let fut = async move { Err(err) };
Ok(Box::pin(fut))
}
}
}
}
pub type WebSocket = WebSocketStream<Upgraded>;
struct WsInfo<P: Protocol> {
connection: WebSocket,
rx: mpsc::UnboundedReceiver<P::ToClient>,
}
pub struct WsHandler<P: Protocol> {
log_target: String,
addr: SocketAddr,
info: Option<WsInfo<P>>,
tx: mpsc::UnboundedSender<P::ToClient>,
}
impl<P: Protocol> WsHandler<P> {
fn new(addr: SocketAddr, websocket: WebSocket) -> Self {
let log_target = format!("WsHandler::{}", addr);
let (tx, rx) = mpsc::unbounded();
let info = WsInfo {
connection: websocket,
rx,
};
Self {
log_target,
addr,
info: Some(info),
tx,
}
}
pub fn addr(&self) -> SocketAddr {
self.addr
}
pub fn worker<A>(&mut self, address: Address<A>) -> WsProcessor<P, A>
where
A: Actor + ActionHandler<WsIncoming<P::ToServer>>,
{
let info = self.info.take().expect("already started");
let log_target = format!("{}::WsProcessor", self.log_target);
WsProcessor {
log_target,
info,
address,
}
}
pub fn send(&self, msg: P::ToClient) {
if let Err(err) = self.tx.unbounded_send(msg) {
log::error!(target: &self.log_target, "Can't send outgoing WS message: {}", err);
}
}
pub fn sender(&self) -> WsSender<P::ToClient> {
WsSender::new(&self.log_target, self.tx.clone())
}
}
pub struct WsProcessor<P: Protocol, A: Actor> {
log_target: String,
info: WsInfo<P>,
address: Address<A>,
}
impl<P, A> TalkerCompatible for WsProcessor<P, A>
where
P: Protocol,
A: Actor + ActionHandler<WsIncoming<P::ToServer>>,
{
type WebSocket = WebSocket;
type Message = tungstenite::Message;
type Error = tungstenite::Error;
type Actor = A;
type Codec = P::Codec;
type Incoming = P::ToServer;
type Outgoing = P::ToClient;
}
#[async_trait]
impl<P, A> LiteTask for WsProcessor<P, A>
where
P: Protocol,
A: Actor + ActionHandler<WsIncoming<P::ToServer>>,
{
type Output = TermReason;
fn log_target(&self) -> &str {
&self.log_target
}
async fn routine(self, stop: StopReceiver) -> Result<Self::Output, Error> {
let mut talker = Talker::<Self>::new(
&self.log_target,
self.address,
self.info.connection,
self.info.rx,
stop,
);
talker.routine().await
}
}