use std::borrow::Cow;
use std::fmt;
use std::pin::Pin;
use std::future::Future;
use std::task::{Context, Poll};
use futures::{future, ready, FutureExt, Sink, Stream, TryFutureExt};
use headers::{Connection, HeaderMapExt, SecWebsocketAccept, SecWebsocketKey, Upgrade};
use hyper::upgrade::OnUpgrade;
use salvo_core::http::header::{SEC_WEBSOCKET_VERSION, UPGRADE};
use salvo_core::http::{HttpError, StatusCode};
use salvo_core::{Error, Request, Response};
use tokio_tungstenite::{
tungstenite::protocol::{self, WebSocketConfig},
WebSocketStream,
};
#[allow(missing_debug_implementations)]
pub struct WsHandler {
config: Option<WebSocketConfig>,
}
impl WsHandler {
pub fn new() -> Self {
WsHandler { config: None }
}
pub fn with_config(config: WebSocketConfig) -> Self {
WsHandler {
config: Some(config),
}
}
pub fn max_send_queue(mut self, max: usize) -> Self {
self.config.get_or_insert_with(WebSocketConfig::default).max_send_queue = Some(max);
self
}
pub fn max_message_size(mut self, max: usize) -> Self {
self.config.get_or_insert_with(WebSocketConfig::default).max_message_size = Some(max);
self
}
pub fn max_frame_size(mut self, max: usize) -> Self {
self.config.get_or_insert_with(|| WebSocketConfig::default()).max_frame_size = Some(max);
self
}
pub fn handle(&self, req: &mut Request, res: &mut Response) -> Result<impl Future<Output=Option<WebSocket>>, HttpError> {
let req_headers = req.headers();
let matched = req_headers.typed_get::<Connection>().map(|conn| conn.contains(UPGRADE)).unwrap_or(false);
if !matched {
tracing::debug!("missing connection upgrade");
return Err(HttpError {
code: StatusCode::BAD_REQUEST,
name: "Bad Request".into(),
summary: "missing connection upgrade".into(),
detail: "".into(),
});
}
let matched = req_headers
.get(UPGRADE)
.and_then(|v| v.to_str().ok())
.map(|v| v.to_lowercase() == "websocket")
.unwrap_or(false);
if !matched {
tracing::debug!("missing upgrade header or it is not equal websocket");
return Err(HttpError {
code: StatusCode::BAD_REQUEST,
name: "Bad Request".into(),
summary: "missing upgrade header or it is not equal websocket".into(),
detail: "".into(),
});
}
let matched = !req_headers
.get(SEC_WEBSOCKET_VERSION)
.and_then(|v| v.to_str().ok())
.map(|v| v == "13")
.unwrap_or(false);
if matched {
tracing::debug!("websocket version is not equal 13");
return Err(HttpError {
code: StatusCode::BAD_REQUEST,
name: "Bad Request".into(),
summary: "websocket version is not equal 13".into(),
detail: "".into(),
});
}
let sec_ws_key = if let Some(key) = req_headers.typed_get::<SecWebsocketKey>() {
key
} else {
tracing::debug!("sec_websocket_key is not exist in request headers");
return Err(HttpError {
code: StatusCode::BAD_REQUEST,
name: "Bad Request".into(),
summary: "sec_websocket_key is not exist in request headers".into(),
detail: "".into(),
});
};
res.set_status_code(StatusCode::SWITCHING_PROTOCOLS);
res.headers_mut().typed_insert(Connection::upgrade());
res.headers_mut().typed_insert(Upgrade::websocket());
res.headers_mut().typed_insert(SecWebsocketAccept::from(sec_ws_key));
if let Some(on_upgrade) = req.extensions_mut().remove::<OnUpgrade>() {
let config = self.config.clone();
let fut = async move {
let ws = on_upgrade
.and_then(move |upgraded| {
tracing::debug!("websocket upgrade complete");
WebSocket::from_raw_socket(upgraded, protocol::Role::Server, config).map(Ok)
})
.await.ok();
ws
};
Ok(fut)
} else {
tracing::debug!("ws couldn't be upgraded since no upgrade state was present");
Err(HttpError {
code: StatusCode::BAD_REQUEST,
name: "Bad Request".into(),
summary: "ws couldn't be upgraded since no upgrade state was present".into(),
detail: "".into(),
})
}
}
}
pub struct WebSocket {
inner: WebSocketStream<hyper::upgrade::Upgraded>,
}
impl WebSocket {
pub(crate) async fn from_raw_socket(upgraded: hyper::upgrade::Upgraded, role: protocol::Role, config: Option<protocol::WebSocketConfig>) -> Self {
WebSocketStream::from_raw_socket(upgraded, role, config)
.map(|inner| WebSocket { inner })
.await
}
pub async fn close(mut self) -> Result<(), Error> {
future::poll_fn(|cx| Pin::new(&mut self).poll_close(cx)).await
}
}
impl Stream for WebSocket {
type Item = Result<Message, Error>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
match ready!(Pin::new(&mut self.inner).poll_next(cx)) {
Some(Ok(item)) => Poll::Ready(Some(Ok(Message { inner: item }))),
Some(Err(e)) => {
tracing::debug!("websocket poll error: {}", e);
Poll::Ready(Some(Err(Error::new(e))))
}
None => {
tracing::trace!("websocket closed");
Poll::Ready(None)
}
}
}
}
impl Sink<Message> for WebSocket {
type Error = Error;
fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
match ready!(Pin::new(&mut self.inner).poll_ready(cx)) {
Ok(()) => Poll::Ready(Ok(())),
Err(e) => Poll::Ready(Err(Error::new(e))),
}
}
fn start_send(mut self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
match Pin::new(&mut self.inner).start_send(item.inner) {
Ok(()) => Ok(()),
Err(e) => {
tracing::debug!("websocket start_send error: {}", e);
Err(Error::new(e))
}
}
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
match ready!(Pin::new(&mut self.inner).poll_flush(cx)) {
Ok(()) => Poll::Ready(Ok(())),
Err(e) => Poll::Ready(Err(Error::new(e))),
}
}
fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
match ready!(Pin::new(&mut self.inner).poll_close(cx)) {
Ok(()) => Poll::Ready(Ok(())),
Err(err) => {
tracing::debug!("websocket close error: {}", err);
Poll::Ready(Err(Error::new(err)))
}
}
}
}
impl fmt::Debug for WebSocket {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_struct("WebSocket").finish()
}
}
#[derive(Eq, PartialEq, Clone)]
pub struct Message {
inner: protocol::Message,
}
impl Message {
pub fn text<S: Into<String>>(s: S) -> Message {
Message {
inner: protocol::Message::text(s),
}
}
pub fn binary<V: Into<Vec<u8>>>(v: V) -> Message {
Message {
inner: protocol::Message::binary(v),
}
}
pub fn ping<V: Into<Vec<u8>>>(v: V) -> Message {
Message {
inner: protocol::Message::Ping(v.into()),
}
}
pub fn close() -> Message {
Message {
inner: protocol::Message::Close(None),
}
}
pub fn close_with(code: impl Into<u16>, reason: impl Into<Cow<'static, str>>) -> Message {
Message {
inner: protocol::Message::Close(Some(protocol::frame::CloseFrame {
code: protocol::frame::coding::CloseCode::from(code.into()),
reason: reason.into(),
})),
}
}
pub fn is_text(&self) -> bool {
self.inner.is_text()
}
pub fn is_binary(&self) -> bool {
self.inner.is_binary()
}
pub fn is_close(&self) -> bool {
self.inner.is_close()
}
pub fn is_ping(&self) -> bool {
self.inner.is_ping()
}
pub fn is_pong(&self) -> bool {
self.inner.is_pong()
}
pub fn close_frame(&self) -> Option<(u16, &str)> {
if let protocol::Message::Close(Some(ref close_frame)) = self.inner {
Some((close_frame.code.into(), close_frame.reason.as_ref()))
} else {
None
}
}
pub fn to_str(&self) -> Result<&str, ()> {
match self.inner {
protocol::Message::Text(ref s) => Ok(s),
_ => Err(()),
}
}
pub fn as_bytes(&self) -> &[u8] {
match self.inner {
protocol::Message::Text(ref s) => s.as_bytes(),
protocol::Message::Binary(ref v) => v,
protocol::Message::Ping(ref v) => v,
protocol::Message::Pong(ref v) => v,
protocol::Message::Close(_) => &[],
}
}
pub fn into_bytes(self) -> Vec<u8> {
self.inner.into_data()
}
}
impl fmt::Debug for Message {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
fmt::Debug::fmt(&self.inner, f)
}
}
impl Into<Vec<u8>> for Message {
fn into(self) -> Vec<u8> {
self.into_bytes()
}
}
#[derive(Debug)]
pub struct MissingConnectionUpgrade;
impl ::std::fmt::Display for MissingConnectionUpgrade {
fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
write!(f, "Connection header did not include 'upgrade'")
}
}
impl ::std::error::Error for MissingConnectionUpgrade {}