chore(monorepo): delta, january, quark
This commit is contained in:
Paul Makles
150
crates/bonfire/src/config.rs
Normal file
150
crates/bonfire/src/config.rs
Normal file
@@ -0,0 +1,150 @@
|
||||
use async_tungstenite::tungstenite::{handshake, Message};
|
||||
use futures::channel::oneshot::Sender;
|
||||
use revolt_quark::{Error, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Enumeration of supported protocol formats
|
||||
#[derive(Debug)]
|
||||
pub enum ProtocolFormat {
|
||||
Json,
|
||||
Msgpack,
|
||||
}
|
||||
|
||||
/// User-provided protocol configuration
|
||||
#[derive(Debug)]
|
||||
pub struct ProtocolConfiguration {
|
||||
protocol_version: i32,
|
||||
format: ProtocolFormat,
|
||||
session_token: Option<String>,
|
||||
}
|
||||
|
||||
impl ProtocolConfiguration {
|
||||
/// Create a new protocol configuration object from provided data
|
||||
pub fn from(
|
||||
protocol_version: i32,
|
||||
format: ProtocolFormat,
|
||||
session_token: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
protocol_version,
|
||||
format,
|
||||
session_token,
|
||||
}
|
||||
}
|
||||
|
||||
/// Decode some WebSocket message into a T: Deserialize using the client's specified protocol format
|
||||
pub fn decode<'a, T: Deserialize<'a>>(&self, msg: &'a Message) -> Result<T> {
|
||||
match self.format {
|
||||
ProtocolFormat::Json => {
|
||||
if let Message::Text(text) = msg {
|
||||
serde_json::from_str(text).map_err(|_| Error::InternalError)
|
||||
} else {
|
||||
Err(Error::InternalError)
|
||||
}
|
||||
}
|
||||
ProtocolFormat::Msgpack => {
|
||||
if let Message::Binary(buf) = msg {
|
||||
rmp_serde::from_slice(buf).map_err(|_| Error::InternalError)
|
||||
} else {
|
||||
Err(Error::InternalError)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Encode T: Serialize into a WebSocket message using the client's specified protocol format
|
||||
pub fn encode<T: Serialize>(&self, data: &T) -> Message {
|
||||
match self.format {
|
||||
ProtocolFormat::Json => {
|
||||
Message::Text(serde_json::to_string(data).expect("Failed to serialise (as json)."))
|
||||
}
|
||||
ProtocolFormat::Msgpack => Message::Binary(
|
||||
rmp_serde::to_vec_named(data).expect("Failed to serialise (as msgpack)."),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the current session token
|
||||
pub fn set_session_token(&mut self, token: String) {
|
||||
self.session_token.replace(token);
|
||||
}
|
||||
|
||||
/// Get the current session token
|
||||
pub fn get_session_token(&self) -> &Option<String> {
|
||||
&self.session_token
|
||||
}
|
||||
|
||||
/// Get the protocol version specified
|
||||
pub fn get_protocol_version(&self) -> i32 {
|
||||
self.protocol_version
|
||||
}
|
||||
|
||||
/// Get the protocol format specified
|
||||
pub fn get_protocol_format(&self) -> &ProtocolFormat {
|
||||
&self.format
|
||||
}
|
||||
}
|
||||
|
||||
/// Object holding one side of a channel for receiving the parsed information
|
||||
pub struct WebsocketHandshakeCallback {
|
||||
sender: Sender<ProtocolConfiguration>,
|
||||
}
|
||||
|
||||
impl WebsocketHandshakeCallback {
|
||||
/// Create a callback using a given sender
|
||||
pub fn from(sender: Sender<ProtocolConfiguration>) -> Self {
|
||||
Self { sender }
|
||||
}
|
||||
}
|
||||
|
||||
impl handshake::server::Callback for WebsocketHandshakeCallback {
|
||||
/// Handle request to create a new WebSocket connection
|
||||
fn on_request(
|
||||
self,
|
||||
request: &handshake::server::Request,
|
||||
response: handshake::server::Response,
|
||||
) -> Result<handshake::server::Response, handshake::server::ErrorResponse> {
|
||||
// Take and parse query parameters from the URI.
|
||||
let query = request.uri().query().unwrap_or_default();
|
||||
let params = querystring::querify(query);
|
||||
|
||||
// Set default values for the protocol.
|
||||
let mut protocol_version = 1;
|
||||
let mut format = ProtocolFormat::Json;
|
||||
let mut session_token = None;
|
||||
|
||||
// Parse and map parameters from key-value to known variables.
|
||||
for (key, value) in params {
|
||||
match key {
|
||||
"version" => {
|
||||
if let Ok(version) = value.parse() {
|
||||
protocol_version = version;
|
||||
}
|
||||
}
|
||||
"format" => match value {
|
||||
"json" => format = ProtocolFormat::Json,
|
||||
"msgpack" => format = ProtocolFormat::Msgpack,
|
||||
_ => {}
|
||||
},
|
||||
"token" => session_token = Some(value.into()),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
// Send configuration information back from this callback.
|
||||
// We have to use a channel as this function does not borrow mutably.
|
||||
if self
|
||||
.sender
|
||||
.send(ProtocolConfiguration {
|
||||
protocol_version,
|
||||
format,
|
||||
session_token,
|
||||
})
|
||||
.is_ok()
|
||||
{
|
||||
Ok(response)
|
||||
} else {
|
||||
Err(handshake::server::ErrorResponse::new(None))
|
||||
}
|
||||
}
|
||||
}
|
||||
19
crates/bonfire/src/database.rs
Normal file
19
crates/bonfire/src/database.rs
Normal file
@@ -0,0 +1,19 @@
|
||||
use once_cell::sync::OnceCell;
|
||||
use revolt_quark::{Database, DatabaseInfo};
|
||||
|
||||
static DBCONN: OnceCell<Database> = OnceCell::new();
|
||||
|
||||
/// Connect Bonfire to the database.
|
||||
pub async fn connect() {
|
||||
let database = DatabaseInfo::Auto
|
||||
.connect()
|
||||
.await
|
||||
.expect("Failed to connect to the database.");
|
||||
|
||||
DBCONN.set(database).expect("Setting `Database`");
|
||||
}
|
||||
|
||||
/// Get a reference to the current database.
|
||||
pub fn get_db() -> &'static Database {
|
||||
DBCONN.get().expect("Valid `Database`")
|
||||
}
|
||||
34
crates/bonfire/src/main.rs
Normal file
34
crates/bonfire/src/main.rs
Normal file
@@ -0,0 +1,34 @@
|
||||
use std::env;
|
||||
|
||||
use async_std::net::TcpListener;
|
||||
use revolt_quark::presence::presence_clear_region;
|
||||
|
||||
#[macro_use]
|
||||
extern crate log;
|
||||
|
||||
pub mod config;
|
||||
|
||||
mod database;
|
||||
mod websocket;
|
||||
|
||||
#[async_std::main]
|
||||
async fn main() {
|
||||
// Configure requirements for Bonfire.
|
||||
let _guard = revolt_quark::setup_logging();
|
||||
database::connect().await;
|
||||
|
||||
// Clean up the current region information.
|
||||
presence_clear_region(None).await;
|
||||
|
||||
// Setup a TCP listener to accept WebSocket connections on.
|
||||
// By default, we bind to port 9000 on all interfaces.
|
||||
let bind = env::var("HOST").unwrap_or_else(|_| "0.0.0.0:9000".into());
|
||||
info!("Listening on host {bind}");
|
||||
let try_socket = TcpListener::bind(bind).await;
|
||||
let listener = try_socket.expect("Failed to bind");
|
||||
|
||||
// Start accepting new connections and spawn a client for each connection.
|
||||
while let Ok((stream, addr)) = listener.accept().await {
|
||||
websocket::spawn_client(database::get_db(), stream, addr);
|
||||
}
|
||||
}
|
||||
253
crates/bonfire/src/websocket.rs
Normal file
253
crates/bonfire/src/websocket.rs
Normal file
@@ -0,0 +1,253 @@
|
||||
use std::net::SocketAddr;
|
||||
|
||||
use futures::{channel::oneshot, pin_mut, select, FutureExt, SinkExt, StreamExt, TryStreamExt};
|
||||
use revolt_quark::{
|
||||
events::{
|
||||
client::EventV1,
|
||||
server::ClientMessage,
|
||||
state::{State, SubscriptionStateChange},
|
||||
},
|
||||
models::{user::UserHint, User},
|
||||
presence::{presence_create_session, presence_delete_session},
|
||||
redis_kiss, Database,
|
||||
};
|
||||
|
||||
use async_std::{net::TcpStream, sync::Mutex, task};
|
||||
|
||||
use crate::config::WebsocketHandshakeCallback;
|
||||
|
||||
/// Spawn a new WebSocket client worker given access to the database,
|
||||
/// the relevant TCP stream and the remote address of the client.
|
||||
pub fn spawn_client(db: &'static Database, stream: TcpStream, addr: SocketAddr) {
|
||||
// Spawn a new Async task to work on.
|
||||
task::spawn(async move {
|
||||
info!("User connected from {addr:?}");
|
||||
|
||||
// Upgrade the TCP connection to a WebSocket connection.
|
||||
// In this process, we also parse any additional parameters given.
|
||||
// e.g. wss://example.com?format=json&version=1
|
||||
let (sender, receiver) = oneshot::channel();
|
||||
if let Ok(ws) = async_tungstenite::accept_hdr_async_with_config(
|
||||
stream,
|
||||
WebsocketHandshakeCallback::from(sender),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
{
|
||||
// Verify we've received a valid config, otherwise we should just drop the connection.
|
||||
if let Ok(mut config) = receiver.await {
|
||||
info!(
|
||||
"User {addr:?} provided protocol configuration (version = {}, format = {:?})",
|
||||
config.get_protocol_version(),
|
||||
config.get_protocol_format()
|
||||
);
|
||||
|
||||
// Split the socket for simultaneously read and write.
|
||||
let (write, mut read) = ws.split();
|
||||
let write = Mutex::new(write);
|
||||
|
||||
// If the user has not provided authentication, request information.
|
||||
if config.get_session_token().is_none() {
|
||||
'outer: while let Ok(message) = read.try_next().await {
|
||||
if let Ok(ClientMessage::Authenticate { token }) =
|
||||
config.decode(message.as_ref().unwrap())
|
||||
{
|
||||
config.set_session_token(token);
|
||||
break 'outer;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Try to authenticate the user.
|
||||
if let Some(token) = config.get_session_token().as_ref() {
|
||||
match User::from_token(db, token, UserHint::Any).await {
|
||||
Ok(user) => {
|
||||
info!("User {addr:?} authenticated as @{}", user.username);
|
||||
|
||||
// Create local state.
|
||||
let mut state = State::from(user);
|
||||
let user_id = state.cache.user_id.clone();
|
||||
|
||||
// Create presence session.
|
||||
let (first_session, session_id) =
|
||||
presence_create_session(&user_id, 0).await;
|
||||
|
||||
// Notify socket we have authenticated.
|
||||
write
|
||||
.lock()
|
||||
.await
|
||||
.send(config.encode(&EventV1::Authenticated))
|
||||
.await
|
||||
.ok();
|
||||
|
||||
// Download required data to local cache and send Ready payload.
|
||||
if let Ok(ready_payload) = state.generate_ready_payload(db).await {
|
||||
write
|
||||
.lock()
|
||||
.await
|
||||
.send(config.encode(&ready_payload))
|
||||
.await
|
||||
.ok();
|
||||
|
||||
// If this was the first session, notify other users that we just went online.
|
||||
if first_session {
|
||||
state.broadcast_presence_change(true).await;
|
||||
}
|
||||
|
||||
// Create a PubSub connection to poll on.
|
||||
let listener = async {
|
||||
if let Ok(mut conn) = redis_kiss::open_pubsub_connection().await
|
||||
{
|
||||
loop {
|
||||
// Check for state changes for subscriptions.
|
||||
match state.apply_state() {
|
||||
SubscriptionStateChange::Reset => {
|
||||
for id in state.iter_subscriptions() {
|
||||
conn.subscribe(id).await.unwrap();
|
||||
}
|
||||
|
||||
#[cfg(debug_assertions)]
|
||||
info!("{addr:?} has reset their subscriptions");
|
||||
}
|
||||
SubscriptionStateChange::Change { add, remove } => {
|
||||
for id in remove {
|
||||
#[cfg(debug_assertions)]
|
||||
info!("{addr:?} unsubscribing from {id}");
|
||||
|
||||
conn.unsubscribe(id).await.unwrap();
|
||||
}
|
||||
|
||||
for id in add {
|
||||
#[cfg(debug_assertions)]
|
||||
info!("{addr:?} subscribing to {id}");
|
||||
|
||||
conn.subscribe(id).await.unwrap();
|
||||
}
|
||||
}
|
||||
SubscriptionStateChange::None => {}
|
||||
}
|
||||
|
||||
// * Debug logging of current subscriptions.
|
||||
/*#[cfg(debug_assertions)]
|
||||
info!(
|
||||
"User {addr:?} is subscribed to {:?}",
|
||||
state
|
||||
.iter_subscriptions()
|
||||
.collect::<Vec<&String>>()
|
||||
);*/
|
||||
|
||||
// Handle incoming events.
|
||||
match conn.on_message().next().await.map(|item| {
|
||||
(
|
||||
item.get_channel_name().to_string(),
|
||||
redis_kiss::decode_payload::<EventV1>(&item),
|
||||
)
|
||||
}) {
|
||||
Some((channel, item)) => {
|
||||
if let Ok(mut event) = item {
|
||||
if state
|
||||
.handle_incoming_event_v1(
|
||||
db, &mut event,
|
||||
)
|
||||
.await
|
||||
&& write.lock().await
|
||||
.send(config.encode(&event))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
warn!("Failed to deserialise an event for {channel}!");
|
||||
}
|
||||
}
|
||||
// No more data, assume we disconnected or otherwise
|
||||
// something bad occurred, so disconnect user.
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
.fuse();
|
||||
|
||||
// Read from WebSocket stream.
|
||||
let worker =
|
||||
async {
|
||||
while let Ok(Some(msg)) = read.try_next().await {
|
||||
if let Ok(payload) = config.decode(&msg) {
|
||||
match payload {
|
||||
ClientMessage::BeginTyping { channel } => {
|
||||
EventV1::ChannelStartTyping {
|
||||
id: channel.clone(),
|
||||
user: user_id.clone(),
|
||||
}
|
||||
.p(channel.clone())
|
||||
.await;
|
||||
}
|
||||
ClientMessage::EndTyping { channel } => {
|
||||
EventV1::ChannelStopTyping {
|
||||
id: channel.clone(),
|
||||
user: user_id.clone(),
|
||||
}
|
||||
.p(channel.clone())
|
||||
.await;
|
||||
}
|
||||
ClientMessage::Ping { data, responded } => {
|
||||
if responded.is_none() {
|
||||
write
|
||||
.lock()
|
||||
.await
|
||||
.send(config.encode(
|
||||
&EventV1::Pong { data },
|
||||
))
|
||||
.await
|
||||
.ok();
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
.fuse();
|
||||
|
||||
// Pin both tasks.
|
||||
pin_mut!(listener, worker);
|
||||
|
||||
// Wait for either disconnect or for listener to die.
|
||||
select!(
|
||||
() = listener => {},
|
||||
() = worker => {}
|
||||
);
|
||||
|
||||
// * Combine the streams back once we are ready to disconnect.
|
||||
/* ws = read.reunite(write).unwrap(); */
|
||||
}
|
||||
|
||||
// Clean up presence session.
|
||||
let last_session = presence_delete_session(&user_id, session_id).await;
|
||||
|
||||
// If this was the last session, notify other users that we just went offline.
|
||||
if last_session {
|
||||
state.broadcast_presence_change(false).await;
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
write.lock().await.send(config.encode(&err)).await.ok();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// * Disconnect the WebSocket if it isn't already.
|
||||
/*ws.close(Some(CloseFrame {
|
||||
code: CloseCode::Normal,
|
||||
reason: std::borrow::Cow::from(""),
|
||||
}))
|
||||
.await
|
||||
.unwrap();*/
|
||||
}
|
||||
|
||||
info!("User disconnected from {addr:?}");
|
||||
});
|
||||
}
|
||||
Reference in New Issue
Block a user