1
0
mirror of https://github.com/actix/actix-extras.git synced 2024-11-25 08:22:43 +01:00
actix-extras/src/ws.rs

430 lines
14 KiB
Rust
Raw Normal View History

2017-10-08 07:41:02 +02:00
//! `WebSocket` support for Actix
//!
//! To setup a `WebSocket`, first do web socket handshake then on success convert `Payload`
//! into a `WsStream` stream and then use `WsWriter` to communicate with the peer.
//!
//! ## Example
//!
//! ```rust
//! extern crate actix;
2017-10-14 16:59:35 +02:00
//! extern crate actix_web;
2017-10-16 00:10:35 +02:00
//!
//! use actix::*;
2017-10-14 16:59:35 +02:00
//! use actix_web::*;
2017-10-08 07:41:02 +02:00
//!
//! // WebSocket Route
//! struct WsRoute;
//!
//! impl Actor for WsRoute {
//! type Context = HttpContext<Self>;
//! }
//!
//! impl Route for WsRoute {
//! type State = ();
//!
2017-10-22 18:13:29 +02:00
//! fn request(req: &mut HttpRequest,
//! payload: Payload, ctx: &mut HttpContext<Self>) -> RouteResult<Self>
2017-10-08 07:41:02 +02:00
//! {
2017-10-09 05:16:48 +02:00
//! // WebSocket handshake
2017-10-11 01:03:32 +02:00
//! match ws::handshake(&req) {
2017-10-09 05:16:48 +02:00
//! Ok(resp) => {
//! // Send handshake response to peer
2017-10-14 01:33:23 +02:00
//! ctx.start(resp);
2017-10-09 05:16:48 +02:00
//! // Map Payload into WsStream
//! ctx.add_stream(ws::WsStream::new(payload));
//! // Start ws messages processing
2017-10-19 08:43:50 +02:00
//! Reply::async(WsRoute)
2017-10-09 05:16:48 +02:00
//! },
//! Err(err) =>
2017-10-14 01:33:23 +02:00
//! Reply::reply(err)
2017-10-08 07:41:02 +02:00
//! }
//! }
//! }
//!
//! // Define Handler for ws::Message message
//! impl StreamHandler<ws::Message> for WsRoute {}
//!
//! impl Handler<ws::Message> for WsRoute {
//! fn handle(&mut self, msg: ws::Message, ctx: &mut HttpContext<Self>)
//! -> Response<Self, ws::Message>
//! {
//! match msg {
//! ws::Message::Ping(msg) => ws::WsWriter::pong(ctx, msg),
2017-10-21 06:08:38 +02:00
//! ws::Message::Text(text) => ws::WsWriter::text(ctx, &text),
2017-10-08 07:41:02 +02:00
//! ws::Message::Binary(bin) => ws::WsWriter::binary(ctx, bin),
//! _ => (),
//! }
//! Self::empty()
//! }
//! }
//!
//! fn main() {}
//! ```
2017-10-08 06:48:00 +02:00
use std::vec::Vec;
2017-10-10 08:07:32 +02:00
use http::{Method, StatusCode, header};
2017-10-24 08:39:01 +02:00
use bytes::BytesMut;
2017-10-08 07:41:02 +02:00
use futures::{Async, Poll, Stream};
2017-10-08 06:48:00 +02:00
2017-10-22 00:21:16 +02:00
use actix::{Actor, ResponseType};
2017-10-08 06:48:00 +02:00
2017-10-24 08:25:32 +02:00
use body::Body;
2017-10-08 06:48:00 +02:00
use context::HttpContext;
2017-10-09 05:16:48 +02:00
use route::Route;
use payload::Payload;
2017-10-08 06:48:00 +02:00
use httpcodes::{HTTPBadRequest, HTTPMethodNotAllowed};
use httprequest::HttpRequest;
2017-10-24 08:25:32 +02:00
use httpresponse::{ConnectionType, HttpResponse};
2017-10-08 06:48:00 +02:00
use wsframe;
2017-10-08 07:41:02 +02:00
use wsproto::*;
2017-10-08 06:48:00 +02:00
2017-10-08 07:41:02 +02:00
#[doc(hidden)]
2017-10-10 08:07:32 +02:00
const SEC_WEBSOCKET_ACCEPT: &'static str = "SEC-WEBSOCKET-ACCEPT";
#[doc(hidden)]
const SEC_WEBSOCKET_KEY: &'static str = "SEC-WEBSOCKET-KEY";
#[doc(hidden)]
const SEC_WEBSOCKET_VERSION: &'static str = "SEC-WEBSOCKET-VERSION";
// #[doc(hidden)]
// const SEC_WEBSOCKET_PROTOCOL: &'static str = "SEC-WEBSOCKET-PROTOCOL";
2017-10-08 06:48:00 +02:00
2017-10-08 08:59:57 +02:00
/// `WebSocket` Message
2017-10-08 06:48:00 +02:00
#[derive(Debug)]
pub enum Message {
Text(String),
Binary(Vec<u8>),
Ping(String),
Pong(String),
Close,
Closed,
Error
}
2017-10-22 00:21:16 +02:00
impl ResponseType for Message {
type Item = ();
type Error = ();
}
2017-10-08 07:41:02 +02:00
/// Prepare `WebSocket` handshake response.
2017-10-08 06:48:00 +02:00
///
2017-10-08 08:59:57 +02:00
/// This function returns handshake `HttpResponse`, ready to send to peer.
2017-10-08 07:41:02 +02:00
/// It does not perform any IO.
2017-10-08 06:48:00 +02:00
///
2017-10-08 07:41:02 +02:00
// /// `protocols` is a sequence of known protocols. On successful handshake,
// /// the returned response headers contain the first protocol in this list
// /// which the server also knows.
2017-10-11 01:03:32 +02:00
pub fn handshake(req: &HttpRequest) -> Result<HttpResponse, HttpResponse> {
2017-10-08 06:48:00 +02:00
// WebSocket accepts only GET
if *req.method() != Method::GET {
2017-10-29 22:50:26 +01:00
return Err(
HTTPMethodNotAllowed
.builder()
.header(header::ALLOW, "GET")
.finish()?)
2017-10-08 06:48:00 +02:00
}
// Check for "UPGRADE" to websocket header
2017-10-10 08:07:32 +02:00
let has_hdr = if let Some(hdr) = req.headers().get(header::UPGRADE) {
if let Ok(s) = hdr.to_str() {
s.to_lowercase().contains("websocket")
} else {
false
}
2017-10-08 06:48:00 +02:00
} else {
false
};
if !has_hdr {
2017-10-29 22:50:26 +01:00
return Err(HTTPBadRequest.with_reason("No WebSocket UPGRADE header found"))
2017-10-08 06:48:00 +02:00
}
// Upgrade connection
2017-10-14 09:11:12 +02:00
if !req.upgrade() {
2017-10-11 01:03:32 +02:00
return Err(HTTPBadRequest.with_reason("No CONNECTION upgrade"))
2017-10-08 06:48:00 +02:00
}
// check supported version
2017-10-10 08:07:32 +02:00
if !req.headers().contains_key(SEC_WEBSOCKET_VERSION) {
2017-10-11 01:03:32 +02:00
return Err(HTTPBadRequest.with_reason("No websocket version header is required"))
2017-10-08 06:48:00 +02:00
}
let supported_ver = {
2017-10-10 08:07:32 +02:00
if let Some(hdr) = req.headers().get(SEC_WEBSOCKET_VERSION) {
hdr == "13" || hdr == "8" || hdr == "7"
} else {
false
2017-10-08 06:48:00 +02:00
}
};
if !supported_ver {
2017-10-11 01:03:32 +02:00
return Err(HTTPBadRequest.with_reason("Unsupported version"))
2017-10-08 06:48:00 +02:00
}
// check client handshake for validity
2017-10-10 08:07:32 +02:00
if !req.headers().contains_key(SEC_WEBSOCKET_KEY) {
2017-10-11 01:03:32 +02:00
return Err(HTTPBadRequest.with_reason("Handshake error"));
2017-10-10 08:07:32 +02:00
}
let key = {
let key = req.headers().get(SEC_WEBSOCKET_KEY).unwrap();
hash_key(key.as_ref())
2017-10-08 06:48:00 +02:00
};
2017-10-11 01:03:32 +02:00
Ok(HttpResponse::builder(StatusCode::SWITCHING_PROTOCOLS)
.connection_type(ConnectionType::Upgrade)
.header(header::UPGRADE, "websocket")
.header(header::TRANSFER_ENCODING, "chunked")
.header(SEC_WEBSOCKET_ACCEPT, key.as_str())
.body(Body::Upgrade)?
2017-10-08 06:48:00 +02:00
)
}
2017-10-08 07:41:02 +02:00
/// Maps `Payload` stream into stream of `ws::Message` items
2017-10-08 06:48:00 +02:00
pub struct WsStream {
rx: Payload,
buf: BytesMut,
2017-10-14 01:33:23 +02:00
closed: bool,
error_sent: bool,
2017-10-08 06:48:00 +02:00
}
impl WsStream {
pub fn new(rx: Payload) -> WsStream {
2017-10-14 01:33:23 +02:00
WsStream { rx: rx, buf: BytesMut::new(), closed: false, error_sent: false }
2017-10-08 06:48:00 +02:00
}
}
impl Stream for WsStream {
type Item = Message;
type Error = ();
fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
let mut done = false;
2017-10-14 01:33:23 +02:00
if !self.closed {
loop {
match self.rx.readany() {
2017-10-27 08:14:33 +02:00
Ok(Async::Ready(Some(chunk))) => {
self.buf.extend(chunk.0)
2017-10-14 01:33:23 +02:00
}
2017-10-27 08:14:33 +02:00
Ok(Async::Ready(None)) => {
done = true;
2017-10-14 01:33:23 +02:00
self.closed = true;
2017-10-21 08:12:36 +02:00
break;
2017-10-14 01:33:23 +02:00
}
2017-10-27 08:14:33 +02:00
Ok(Async::NotReady) => break,
Err(_) => {
2017-10-21 08:12:36 +02:00
self.closed = true;
break;
2017-10-14 01:33:23 +02:00
}
2017-10-09 05:16:48 +02:00
}
2017-10-08 06:48:00 +02:00
}
2017-10-09 05:16:48 +02:00
}
2017-10-08 06:48:00 +02:00
2017-10-09 05:16:48 +02:00
loop {
2017-10-08 06:48:00 +02:00
match wsframe::Frame::parse(&mut self.buf) {
Ok(Some(frame)) => {
2017-10-11 01:03:32 +02:00
trace!("WsFrame {}", frame);
2017-10-08 07:41:02 +02:00
let (_finished, opcode, payload) = frame.unpack();
2017-10-08 06:48:00 +02:00
match opcode {
OpCode::Continue => continue,
OpCode::Bad =>
return Ok(Async::Ready(Some(Message::Error))),
2017-10-21 08:12:36 +02:00
OpCode::Close => {
self.closed = true;
self.error_sent = true;
return Ok(Async::Ready(Some(Message::Closed)))
},
2017-10-08 06:48:00 +02:00
OpCode::Ping =>
return Ok(Async::Ready(Some(
Message::Ping(String::from_utf8_lossy(&payload).into())))),
OpCode::Pong =>
return Ok(Async::Ready(Some(
Message::Pong(String::from_utf8_lossy(&payload).into())))),
OpCode::Binary =>
return Ok(Async::Ready(Some(Message::Binary(payload)))),
OpCode::Text => {
match String::from_utf8(payload) {
Ok(s) =>
return Ok(Async::Ready(Some(Message::Text(s)))),
2017-10-08 07:41:02 +02:00
Err(_) =>
2017-10-08 06:48:00 +02:00
return Ok(Async::Ready(Some(Message::Error))),
}
}
}
}
2017-10-14 01:33:23 +02:00
Ok(None) => {
if done {
return Ok(Async::Ready(None))
} else if self.closed {
if !self.error_sent {
self.error_sent = true;
return Ok(Async::Ready(Some(Message::Closed)))
} else {
return Ok(Async::Ready(None))
}
} else {
return Ok(Async::NotReady)
}
2017-10-08 06:48:00 +02:00
},
2017-10-14 01:33:23 +02:00
Err(_) => {
self.closed = true;
self.error_sent = true;
return Ok(Async::Ready(Some(Message::Error)));
}
2017-10-08 06:48:00 +02:00
}
}
}
}
/// `WebSocket` writer
pub struct WsWriter;
impl WsWriter {
2017-10-08 07:41:02 +02:00
/// Send text frame
2017-10-21 02:16:17 +02:00
pub fn text<A>(ctx: &mut HttpContext<A>, text: &str)
2017-10-08 06:48:00 +02:00
where A: Actor<Context=HttpContext<A>> + Route
{
2017-10-21 02:16:17 +02:00
let mut frame = wsframe::Frame::message(Vec::from(text), OpCode::Text, true);
2017-10-08 06:48:00 +02:00
let mut buf = Vec::new();
frame.format(&mut buf).unwrap();
2017-10-24 08:39:01 +02:00
ctx.write(buf);
2017-10-08 06:48:00 +02:00
}
2017-10-08 07:41:02 +02:00
/// Send binary frame
2017-10-08 06:48:00 +02:00
pub fn binary<A>(ctx: &mut HttpContext<A>, data: Vec<u8>)
where A: Actor<Context=HttpContext<A>> + Route
{
let mut frame = wsframe::Frame::message(data, OpCode::Binary, true);
let mut buf = Vec::new();
frame.format(&mut buf).unwrap();
2017-10-24 08:39:01 +02:00
ctx.write(buf);
2017-10-08 06:48:00 +02:00
}
2017-10-08 07:41:02 +02:00
/// Send ping frame
2017-10-08 06:48:00 +02:00
pub fn ping<A>(ctx: &mut HttpContext<A>, message: String)
where A: Actor<Context=HttpContext<A>> + Route
{
2017-10-09 02:47:41 +02:00
let mut frame = wsframe::Frame::message(
Vec::from(message.as_str()), OpCode::Ping, true);
2017-10-08 06:48:00 +02:00
let mut buf = Vec::new();
frame.format(&mut buf).unwrap();
2017-10-24 08:39:01 +02:00
ctx.write(buf);
2017-10-08 06:48:00 +02:00
}
2017-10-08 07:41:02 +02:00
/// Send pong frame
2017-10-08 06:48:00 +02:00
pub fn pong<A>(ctx: &mut HttpContext<A>, message: String)
where A: Actor<Context=HttpContext<A>> + Route
{
2017-10-09 02:47:41 +02:00
let mut frame = wsframe::Frame::message(
Vec::from(message.as_str()), OpCode::Pong, true);
2017-10-08 06:48:00 +02:00
let mut buf = Vec::new();
frame.format(&mut buf).unwrap();
2017-10-24 08:39:01 +02:00
ctx.write(buf);
2017-10-08 06:48:00 +02:00
}
}
2017-10-23 02:33:24 +02:00
#[cfg(test)]
mod tests {
use http::{Method, HeaderMap, StatusCode, Version, header};
2017-10-23 02:33:24 +02:00
use super::{HttpRequest, SEC_WEBSOCKET_VERSION, SEC_WEBSOCKET_KEY, handshake};
#[test]
fn test_handshake() {
let req = HttpRequest::new(Method::POST, "/".to_owned(),
Version::HTTP_11, HeaderMap::new(), String::new());
2017-10-23 02:33:24 +02:00
match handshake(&req) {
Err(err) => assert_eq!(err.status(), StatusCode::METHOD_NOT_ALLOWED),
_ => panic!("should not happen"),
}
let req = HttpRequest::new(Method::GET, "/".to_owned(),
Version::HTTP_11, HeaderMap::new(), String::new());
2017-10-23 02:33:24 +02:00
match handshake(&req) {
Err(err) => assert_eq!(err.status(), StatusCode::METHOD_NOT_ALLOWED),
_ => panic!("should not happen"),
}
let mut headers = HeaderMap::new();
headers.insert(header::UPGRADE,
header::HeaderValue::from_static("test"));
let req = HttpRequest::new(Method::GET, "/".to_owned(),
Version::HTTP_11, headers, String::new());
2017-10-23 02:33:24 +02:00
match handshake(&req) {
Err(err) => assert_eq!(err.status(), StatusCode::METHOD_NOT_ALLOWED),
_ => panic!("should not happen"),
}
let mut headers = HeaderMap::new();
headers.insert(header::UPGRADE,
header::HeaderValue::from_static("websocket"));
let req = HttpRequest::new(Method::GET, "/".to_owned(),
Version::HTTP_11, headers, String::new());
2017-10-23 02:33:24 +02:00
match handshake(&req) {
Err(err) => assert_eq!(err.status(), StatusCode::BAD_REQUEST),
_ => panic!("should not happen"),
}
let mut headers = HeaderMap::new();
headers.insert(header::UPGRADE,
header::HeaderValue::from_static("websocket"));
headers.insert(header::CONNECTION,
header::HeaderValue::from_static("upgrade"));
let req = HttpRequest::new(Method::GET, "/".to_owned(),
Version::HTTP_11, headers, String::new());
2017-10-23 02:33:24 +02:00
match handshake(&req) {
Err(err) => assert_eq!(err.status(), StatusCode::BAD_REQUEST),
_ => panic!("should not happen"),
}
let mut headers = HeaderMap::new();
headers.insert(header::UPGRADE,
header::HeaderValue::from_static("websocket"));
headers.insert(header::CONNECTION,
header::HeaderValue::from_static("upgrade"));
headers.insert(SEC_WEBSOCKET_VERSION,
header::HeaderValue::from_static("5"));
let req = HttpRequest::new(Method::GET, "/".to_owned(),
Version::HTTP_11, headers, String::new());
2017-10-23 02:33:24 +02:00
match handshake(&req) {
Err(err) => assert_eq!(err.status(), StatusCode::BAD_REQUEST),
_ => panic!("should not happen"),
}
let mut headers = HeaderMap::new();
headers.insert(header::UPGRADE,
header::HeaderValue::from_static("websocket"));
headers.insert(header::CONNECTION,
header::HeaderValue::from_static("upgrade"));
headers.insert(SEC_WEBSOCKET_VERSION,
header::HeaderValue::from_static("13"));
let req = HttpRequest::new(Method::GET, "/".to_owned(),
Version::HTTP_11, headers, String::new());
2017-10-23 02:33:24 +02:00
match handshake(&req) {
Err(err) => assert_eq!(err.status(), StatusCode::BAD_REQUEST),
_ => panic!("should not happen"),
}
let mut headers = HeaderMap::new();
headers.insert(header::UPGRADE,
header::HeaderValue::from_static("websocket"));
headers.insert(header::CONNECTION,
header::HeaderValue::from_static("upgrade"));
headers.insert(SEC_WEBSOCKET_VERSION,
header::HeaderValue::from_static("13"));
headers.insert(SEC_WEBSOCKET_KEY,
header::HeaderValue::from_static("13"));
let req = HttpRequest::new(Method::GET, "/".to_owned(),
Version::HTTP_11, headers, String::new());
2017-10-23 02:33:24 +02:00
match handshake(&req) {
Ok(resp) => {
assert_eq!(resp.status(), StatusCode::SWITCHING_PROTOCOLS)
},
_ => panic!("should not happen"),
}
}
}