|
|
|
@ -6,7 +6,7 @@ use crate::{
|
|
|
|
|
use actix::prelude::*;
|
|
|
|
|
use actix_web::{web, Error, HttpRequest, HttpResponse};
|
|
|
|
|
use actix_web_actors::ws;
|
|
|
|
|
use lemmy_utils::{utils::get_ip, ConnectionId, IpAddr};
|
|
|
|
|
use lemmy_utils::{rate_limit::RateLimit, utils::get_ip, ConnectionId, IpAddr};
|
|
|
|
|
use std::time::{Duration, Instant};
|
|
|
|
|
use tracing::{debug, error, info};
|
|
|
|
|
|
|
|
|
@ -20,6 +20,7 @@ pub async fn chat_route(
|
|
|
|
|
req: HttpRequest,
|
|
|
|
|
stream: web::Payload,
|
|
|
|
|
context: web::Data<LemmyContext>,
|
|
|
|
|
rate_limiter: web::Data<RateLimit>,
|
|
|
|
|
) -> Result<HttpResponse, Error> {
|
|
|
|
|
ws::start(
|
|
|
|
|
WsSession {
|
|
|
|
@ -27,6 +28,7 @@ pub async fn chat_route(
|
|
|
|
|
id: 0,
|
|
|
|
|
hb: Instant::now(),
|
|
|
|
|
ip: get_ip(&req.connection_info()),
|
|
|
|
|
rate_limiter: rate_limiter.as_ref().to_owned(),
|
|
|
|
|
},
|
|
|
|
|
&req,
|
|
|
|
|
stream,
|
|
|
|
@ -41,6 +43,8 @@ struct WsSession {
|
|
|
|
|
/// Client must send ping at least once per 10 seconds (CLIENT_TIMEOUT),
|
|
|
|
|
/// otherwise we drop connection.
|
|
|
|
|
hb: Instant,
|
|
|
|
|
/// A rate limiter for websocket joins
|
|
|
|
|
rate_limiter: RateLimit,
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
impl Actor for WsSession {
|
|
|
|
@ -57,6 +61,11 @@ impl Actor for WsSession {
|
|
|
|
|
// before processing any other events.
|
|
|
|
|
// across all routes within application
|
|
|
|
|
let addr = ctx.address();
|
|
|
|
|
|
|
|
|
|
if !self.rate_limit_check(ctx) {
|
|
|
|
|
return;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
self
|
|
|
|
|
.cs_addr
|
|
|
|
|
.send(Connect {
|
|
|
|
@ -98,6 +107,10 @@ impl Handler<WsMessage> for WsSession {
|
|
|
|
|
/// WebSocket message handler
|
|
|
|
|
impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for WsSession {
|
|
|
|
|
fn handle(&mut self, result: Result<ws::Message, ws::ProtocolError>, ctx: &mut Self::Context) {
|
|
|
|
|
if !self.rate_limit_check(ctx) {
|
|
|
|
|
return;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
let message = match result {
|
|
|
|
|
Ok(m) => m,
|
|
|
|
|
Err(e) => {
|
|
|
|
@ -169,4 +182,14 @@ impl WsSession {
|
|
|
|
|
ctx.ping(b"");
|
|
|
|
|
});
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// Check the rate limit, and stop the ctx if it fails
|
|
|
|
|
fn rate_limit_check(&mut self, ctx: &mut ws::WebsocketContext<Self>) -> bool {
|
|
|
|
|
let check = self.rate_limiter.message().check(self.ip.to_owned());
|
|
|
|
|
if !check {
|
|
|
|
|
debug!("Websocket join with IP: {} has been rate limited.", self.ip);
|
|
|
|
|
ctx.stop()
|
|
|
|
|
}
|
|
|
|
|
check
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|