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, Addr, Handler, Message, Syn, Unsync, SpawnHandle}; use actix::fut::ActorFuture; use actix::dev::{ContextImpl, ToEnvelope, SyncEnvelope}; use body::{Body, Binary}; use error::{Error, 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)] #[inline] fn unsync_address(&mut self) -> Addr { self.inner.unsync_address() } #[doc(hidden)] #[inline] fn sync_address(&mut self) -> Addr { 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: Binary) { if !self.disconnected { if self.stream.is_none() { self.stream = Some(SmallVec::new()); } let stream = self.stream.as_mut().unwrap(); stream.push(ContextFrame::Chunk(Some(data))); } 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 #[inline] pub fn text>(&mut self, text: T) { self.write(Frame::message(text.into(), OpCode::Text, true, false)); } /// Send binary frame #[inline] pub fn binary>(&mut self, data: B) { self.write(Frame::message(data, OpCode::Binary, true, false)); } /// Send ping frame #[inline] pub fn ping(&mut self, message: &str) { self.write(Frame::message(Vec::from(message), OpCode::Ping, true, false)); } /// Send pong frame #[inline] pub fn pong(&mut self, message: &str) { self.write(Frame::message(Vec::from(message), OpCode::Pong, true, false)); } /// Send close frame #[inline] pub fn close(&mut self, code: CloseCode, reason: &str) { self.write(Frame::close(code, reason, false)); } /// 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 } #[inline] 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)); } /// Handle of the running future /// /// SpawnHandle is the handle returned by `AsyncContext::spawn()` method. pub fn handle(&self) -> SpawnHandle { self.inner.curr_handle() } } 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() && self.inner.poll(ctx).is_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> + Handler, M: Message + Send + 'static, M::Result: Send { fn pack(msg: M, tx: Option>) -> SyncEnvelope { SyncEnvelope::new(msg, tx) } } impl From> for Body where A: Actor>, S: 'static { fn from(ctx: WebsocketContext) -> Body { Body::Actor(Box::new(ctx)) } }