server/src/routes/websocket.rs

   1 use super::*;
   2 use crate::websocket::server::*;
   3 use actix_web::{Error, Result};
   4
   5 pub fn config(cfg: &mut web::ServiceConfig) {
   6   cfg.service(web::resource("/api/v1/ws").to(chat_route));
   7 }
   8
   9 /// How often heartbeat pings are sent
  10 const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5);
  11 /// How long before lack of client response causes a timeout
  12 const CLIENT_TIMEOUT: Duration = Duration::from_secs(10);
  13
  14 /// Entry point for our route
  15 async fn chat_route(
  16   req: HttpRequest,
  17   stream: web::Payload,
  18   chat_server: web::Data<Addr<ChatServer>>,
  19 ) -> Result<HttpResponse, Error> {
  20   ws::start(
  21     WSSession {
  22       cs_addr: chat_server.get_ref().to_owned(),
  23       id: 0,
  24       hb: Instant::now(),
  25       ip: get_ip(&req),
  26     },
  27     &req,
  28     stream,
  29   )
  30 }
  31
  32 struct WSSession {
  33   cs_addr: Addr<ChatServer>,
  34   /// unique session id
  35   id: usize,
  36   ip: String,
  37   /// Client must send ping at least once per 10 seconds (CLIENT_TIMEOUT),
  38   /// otherwise we drop connection.
  39   hb: Instant,
  40   // db: Pool<ConnectionManager<PgConnection>>,
  41 }
  42
  43 impl Actor for WSSession {
  44   type Context = ws::WebsocketContext<Self>;
  45
  46   /// Method is called on actor start.
  47   /// We register ws session with ChatServer
  48   fn started(&mut self, ctx: &mut Self::Context) {
  49     // we'll start heartbeat process on session start.
  50     self.hb(ctx);
  51
  52     // register self in chat server. `AsyncContext::wait` register
  53     // future within context, but context waits until this future resolves
  54     // before processing any other events.
  55     // across all routes within application
  56     let addr = ctx.address();
  57     self
  58       .cs_addr
  59       .send(Connect {
  60         addr: addr.recipient(),
  61         ip: self.ip.to_owned(),
  62       })
  63       .into_actor(self)
  64       .then(|res, act, ctx| {
  65         match res {
  66           Ok(res) => act.id = res,
  67           // something is wrong with chat server
  68           _ => ctx.stop(),
  69         }
  70         actix::fut::ready(())
  71       })
  72       .wait(ctx);
  73   }
  74
  75   fn stopping(&mut self, _ctx: &mut Self::Context) -> Running {
  76     // notify chat server
  77     self.cs_addr.do_send(Disconnect {
  78       id: self.id,
  79       ip: self.ip.to_owned(),
  80     });
  81     Running::Stop
  82   }
  83 }
  84
  85 /// Handle messages from chat server, we simply send it to peer websocket
  86 /// These are room messages, IE sent to others in the room
  87 impl Handler<WSMessage> for WSSession {
  88   type Result = ();
  89
  90   fn handle(&mut self, msg: WSMessage, ctx: &mut Self::Context) {
  91     ctx.text(msg.0);
  92   }
  93 }
  94
  95 /// WebSocket message handler
  96 impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for WSSession {
  97   fn handle(&mut self, result: Result<ws::Message, ws::ProtocolError>, ctx: &mut Self::Context) {
  98     let message = match result {
  99       Ok(m) => m,
 100       Err(e) => {
 101         error!("{}", e);
 102         return;
 103       }
 104     };
 105     match message {
 106       ws::Message::Ping(msg) => {
 107         self.hb = Instant::now();
 108         ctx.pong(&msg);
 109       }
 110       ws::Message::Pong(_) => {
 111         self.hb = Instant::now();
 112       }
 113       ws::Message::Text(text) => {
 114         let m = text.trim().to_owned();
 115         info!("Message received: {:?} from id: {}", &m, self.id);
 116
 117         self
 118           .cs_addr
 119           .send(StandardMessage {
 120             id: self.id,
 121             msg: m,
 122           })
 123           .into_actor(self)
 124           .then(|res, _, ctx| {
 125             match res {
 126               Ok(Ok(res)) => ctx.text(res),
 127               Ok(Err(e)) => error!("{}", e),
 128               Err(e) => error!("{}", &e),
 129             }
 130             actix::fut::ready(())
 131           })
 132           .wait(ctx);
 133       }
 134       ws::Message::Binary(_bin) => info!("Unexpected binary"),
 135       ws::Message::Close(_) => {
 136         ctx.stop();
 137       }
 138       _ => {}
 139     }
 140   }
 141 }
 142
 143 impl WSSession {
 144   /// helper method that sends ping to client every second.
 145   ///
 146   /// also this method checks heartbeats from client
 147   fn hb(&self, ctx: &mut ws::WebsocketContext<Self>) {
 148     ctx.run_interval(HEARTBEAT_INTERVAL, |act, ctx| {
 149       // check client heartbeats
 150       if Instant::now().duration_since(act.hb) > CLIENT_TIMEOUT {
 151         // heartbeat timed out
 152         error!("Websocket Client heartbeat failed, disconnecting!");
 153
 154         // notify chat server
 155         act.cs_addr.do_send(Disconnect {
 156           id: act.id,
 157           ip: act.ip.to_owned(),
 158         });
 159
 160         // stop actor
 161         ctx.stop();
 162
 163         // don't try to send a ping
 164         return;
 165       }
 166
 167       ctx.ping(b"");
 168     });
 169   }
 170 }