use std::io;
use derive_more::{Display, Error, From};
use http::{header, Method, StatusCode};
use crate::{
error::ResponseError,
header::HeaderValue,
message::RequestHead,
response::{Response, ResponseBuilder},
};
mod codec;
mod dispatcher;
mod frame;
mod mask;
mod proto;
pub use self::codec::{Codec, Frame, Item, Message};
pub use self::dispatcher::Dispatcher;
pub use self::frame::Parser;
pub use self::proto::{hash_key, CloseCode, CloseReason, OpCode};
#[derive(Debug, Display, From, Error)]
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 = "A payload reached size limit.")]
Overflow,
#[display(fmt = "Continuation is not started.")]
ContinuationNotStarted,
#[display(fmt = "Received new continuation but it is already started.")]
ContinuationStarted,
#[display(fmt = "Unknown continuation fragment: {}.", _0)]
ContinuationFragment(#[error(not(source))] OpCode),
#[display(fmt = "I/O error: {}", _0)]
Io(io::Error),
}
impl ResponseError for ProtocolError {}
#[derive(PartialEq, Debug, Display)]
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 ResponseError for HandshakeError {
fn error_response(&self) -> Response {
match self {
HandshakeError::GetMethodRequired => Response::MethodNotAllowed()
.insert_header((header::ALLOW, "GET"))
.finish(),
HandshakeError::NoWebsocketUpgrade => Response::BadRequest()
.reason("No WebSocket Upgrade header found")
.finish(),
HandshakeError::NoConnectionUpgrade => Response::BadRequest()
.reason("No Connection upgrade")
.finish(),
HandshakeError::NoVersionHeader => Response::BadRequest()
.reason("WebSocket version header is required")
.finish(),
HandshakeError::UnsupportedVersion => Response::BadRequest()
.reason("Unsupported WebSocket version")
.finish(),
HandshakeError::BadWebsocketKey => {
Response::BadRequest().reason("Handshake error").finish()
}
}
}
}
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::TRANSFER_ENCODING, "chunked"))
.insert_header((
header::SEC_WEBSOCKET_ACCEPT,
HeaderValue::from_bytes(&key).unwrap(),
))
.take()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::test::TestRequest;
use http::{header, Method};
#[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_wserror_http_response() {
let resp: Response = HandshakeError::GetMethodRequired.error_response();
assert_eq!(resp.status(), StatusCode::METHOD_NOT_ALLOWED);
let resp: Response = HandshakeError::NoWebsocketUpgrade.error_response();
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
let resp: Response = HandshakeError::NoConnectionUpgrade.error_response();
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
let resp: Response = HandshakeError::NoVersionHeader.error_response();
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
let resp: Response = HandshakeError::UnsupportedVersion.error_response();
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
let resp: Response = HandshakeError::BadWebsocketKey.error_response();
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
}
}