Automatically close unresponsive websockets
This commit is contained in:
parent
dff4384bd8
commit
a866deb3e4
@ -1,4 +1,5 @@
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use actix::{Actor, ActorContext, AsyncContext, Handler, Message, StreamHandler};
|
||||
use actix_web::{Error, HttpRequest, HttpResponse, web};
|
||||
@ -9,6 +10,12 @@ use tokio::net::TcpStream;
|
||||
|
||||
use crate::args::Args;
|
||||
|
||||
/// How often heartbeat pings are sent
|
||||
const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5);
|
||||
|
||||
/// How long before lack of client response causes a timeout
|
||||
const CLIENT_TIMEOUT: Duration = Duration::from_secs(60);
|
||||
|
||||
#[derive(Message)]
|
||||
#[rtype(result = "bool")]
|
||||
pub struct DataForWebSocket(Vec<u8>);
|
||||
@ -19,15 +26,42 @@ struct RelayWS {
|
||||
tcp_read: Option<OwnedReadHalf>,
|
||||
tcp_write: OwnedWriteHalf,
|
||||
|
||||
// TODO : add disconnect after ping timeout
|
||||
// Client must respond to ping at a specific interval, otherwise we drop connection
|
||||
hb: Instant,
|
||||
|
||||
// TODO : handle socket close
|
||||
}
|
||||
|
||||
impl RelayWS {
|
||||
/// helper method that sends ping to client every second.
|
||||
///
|
||||
/// also this method checks heartbeats from client
|
||||
fn hb(&self, ctx: &mut ws::WebsocketContext<Self>) {
|
||||
ctx.run_interval(HEARTBEAT_INTERVAL, |act, ctx| {
|
||||
// check client heartbeats
|
||||
if Instant::now().duration_since(act.hb) > CLIENT_TIMEOUT {
|
||||
// heartbeat timed out
|
||||
log::warn!("WebSocket Client heartbeat failed, disconnecting!");
|
||||
|
||||
// stop actor
|
||||
ctx.stop();
|
||||
|
||||
// don't try to send a ping
|
||||
return;
|
||||
}
|
||||
|
||||
log::debug!("Send ping message...");
|
||||
ctx.ping(b"");
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
impl Actor for RelayWS {
|
||||
type Context = ws::WebsocketContext<Self>;
|
||||
|
||||
fn started(&mut self, ctx: &mut Self::Context) {
|
||||
self.hb(ctx);
|
||||
|
||||
// Start to read on remote socket
|
||||
let mut read_half = self.tcp_read.take().unwrap();
|
||||
let addr = ctx.address();
|
||||
@ -67,6 +101,7 @@ impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for RelayWS {
|
||||
fn handle(&mut self, msg: Result<ws::Message, ws::ProtocolError>, ctx: &mut Self::Context) {
|
||||
match msg {
|
||||
Ok(ws::Message::Ping(msg)) => ctx.pong(&msg),
|
||||
Ok(ws::Message::Pong(_)) => self.hb = Instant::now(),
|
||||
Ok(ws::Message::Text(text)) => ctx.text(text),
|
||||
Ok(ws::Message::Close(_reason)) => ctx.stop(),
|
||||
Ok(ws::Message::Binary(data)) => {
|
||||
@ -125,7 +160,7 @@ pub async fn relay_ws(req: HttpRequest, stream: web::Payload,
|
||||
}
|
||||
};
|
||||
|
||||
let relay = RelayWS { tcp_read: Some(tcp_read), tcp_write };
|
||||
let relay = RelayWS { tcp_read: Some(tcp_read), tcp_write, hb: Instant::now() };
|
||||
let resp = ws::start(relay, &req, stream);
|
||||
log::info!("Opening new WS connection for {:?} to {}", req.peer_addr(), upstream_addr);
|
||||
resp
|
||||
|
Loading…
Reference in New Issue
Block a user