feat: Initialize server websocket

Related-to: #2
Reviewed-by: Amjad Alsharafi <me@amjad.alsharafi.dev>
Signed-off-by: Awiteb <a@4rs.nl>
This commit is contained in:
Awiteb 2024-07-05 02:19:58 +03:00
parent cd2a9ea03e
commit c0d5efe0c3
Signed by: awiteb
GPG key ID: 3F6B55640AA6682F
8 changed files with 491 additions and 12 deletions

View file

@ -142,17 +142,14 @@ impl OnlineUsersExt for OnlineUsers {
} }
async fn disconnect_inactive_users(&self) { async fn disconnect_inactive_users(&self) {
let now = Utc::now().timestamp(); self.write().await.retain(|_, u| {
let is_inactive = // if we send ping and the client doesn't send pong
|u: &SocketUserData| u.pinged_at > u.ponged_at && now - u.pinged_at.timestamp() >= 5; if u.pinged_at > u.ponged_at {
self.read()
.await
.iter()
.filter(|(_, u)| is_inactive(u))
.for_each(|(_, u)| {
log::info!("Disconnected from {}, inactive", u.public_key); log::info!("Disconnected from {}, inactive", u.public_key);
u.sender.close_channel() u.sender.close_channel();
return false;
}
true
}); });
self.write().await.retain(|_, u| !is_inactive(u));
} }
} }

View file

@ -30,6 +30,7 @@ mod middlewares;
mod routes; mod routes;
mod schemas; mod schemas;
mod utils; mod utils;
mod websocket;
/// Nonce cache type, used to store nonces for a certain amount of time /// Nonce cache type, used to store nonces for a certain amount of time
pub type NonceCache = Mutex<HashMap<[u8; 16], i64>>; pub type NonceCache = Mutex<HashMap<[u8; 16], i64>>;

View file

@ -25,7 +25,7 @@ use salvo::rate_limiter::{BasicQuota, FixedGuard, MokaStore, RateLimiter, Remote
use salvo::{catcher::Catcher, logging::Logger, prelude::*}; use salvo::{catcher::Catcher, logging::Logger, prelude::*};
use crate::schemas::MessageSchema; use crate::schemas::MessageSchema;
use crate::{middlewares, NonceCache}; use crate::{middlewares, websocket, NonceCache};
mod user; mod user;
@ -141,6 +141,7 @@ pub fn service(conn: sea_orm::DatabaseConnection, config: &Config) -> Service {
let router = Router::new() let router = Router::new()
.push(Router::with_path("user").push(user::route())) .push(Router::with_path("user").push(user::route()))
.push(Router::with_path("ws").push(websocket::route()))
.hoop(middlewares::add_server_headers) .hoop(middlewares::add_server_headers)
.hoop(Logger::new()) .hoop(Logger::new())
.hoop( .hoop(

View file

@ -0,0 +1,57 @@
// OxideTalis Messaging Protocol homeserver implementation
// Copyright (C) 2024 OxideTalis Developers <otmp@4rs.nl>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <https://gnu.org/licenses/agpl-3.0>.
//! Websocket errors
/// Result type of websocket
pub type WsResult<T> = Result<T, WsError>;
/// Websocket errors, returned in the websocket communication
#[derive(Debug)]
pub enum WsError {
/// The signature is invalid
InvalidSignature,
/// Message type must be text
NotTextMessage,
/// Invalid json data
InvalidJsonData,
/// Unknown client event
UnknownClientEvent,
}
impl WsError {
/// Returns error name
pub const fn name(&self) -> &'static str {
match self {
WsError::InvalidSignature => "InvalidSignature",
WsError::NotTextMessage => "NotTextMessage",
WsError::InvalidJsonData => "InvalidJsonData",
WsError::UnknownClientEvent => "UnknownClientEvent",
}
}
/// Returns the error reason
pub const fn reason(&self) -> &'static str {
match self {
WsError::InvalidSignature => "Invalid event signature",
WsError::NotTextMessage => "The websocket message must be text message",
WsError::InvalidJsonData => "Received invalid json data, the text must be valid json",
WsError::UnknownClientEvent => {
"Unknown client event, the event is not recognized by the server"
}
}
}
}

View file

@ -0,0 +1,61 @@
// OxideTalis Messaging Protocol homeserver implementation
// Copyright (C) 2024 OxideTalis Developers <otmp@4rs.nl>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <https://gnu.org/licenses/agpl-3.0>.
//! Events that the client send it
use oxidetalis_core::types::Signature;
use serde::{Deserialize, Serialize};
use crate::{utils, NonceCache};
/// Client websocket event
#[derive(Deserialize, Clone, Debug)]
pub struct ClientEvent {
pub event: ClientEventType,
signature: Signature,
}
/// Client websocket event type
#[derive(Serialize, Deserialize, Clone, Eq, PartialEq, Debug)]
#[serde(rename_all = "PascalCase", tag = "event", content = "data")]
pub enum ClientEventType {
/// Ping event
Ping { timestamp: u64 },
/// Pong event
Pong { timestamp: u64 },
}
impl ClientEventType {
/// Returns event data as json bytes
pub fn data(&self) -> Vec<u8> {
serde_json::to_value(self).expect("can't fail")["data"]
.to_string()
.into_bytes()
}
}
impl ClientEvent {
/// Verify the signature of the event
pub fn verify_signature(
&self,
shared_secret: &[u8; 32],
nonce_cache: &NonceCache,
nonce_limit: &usize,
) -> bool {
utils::is_valid_nonce(&self.signature, nonce_cache, nonce_limit)
&& self.signature.verify(&self.event.data(), shared_secret)
}
}

View file

@ -0,0 +1,23 @@
// OxideTalis Messaging Protocol homeserver implementation
// Copyright (C) 2024 OxideTalis Developers <otmp@4rs.nl>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <https://gnu.org/licenses/agpl-3.0>.
//! Server and client websocket events
mod client;
mod server;
pub use client::*;
pub use server::*;

View file

@ -0,0 +1,117 @@
// OxideTalis Messaging Protocol homeserver implementation
// Copyright (C) 2024 OxideTalis Developers <otmp@4rs.nl>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <https://gnu.org/licenses/agpl-3.0>.
//! Events that the server send it
use std::marker::PhantomData;
use chrono::Utc;
use oxidetalis_core::{cipher::K256Secret, types::Signature};
use salvo::websocket::Message;
use serde::Serialize;
use crate::websocket::errors::WsError;
/// Signed marker, used to indicate that the event is signed
pub struct Signed;
/// Unsigned marker, used to indicate that the event is unsigned
pub struct Unsigned;
/// Server websocket event
#[derive(Serialize, Clone, Debug)]
pub struct ServerEvent<T> {
#[serde(flatten)]
event: ServerEventType,
signature: Signature,
#[serde(skip)]
phantom: PhantomData<T>,
}
/// server websocket event type
#[derive(Serialize, Clone, Eq, PartialEq, Debug)]
#[serde(rename_all = "PascalCase")]
pub enum ServerEventType {
/// Ping event
Ping { timestamp: u64 },
/// Pong event
Pong { timestamp: u64 },
/// Error event
Error {
name: &'static str,
reason: &'static str,
},
}
impl ServerEventType {
/// Returns event data as json bytes
pub fn data(&self) -> Vec<u8> {
serde_json::to_value(self).expect("can't fail")["data"]
.to_string()
.into_bytes()
}
}
impl ServerEvent<Unsigned> {
/// Creates new [`ServerEvent`]
pub fn new(event: ServerEventType) -> Self {
Self {
event,
signature: Signature::from([0u8; 56]),
phantom: PhantomData,
}
}
/// Creates ping event
pub fn ping() -> Self {
Self::new(ServerEventType::Ping {
timestamp: Utc::now().timestamp() as u64,
})
}
/// Creates pong event
pub fn pong() -> Self {
Self::new(ServerEventType::Pong {
timestamp: Utc::now().timestamp() as u64,
})
}
/// Sign the event
pub fn sign(self, shared_secret: &[u8; 32]) -> ServerEvent<Signed> {
ServerEvent::<Signed> {
signature: K256Secret::sign_with_shared_secret(
&serde_json::to_vec(&self.event.data()).expect("Can't fail"),
shared_secret,
),
event: self.event,
phantom: PhantomData,
}
}
}
impl From<&ServerEvent<Signed>> for Message {
fn from(value: &ServerEvent<Signed>) -> Self {
Message::text(serde_json::to_string(value).expect("This can't fail"))
}
}
impl From<WsError> for ServerEvent<Unsigned> {
fn from(err: WsError) -> Self {
ServerEvent::new(ServerEventType::Error {
name: err.name(),
reason: err.reason(),
})
}
}

View file

@ -0,0 +1,222 @@
// OxideTalis Messaging Protocol homeserver implementation
// Copyright (C) 2024 OxideTalis Developers <otmp@4rs.nl>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <https://gnu.org/licenses/agpl-3.0>.
use std::{collections::HashMap, sync::Arc, time::Duration};
use chrono::Utc;
use errors::{WsError, WsResult};
use futures::{channel::mpsc, FutureExt, StreamExt, TryStreamExt};
use once_cell::sync::Lazy;
use oxidetalis_core::{cipher::K256Secret, types::PublicKey};
use salvo::{
handler,
http::StatusError,
websocket::{Message, WebSocket, WebSocketUpgrade},
Depot,
Request,
Response,
Router,
};
use tokio::{sync::RwLock, task::spawn as tokio_spawn, time::sleep as tokio_sleep};
mod errors;
mod events;
pub use events::*;
use uuid::Uuid;
use crate::{
extensions::{DepotExt, OnlineUsersExt},
middlewares,
utils,
NonceCache,
};
/// Online users type
pub type OnlineUsers = RwLock<HashMap<Uuid, SocketUserData>>;
/// List of online users, users that are connected to the server
// FIXME: Use `std::sync::LazyLock` after it becomes stable in `1.80.0`
static ONLINE_USERS: Lazy<OnlineUsers> = Lazy::new(OnlineUsers::default);
/// A user connected to the server
pub struct SocketUserData {
/// Sender to send messages to the user
pub sender: mpsc::UnboundedSender<salvo::Result<Message>>,
/// User public key
pub public_key: PublicKey,
/// Time that the user pinged at
pub pinged_at: chrono::DateTime<Utc>,
/// Time that the user ponged at
pub ponged_at: chrono::DateTime<Utc>,
/// User shared secret
pub shared_secret: [u8; 32],
}
impl SocketUserData {
/// Creates new [`SocketUserData`]
pub fn new(
public_key: PublicKey,
shared_secret: [u8; 32],
sender: mpsc::UnboundedSender<salvo::Result<Message>>,
) -> Self {
let now = Utc::now();
Self {
sender,
public_key,
shared_secret,
pinged_at: now,
ponged_at: now,
}
}
}
/// WebSocket handler, that handles the user connection
#[handler]
pub async fn user_connected(
req: &mut Request,
res: &mut Response,
depot: &Depot,
) -> Result<(), StatusError> {
let nonce_cache = depot.nonce_cache();
let nonce_limit = *depot.nonce_cache_size();
let public_key =
utils::extract_public_key(req).expect("The public key was checked in the middleware");
// FIXME: The config should hold `K256Secret` not `PrivateKey`
let shared_secret =
K256Secret::from_privkey(&depot.config().server.private_key).shared_secret(&public_key);
WebSocketUpgrade::new()
.upgrade(req, res, move |ws| {
handle_socket(ws, nonce_cache, nonce_limit, public_key, shared_secret)
})
.await
}
/// Handle the websocket connection
async fn handle_socket(
ws: WebSocket,
nonce_cache: Arc<NonceCache>,
nonce_limit: usize,
user_public_key: PublicKey,
user_shared_secret: [u8; 32],
) {
let (user_ws_sender, mut user_ws_receiver) = ws.split();
let (sender, receiver) = mpsc::unbounded();
let receiver = receiver.into_stream();
let fut = receiver.forward(user_ws_sender).map(|result| {
if let Err(err) = result {
log::error!("websocket send error: {err}");
}
});
tokio_spawn(fut);
let conn_id = Uuid::new_v4();
let user = SocketUserData::new(user_public_key, user_shared_secret, sender.clone());
ONLINE_USERS.add_user(&conn_id, user).await;
log::info!("New user connected: ConnId(={conn_id}) PublicKey(={user_public_key})");
let fut = async move {
while let Some(Ok(msg)) = user_ws_receiver.next().await {
match handle_ws_msg(msg, &nonce_cache, &nonce_limit, &user_shared_secret) {
Ok(event) => {
if let Some(server_event) = handle_events(event, &conn_id).await {
if let Err(err) = sender.unbounded_send(Ok(Message::from(
&server_event.sign(&user_shared_secret),
))) {
log::error!("Websocket Error: {err}");
break;
}
};
}
Err(err) => {
if let Err(err) = sender.unbounded_send(Ok(Message::from(
&ServerEvent::from(err).sign(&user_shared_secret),
))) {
log::error!("Websocket Error: {err}");
break;
};
}
};
}
user_disconnected(&conn_id, &user_public_key).await;
};
tokio_spawn(fut);
}
/// Handle websocket msg
fn handle_ws_msg(
msg: Message,
nonce_cache: &NonceCache,
nonce_limit: &usize,
shared_secret: &[u8; 32],
) -> WsResult<ClientEvent> {
let Ok(text) = msg.to_str() else {
return Err(WsError::NotTextMessage);
};
let event = serde_json::from_str::<ClientEvent>(text).map_err(|err| {
if err.is_data() {
WsError::UnknownClientEvent
} else {
WsError::InvalidJsonData
}
})?;
if !event.verify_signature(shared_secret, nonce_cache, nonce_limit) {
return Err(WsError::InvalidSignature);
}
Ok(event)
}
/// Handle user events, and return the server event if needed
async fn handle_events(event: ClientEvent, conn_id: &Uuid) -> Option<ServerEvent<Unsigned>> {
match &event.event {
ClientEventType::Ping { .. } => Some(ServerEvent::pong()),
ClientEventType::Pong { .. } => {
ONLINE_USERS.update_pong(conn_id).await;
None
}
}
}
/// Handle user disconnected
async fn user_disconnected(conn_id: &Uuid, public_key: &PublicKey) {
ONLINE_USERS.remove_user(conn_id).await;
log::debug!("User disconnect: ConnId(={conn_id}) PublicKey(={public_key})");
}
pub fn route() -> Router {
let users_pinger = async {
/// Seconds to wait for pongs, before disconnecting the user
const WAIT_FOR_PONGS_SECS: u32 = 10;
/// Seconds to sleep between pings (10 minutes)
const SLEEP_SECS: u32 = 60 * 10;
loop {
log::debug!("Start pinging online users");
ONLINE_USERS.ping_all().await;
tokio_sleep(Duration::from_secs(u64::from(WAIT_FOR_PONGS_SECS))).await;
ONLINE_USERS.disconnect_inactive_users().await;
log::debug!("Done pinging online users and disconnected inactive ones");
tokio_sleep(Duration::from_secs(u64::from(SLEEP_SECS))).await;
}
};
tokio_spawn(users_pinger);
Router::new()
.push(Router::with_path("chat").get(user_connected))
.hoop(middlewares::signature_check)
.hoop(middlewares::public_key_check)
}