use std::mem; use futures::{Async, Poll}; use futures::sync::oneshot::Sender; use futures::unsync::oneshot; use smallvec::SmallVec; use actix::{Actor, ActorState, ActorContext, AsyncContext, Address, SyncAddress, Handler, Subscriber, ResponseType, SpawnHandle}; use actix::fut::ActorFuture; use actix::dev::{queue, AsyncContextApi, ContextImpl, ContextProtocol, Envelope, ToEnvelope, RemoteEnvelope}; use body::{Body, Binary}; use error::{Error, Result, ErrorInternalServerError}; use httprequest::HttpRequest; use context::{Frame as ContextFrame, ActorHttpContext, Drain}; use ws::frame::Frame; use ws::proto::{OpCode, CloseCode}; /// Http actor execution context pub struct WebsocketContext where A: Actor>, { inner: ContextImpl, stream: Option>, request: HttpRequest, disconnected: bool, } impl ActorContext for WebsocketContext where A: Actor { fn stop(&mut self) { self.inner.stop(); } fn terminate(&mut self) { self.inner.terminate() } fn state(&self) -> ActorState { self.inner.state() } } impl AsyncContext for WebsocketContext where A: Actor { fn spawn(&mut self, fut: F) -> SpawnHandle where F: ActorFuture + 'static { self.inner.spawn(fut) } fn wait(&mut self, fut: F) where F: ActorFuture + 'static { self.inner.wait(fut) } #[doc(hidden)] #[inline] fn waiting(&self) -> bool { self.inner.waiting() || self.inner.state() == ActorState::Stopping || self.inner.state() == ActorState::Stopped } fn cancel_future(&mut self, handle: SpawnHandle) -> bool { self.inner.cancel_future(handle) } } #[doc(hidden)] impl AsyncContextApi for WebsocketContext where A: Actor { #[inline] fn unsync_sender(&mut self) -> queue::unsync::UnboundedSender> { self.inner.unsync_sender() } #[inline] fn unsync_address(&mut self) -> Address { self.inner.unsync_address() } #[inline] fn sync_address(&mut self) -> SyncAddress { self.inner.sync_address() } } impl WebsocketContext where A: Actor { #[inline] pub fn new(req: HttpRequest, actor: A) -> WebsocketContext { WebsocketContext::from_request(req).actor(actor) } pub fn from_request(req: HttpRequest) -> WebsocketContext { WebsocketContext { inner: ContextImpl::new(None), stream: None, request: req, disconnected: false, } } #[inline] pub fn actor(mut self, actor: A) -> WebsocketContext { self.inner.set_actor(actor); self } } impl WebsocketContext where A: Actor { /// Write payload #[inline] fn write>(&mut self, data: B) { if !self.disconnected { self.add_frame(ContextFrame::Chunk(Some(data.into()))); } else { warn!("Trying to write to disconnected response"); } } /// Shared application state #[inline] pub fn state(&self) -> &S { self.request.state() } /// Incoming request #[inline] pub fn request(&mut self) -> &mut HttpRequest { &mut self.request } /// Send text frame pub fn text(&mut self, text: &str) { let mut frame = Frame::message(Vec::from(text), OpCode::Text, true); let mut buf = Vec::new(); frame.format(&mut buf).unwrap(); self.write(buf); } /// Send binary frame pub fn binary>(&mut self, data: B) { let mut frame = Frame::message(data, OpCode::Binary, true); let mut buf = Vec::new(); frame.format(&mut buf).unwrap(); self.write(buf); } /// Send ping frame pub fn ping(&mut self, message: &str) { let mut frame = Frame::message(Vec::from(message), OpCode::Ping, true); let mut buf = Vec::new(); frame.format(&mut buf).unwrap(); self.write(buf); } /// Send pong frame pub fn pong(&mut self, message: &str) { let mut frame = Frame::message(Vec::from(message), OpCode::Pong, true); let mut buf = Vec::new(); frame.format(&mut buf).unwrap(); self.write(buf); } /// Send close frame pub fn close(&mut self, code: CloseCode, reason: &str) { let mut frame = Frame::close(code, reason); let mut buf = Vec::new(); frame.format(&mut buf).unwrap(); self.write(buf); } /// Returns drain future pub fn drain(&mut self) -> Drain { let (tx, rx) = oneshot::channel(); self.inner.modify(); self.add_frame(ContextFrame::Drain(tx)); Drain::new(rx) } /// Check if connection still open #[inline] pub fn connected(&self) -> bool { !self.disconnected } fn add_frame(&mut self, frame: ContextFrame) { if self.stream.is_none() { self.stream = Some(SmallVec::new()); } self.stream.as_mut().map(|s| s.push(frame)); } } impl WebsocketContext where A: Actor { #[inline] #[doc(hidden)] pub fn subscriber(&mut self) -> Box> where A: Handler, M: ResponseType + 'static { self.inner.subscriber() } #[inline] #[doc(hidden)] pub fn sync_subscriber(&mut self) -> Box + Send> where A: Handler, M: ResponseType + Send + 'static, M::Item: Send, M::Error: Send, { self.inner.sync_subscriber() } } impl ActorHttpContext for WebsocketContext where A: Actor, S: 'static { #[inline] fn disconnected(&mut self) { self.disconnected = true; self.stop(); } fn poll(&mut self) -> Poll>, Error> { let ctx: &mut WebsocketContext = unsafe { mem::transmute(self as &mut WebsocketContext) }; if self.inner.alive() { match self.inner.poll(ctx) { Ok(Async::NotReady) | Ok(Async::Ready(())) => (), Err(_) => return Err(ErrorInternalServerError("error").into()), } } // frames if let Some(data) = self.stream.take() { Ok(Async::Ready(Some(data))) } else if self.inner.alive() { Ok(Async::NotReady) } else { Ok(Async::Ready(None)) } } } impl ToEnvelope for WebsocketContext where A: Actor>, { #[inline] fn pack(msg: M, tx: Option>>, channel_on_drop: bool) -> Envelope where A: Handler, M: ResponseType + Send + 'static, M::Item: Send, M::Error: Send { RemoteEnvelope::new(msg, tx, channel_on_drop).into() } } impl From> for Body where A: Actor>, S: 'static { fn from(ctx: WebsocketContext) -> Body { Body::Actor(Box::new(ctx)) } }