use std::io;
use derive_more::{Display, Error, From};
use http::{header, Method, StatusCode};
use crate::{body::BoxBody, header::HeaderValue, RequestHead, Response, ResponseBuilder};
mod codec;
mod dispatcher;
mod frame;
mod mask;
mod proto;
pub use self::{
codec::{Codec, Frame, Item, Message},
dispatcher::Dispatcher,
frame::Parser,
proto::{hash_key, CloseCode, CloseReason, OpCode},
};
#[derive(Debug, Display, Error, From)]
pub enum ProtocolError {
#[display(fmt = "received an unmasked frame from client")]
UnmaskedFrame,
#[display(fmt = "received a masked frame from server")]
MaskedFrame,
#[display(fmt = "invalid opcode ({})", _0)]
InvalidOpcode(#[error(not(source))] u8),
#[display(fmt = "invalid control frame length ({})", _0)]
InvalidLength(#[error(not(source))] usize),
#[display(fmt = "bad opcode")]
BadOpCode,
#[display(fmt = "payload reached size limit")]
Overflow,
#[display(fmt = "continuation has not started")]
ContinuationNotStarted,
#[display(fmt = "received new continuation but it has already started")]
ContinuationStarted,
#[display(fmt = "unknown continuation fragment: {}", _0)]
ContinuationFragment(#[error(not(source))] OpCode),
#[display(fmt = "I/O error: {}", _0)]
Io(io::Error),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Display, Error)]
pub enum HandshakeError {
#[display(fmt = "method not allowed")]
GetMethodRequired,
#[display(fmt = "WebSocket upgrade is expected")]
NoWebsocketUpgrade,
#[display(fmt = "connection upgrade is expected")]
NoConnectionUpgrade,
#[display(fmt = "WebSocket version header is required")]
NoVersionHeader,
#[display(fmt = "unsupported WebSocket version")]
UnsupportedVersion,
#[display(fmt = "unknown WebSocket key")]
BadWebsocketKey,
}
impl From<HandshakeError> for Response<BoxBody> {
fn from(err: HandshakeError) -> Self {
match err {
HandshakeError::GetMethodRequired => {
let mut res = Response::new(StatusCode::METHOD_NOT_ALLOWED);
#[allow(clippy::declare_interior_mutable_const)]
const HV_GET: HeaderValue = HeaderValue::from_static("GET");
res.headers_mut().insert(header::ALLOW, HV_GET);
res
}
HandshakeError::NoWebsocketUpgrade => {
let mut res = Response::bad_request();
res.head_mut().reason = Some("No WebSocket Upgrade header found");
res
}
HandshakeError::NoConnectionUpgrade => {
let mut res = Response::bad_request();
res.head_mut().reason = Some("No Connection upgrade");
res
}
HandshakeError::NoVersionHeader => {
let mut res = Response::bad_request();
res.head_mut().reason = Some("WebSocket version header is required");
res
}
HandshakeError::UnsupportedVersion => {
let mut res = Response::bad_request();
res.head_mut().reason = Some("Unsupported WebSocket version");
res
}
HandshakeError::BadWebsocketKey => {
let mut res = Response::bad_request();
res.head_mut().reason = Some("Handshake error");
res
}
}
}
}
impl From<&HandshakeError> for Response<BoxBody> {
fn from(err: &HandshakeError) -> Self {
(*err).into()
}
}
pub fn handshake(req: &RequestHead) -> Result<ResponseBuilder, HandshakeError> {
verify_handshake(req)?;
Ok(handshake_response(req))
}
pub fn verify_handshake(req: &RequestHead) -> Result<(), HandshakeError> {
if req.method != Method::GET {
return Err(HandshakeError::GetMethodRequired);
}
let has_hdr = if let Some(hdr) = req.headers().get(header::UPGRADE) {
if let Ok(s) = hdr.to_str() {
s.to_ascii_lowercase().contains("websocket")
} else {
false
}
} else {
false
};
if !has_hdr {
return Err(HandshakeError::NoWebsocketUpgrade);
}
if !req.upgrade() {
return Err(HandshakeError::NoConnectionUpgrade);
}
if !req.headers().contains_key(header::SEC_WEBSOCKET_VERSION) {
return Err(HandshakeError::NoVersionHeader);
}
let supported_ver = {
if let Some(hdr) = req.headers().get(header::SEC_WEBSOCKET_VERSION) {
hdr == "13" || hdr == "8" || hdr == "7"
} else {
false
}
};
if !supported_ver {
return Err(HandshakeError::UnsupportedVersion);
}
if !req.headers().contains_key(header::SEC_WEBSOCKET_KEY) {
return Err(HandshakeError::BadWebsocketKey);
}
Ok(())
}
pub fn handshake_response(req: &RequestHead) -> ResponseBuilder {
let key = {
let key = req.headers().get(header::SEC_WEBSOCKET_KEY).unwrap();
proto::hash_key(key.as_ref())
};
Response::build(StatusCode::SWITCHING_PROTOCOLS)
.upgrade("websocket")
.insert_header((
header::SEC_WEBSOCKET_ACCEPT,
HeaderValue::from_bytes(&key).unwrap(),
))
.take()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{header, test::TestRequest};
#[test]
fn test_handshake() {
let req = TestRequest::default().method(Method::POST).finish();
assert_eq!(
HandshakeError::GetMethodRequired,
verify_handshake(req.head()).unwrap_err(),
);
let req = TestRequest::default().finish();
assert_eq!(
HandshakeError::NoWebsocketUpgrade,
verify_handshake(req.head()).unwrap_err(),
);
let req = TestRequest::default()
.insert_header((header::UPGRADE, header::HeaderValue::from_static("test")))
.finish();
assert_eq!(
HandshakeError::NoWebsocketUpgrade,
verify_handshake(req.head()).unwrap_err(),
);
let req = TestRequest::default()
.insert_header((
header::UPGRADE,
header::HeaderValue::from_static("websocket"),
))
.finish();
assert_eq!(
HandshakeError::NoConnectionUpgrade,
verify_handshake(req.head()).unwrap_err(),
);
let req = TestRequest::default()
.insert_header((
header::UPGRADE,
header::HeaderValue::from_static("websocket"),
))
.insert_header((
header::CONNECTION,
header::HeaderValue::from_static("upgrade"),
))
.finish();
assert_eq!(
HandshakeError::NoVersionHeader,
verify_handshake(req.head()).unwrap_err(),
);
let req = TestRequest::default()
.insert_header((
header::UPGRADE,
header::HeaderValue::from_static("websocket"),
))
.insert_header((
header::CONNECTION,
header::HeaderValue::from_static("upgrade"),
))
.insert_header((
header::SEC_WEBSOCKET_VERSION,
header::HeaderValue::from_static("5"),
))
.finish();
assert_eq!(
HandshakeError::UnsupportedVersion,
verify_handshake(req.head()).unwrap_err(),
);
let req = TestRequest::default()
.insert_header((
header::UPGRADE,
header::HeaderValue::from_static("websocket"),
))
.insert_header((
header::CONNECTION,
header::HeaderValue::from_static("upgrade"),
))
.insert_header((
header::SEC_WEBSOCKET_VERSION,
header::HeaderValue::from_static("13"),
))
.finish();
assert_eq!(
HandshakeError::BadWebsocketKey,
verify_handshake(req.head()).unwrap_err(),
);
let req = TestRequest::default()
.insert_header((
header::UPGRADE,
header::HeaderValue::from_static("websocket"),
))
.insert_header((
header::CONNECTION,
header::HeaderValue::from_static("upgrade"),
))
.insert_header((
header::SEC_WEBSOCKET_VERSION,
header::HeaderValue::from_static("13"),
))
.insert_header((
header::SEC_WEBSOCKET_KEY,
header::HeaderValue::from_static("13"),
))
.finish();
assert_eq!(
StatusCode::SWITCHING_PROTOCOLS,
handshake_response(req.head()).finish().status()
);
}
#[test]
fn test_ws_error_http_response() {
let resp: Response<BoxBody> = HandshakeError::GetMethodRequired.into();
assert_eq!(resp.status(), StatusCode::METHOD_NOT_ALLOWED);
let resp: Response<BoxBody> = HandshakeError::NoWebsocketUpgrade.into();
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
let resp: Response<BoxBody> = HandshakeError::NoConnectionUpgrade.into();
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
let resp: Response<BoxBody> = HandshakeError::NoVersionHeader.into();
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
let resp: Response<BoxBody> = HandshakeError::UnsupportedVersion.into();
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
let resp: Response<BoxBody> = HandshakeError::BadWebsocketKey.into();
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
}
}