use chrono::{DateTime, Duration, SubsecRound, Utc};
use diesel::prelude::*;
use diesel_async::pooled_connection::deadpool::Pool;
use diesel_async::{AsyncConnection as _, RunQueryDsl};
use uuid::Uuid;
use autumn_web::RuntimeConnection;
use crate::config::DEFAULT_ROOM_MAX_PARTICIPANTS;
use crate::rooms::{
JoinRecord, MAX_ROOMS, ParticipantView, ReapFuture, ReapStats, RoomError, RoomSnapshot,
RoomStore, RoomStoreFuture, SessionToken, renewed_expiry, validate_room_segment,
};
diesel::table! {
media_rooms (namespace, room_id) {
namespace -> diesel::sql_types::Text,
room_id -> diesel::sql_types::Text,
max_participants -> diesel::sql_types::Integer,
created_at -> diesel::sql_types::Timestamp,
}
}
diesel::table! {
media_room_participants (namespace, room_id, participant_id) {
namespace -> diesel::sql_types::Text,
room_id -> diesel::sql_types::Text,
participant_id -> diesel::sql_types::Text,
display_name -> diesel::sql_types::Nullable<diesel::sql_types::Text>,
token -> diesel::sql_types::Text,
joined_at -> diesel::sql_types::Timestamp,
token_expires_at -> diesel::sql_types::Timestamp,
last_seen_at -> diesel::sql_types::Timestamp,
}
}
diesel::allow_tables_to_appear_in_same_query!(media_rooms, media_room_participants);
#[derive(Queryable, Selectable, Insertable)]
#[diesel(table_name = media_rooms)]
struct RoomRow {
namespace: String,
room_id: String,
max_participants: i32,
created_at: chrono::NaiveDateTime,
}
#[derive(Queryable, Selectable, Insertable)]
#[diesel(table_name = media_room_participants)]
struct ParticipantRow {
namespace: String,
room_id: String,
participant_id: String,
display_name: Option<String>,
token: String,
joined_at: chrono::NaiveDateTime,
token_expires_at: chrono::NaiveDateTime,
last_seen_at: chrono::NaiveDateTime,
}
pub struct DbRoomStore {
pool: Pool<RuntimeConnection>,
hard_cap: usize,
max_rooms: usize,
}
impl DbRoomStore {
#[must_use]
pub const fn new(pool: Pool<RuntimeConnection>, hard_cap: usize) -> Self {
Self {
pool,
hard_cap,
max_rooms: MAX_ROOMS,
}
}
#[must_use]
pub const fn with_max_rooms(mut self, max_rooms: usize) -> Self {
self.max_rooms = max_rooms;
self
}
}
fn map_db_err<E: std::fmt::Display>(err: E) -> RoomError {
tracing::warn!(error = %err, "media rooms: db operation failed");
RoomError::Store
}
enum TxError {
Room(RoomError),
Db(diesel::result::Error),
}
impl From<diesel::result::Error> for TxError {
fn from(err: diesel::result::Error) -> Self {
Self::Db(err)
}
}
impl From<RoomError> for TxError {
fn from(err: RoomError) -> Self {
Self::Room(err)
}
}
impl TxError {
fn into_room(self) -> RoomError {
match self {
Self::Room(err) => err,
Self::Db(err) => map_db_err(err),
}
}
}
async fn lock_room(
conn: &mut RuntimeConnection,
namespace: &str,
room_id: &str,
) -> QueryResult<bool> {
let locked = diesel::update(
media_rooms::table.filter(
media_rooms::namespace
.eq(namespace)
.and(media_rooms::room_id.eq(room_id)),
),
)
.set(media_rooms::max_participants.eq(media_rooms::max_participants))
.execute(conn)
.await?;
Ok(locked > 0)
}
async fn detached<T, Fut>(work: impl FnOnce(Caller) -> Fut) -> Result<T, RoomError>
where
T: Send + 'static,
Fut: Future<Output = Result<T, RoomError>> + Send + 'static,
{
let (_alive, watch) = tokio::sync::oneshot::channel::<()>();
match tokio::spawn(work(Caller(watch))).await {
Ok(result) => result,
Err(err) if err.is_panic() => std::panic::resume_unwind(err.into_panic()),
Err(err) => Err(map_db_err(err)),
}
}
struct Caller(tokio::sync::oneshot::Receiver<()>);
impl Caller {
fn is_gone(&mut self) -> bool {
matches!(
self.0.try_recv(),
Err(tokio::sync::oneshot::error::TryRecvError::Closed)
)
}
async fn checkout(
&mut self,
pool: &Pool<RuntimeConnection>,
) -> Result<diesel_async::pooled_connection::deadpool::Object<RuntimeConnection>, RoomError>
{
tokio::select! {
conn = pool.get() => conn.map_err(map_db_err),
_ = &mut self.0 => Err(RoomError::Store),
}
}
}
async fn check_seat(
conn: &mut RuntimeConnection,
namespace: &str,
room_id: &str,
participant_id: &str,
token: &str,
) -> Result<(), TxError> {
let room_exists: Option<String> = media_rooms::table
.filter(
media_rooms::namespace
.eq(namespace)
.and(media_rooms::room_id.eq(room_id)),
)
.select(media_rooms::room_id)
.first(conn)
.await
.optional()?;
if room_exists.is_none() {
return Err(RoomError::RoomNotFound.into());
}
let stored: String = media_room_participants::table
.filter(
media_room_participants::namespace
.eq(namespace)
.and(media_room_participants::room_id.eq(room_id))
.and(media_room_participants::participant_id.eq(participant_id)),
)
.select(media_room_participants::token)
.first(conn)
.await
.optional()?
.ok_or(RoomError::ParticipantNotFound)?;
if !autumn_web::auth::constant_time_eq(token.as_bytes(), stored.as_bytes()) {
return Err(RoomError::Unauthorized.into());
}
Ok(())
}
fn snapshot_from(room: &RoomRow, rows: &[ParticipantRow]) -> RoomSnapshot {
let mut participants: Vec<ParticipantView> = rows
.iter()
.map(|row| ParticipantView {
id: row.participant_id.clone(),
display_name: row.display_name.clone(),
joined_at: row.joined_at.and_utc(),
})
.collect();
participants.sort_by(|a, b| a.joined_at.cmp(&b.joined_at).then_with(|| a.id.cmp(&b.id)));
RoomSnapshot {
id: room.room_id.clone(),
namespace: room.namespace.clone(),
max_participants: usize::try_from(room.max_participants).unwrap_or(0),
created_at: room.created_at.and_utc(),
participants,
}
}
#[allow(clippy::needless_lifetimes)] impl RoomStore for DbRoomStore {
fn create_room<'a>(
&'a self,
namespace: &'a str,
max_participants: usize,
) -> RoomStoreFuture<'a, RoomSnapshot> {
Box::pin(async move {
if !namespace.is_empty() {
validate_room_segment(namespace)?;
}
let ceiling = self.hard_cap.min(DEFAULT_ROOM_MAX_PARTICIPANTS);
if max_participants == 0 || max_participants > ceiling {
return Err(RoomError::InvalidMaxParticipants {
requested: max_participants,
cap: ceiling,
});
}
let mut conn = self.pool.get().await.map_err(map_db_err)?;
let count: i64 = media_rooms::table
.count()
.get_result(&mut conn)
.await
.map_err(map_db_err)?;
if usize::try_from(count).unwrap_or(usize::MAX) >= self.max_rooms {
return Err(RoomError::RegistryFull {
max: self.max_rooms,
});
}
let id = Uuid::new_v4().to_string();
let now = Utc::now();
let row = RoomRow {
namespace: namespace.to_owned(),
room_id: id.clone(),
max_participants: i32::try_from(max_participants).unwrap_or(i32::MAX),
created_at: now.naive_utc(),
};
diesel::insert_into(media_rooms::table)
.values(&row)
.execute(&mut conn)
.await
.map_err(map_db_err)?;
Ok(RoomSnapshot {
id,
namespace: namespace.to_owned(),
max_participants,
created_at: now,
participants: Vec::new(),
})
})
}
fn join_room<'a>(
&'a self,
namespace: &'a str,
room_id: &'a str,
display_name: Option<String>,
token_ttl: Duration,
) -> RoomStoreFuture<'a, JoinRecord> {
let pool = self.pool.clone();
let (namespace, room_id) = (namespace.to_owned(), room_id.to_owned());
Box::pin(detached(move |mut caller| async move {
let (namespace, room_id) = (namespace.as_str(), room_id.as_str());
let mut conn = caller.checkout(&pool).await?;
conn.transaction(async move |conn| {
if !lock_room(conn, namespace, room_id).await? {
return Err(RoomError::RoomNotFound.into());
}
let room: RoomRow = media_rooms::table
.filter(
media_rooms::namespace
.eq(namespace)
.and(media_rooms::room_id.eq(room_id)),
)
.select(RoomRow::as_select())
.first(conn)
.await?;
let max = usize::try_from(room.max_participants).unwrap_or(0);
let seats: i64 = media_room_participants::table
.filter(
media_room_participants::namespace
.eq(namespace)
.and(media_room_participants::room_id.eq(room_id)),
)
.count()
.get_result(conn)
.await?;
if usize::try_from(seats).unwrap_or(usize::MAX) >= max {
return Err(RoomError::RoomFull { max }.into());
}
if caller.is_gone() {
return Err(RoomError::Store.into());
}
let now = Utc::now();
let participant_id = Uuid::new_v4().to_string();
let token = SessionToken::generate();
let token_expires_at = now + token_ttl;
let new = ParticipantRow {
namespace: namespace.to_owned(),
room_id: room_id.to_owned(),
participant_id: participant_id.clone(),
display_name,
token: token.expose().to_owned(),
joined_at: now.naive_utc(),
token_expires_at: token_expires_at.naive_utc(),
last_seen_at: now.naive_utc(),
};
diesel::insert_into(media_room_participants::table)
.values(&new)
.execute(conn)
.await?;
let rows: Vec<ParticipantRow> = media_room_participants::table
.filter(
media_room_participants::namespace
.eq(namespace)
.and(media_room_participants::room_id.eq(room_id)),
)
.select(ParticipantRow::as_select())
.load(conn)
.await?;
Ok(JoinRecord {
participant_id,
token,
token_expires_at,
room: snapshot_from(&room, &rows),
})
})
.await
.map_err(TxError::into_room)
}))
}
fn leave_room<'a>(
&'a self,
namespace: &'a str,
room_id: &'a str,
participant_id: &'a str,
token: &'a str,
) -> RoomStoreFuture<'a, ()> {
let pool = self.pool.clone();
let owned = [namespace, room_id, participant_id, token].map(str::to_owned);
Box::pin(detached(move |mut caller| async move {
let [namespace, room_id, participant_id, token] = owned.each_ref().map(String::as_str);
let mut conn = caller.checkout(&pool).await?;
check_seat(&mut conn, namespace, room_id, participant_id, token)
.await
.map_err(TxError::into_room)?;
conn.transaction(async move |conn| {
if !lock_room(conn, namespace, room_id).await? {
return Err(RoomError::RoomNotFound.into());
}
check_seat(conn, namespace, room_id, participant_id, token).await?;
diesel::delete(
media_room_participants::table.filter(
media_room_participants::namespace
.eq(namespace)
.and(media_room_participants::room_id.eq(room_id))
.and(media_room_participants::participant_id.eq(participant_id)),
),
)
.execute(conn)
.await?;
let remaining: i64 = media_room_participants::table
.filter(
media_room_participants::namespace
.eq(namespace)
.and(media_room_participants::room_id.eq(room_id)),
)
.count()
.get_result(conn)
.await?;
if remaining == 0 {
diesel::delete(
media_rooms::table.filter(
media_rooms::namespace
.eq(namespace)
.and(media_rooms::room_id.eq(room_id)),
),
)
.execute(conn)
.await?;
}
Ok(())
})
.await
.map_err(TxError::into_room)
}))
}
fn roster<'a>(
&'a self,
namespace: &'a str,
room_id: &'a str,
auth_token: &'a str,
session_max: Duration,
) -> RoomStoreFuture<'a, RoomSnapshot> {
Box::pin(async move {
let mut conn = self.pool.get().await.map_err(map_db_err)?;
let room: RoomRow = media_rooms::table
.filter(
media_rooms::namespace
.eq(namespace)
.and(media_rooms::room_id.eq(room_id)),
)
.select(RoomRow::as_select())
.first(&mut conn)
.await
.optional()
.map_err(map_db_err)?
.ok_or(RoomError::RoomNotFound)?;
let rows: Vec<ParticipantRow> = media_room_participants::table
.filter(
media_room_participants::namespace
.eq(namespace)
.and(media_room_participants::room_id.eq(room_id)),
)
.select(ParticipantRow::as_select())
.load(&mut conn)
.await
.map_err(map_db_err)?;
let now = Utc::now();
let member_id = rows
.iter()
.find(|row| {
autumn_web::auth::constant_time_eq(auth_token.as_bytes(), row.token.as_bytes())
})
.filter(|row| now < row.joined_at.and_utc() + session_max)
.map(|row| row.participant_id.clone())
.ok_or(RoomError::RoomNotFound)?;
diesel::update(
media_room_participants::table.filter(
media_room_participants::namespace
.eq(namespace)
.and(media_room_participants::room_id.eq(room_id))
.and(media_room_participants::participant_id.eq(&member_id)),
),
)
.set(media_room_participants::last_seen_at.eq(now.naive_utc()))
.execute(&mut conn)
.await
.map_err(map_db_err)?;
Ok(snapshot_from(&room, &rows))
})
}
fn heartbeat<'a>(
&'a self,
namespace: &'a str,
room_id: &'a str,
participant_id: &'a str,
token: &'a str,
token_ttl: Duration,
session_max: Duration,
) -> RoomStoreFuture<'a, DateTime<Utc>> {
Box::pin(async move {
let mut conn = self.pool.get().await.map_err(map_db_err)?;
let (stored, joined_at): (String, chrono::NaiveDateTime) =
media_room_participants::table
.filter(
media_room_participants::namespace
.eq(namespace)
.and(media_room_participants::room_id.eq(room_id))
.and(media_room_participants::participant_id.eq(participant_id)),
)
.select((
media_room_participants::token,
media_room_participants::joined_at,
))
.first(&mut conn)
.await
.optional()
.map_err(map_db_err)?
.ok_or(RoomError::RoomNotFound)?;
if !autumn_web::auth::constant_time_eq(token.as_bytes(), stored.as_bytes()) {
return Err(RoomError::RoomNotFound);
}
let now = Utc::now();
let renewed = renewed_expiry(joined_at.and_utc(), now, token_ttl, session_max)
.ok_or(RoomError::RoomNotFound)?;
let renewed = renewed.trunc_subsecs(6);
let updated = diesel::update(
media_room_participants::table.filter(
media_room_participants::namespace
.eq(namespace)
.and(media_room_participants::room_id.eq(room_id))
.and(media_room_participants::participant_id.eq(participant_id)),
),
)
.set((
media_room_participants::last_seen_at.eq(now.naive_utc()),
media_room_participants::token_expires_at.eq(renewed.naive_utc()),
))
.execute(&mut conn)
.await
.map_err(map_db_err)?;
if updated == 0 {
return Err(RoomError::RoomNotFound);
}
Ok(renewed)
})
}
fn reap_stale(&self, now: DateTime<Utc>, idle_ttl: Duration) -> ReapFuture<'_> {
Box::pin(async move {
let mut stats = ReapStats::default();
let cutoff = (now - idle_ttl).naive_utc();
let mut conn = match self.pool.get().await {
Ok(conn) => conn,
Err(err) => {
tracing::warn!(error = %err, "media rooms: reaper db checkout failed");
return stats;
}
};
match diesel::delete(
media_room_participants::table
.filter(media_room_participants::last_seen_at.lt(cutoff)),
)
.execute(&mut conn)
.await
{
Ok(reaped) => stats.participants_reaped = reaped,
Err(err) => {
tracing::warn!(error = %err, "media rooms: reaper participant sweep failed");
return stats;
}
}
let empty = || {
diesel::dsl::not(diesel::dsl::exists(
media_room_participants::table.filter(
media_room_participants::namespace
.eq(media_rooms::namespace)
.and(media_room_participants::room_id.eq(media_rooms::room_id)),
),
))
};
let swept = conn
.transaction(async |conn| {
diesel::update(
media_rooms::table.filter(
media_rooms::created_at
.lt(cutoff)
.and(media_rooms::max_participants.gt(0))
.and(empty()),
),
)
.set(media_rooms::max_participants.eq(media_rooms::max_participants * -1))
.execute(conn)
.await?;
let reaped = diesel::delete(
media_rooms::table.filter(
media_rooms::created_at
.lt(cutoff)
.and(media_rooms::max_participants.lt(0))
.and(empty()),
),
)
.execute(conn)
.await?;
diesel::update(
media_rooms::table.filter(
media_rooms::created_at
.lt(cutoff)
.and(media_rooms::max_participants.lt(0)),
),
)
.set(media_rooms::max_participants.eq(media_rooms::max_participants * -1))
.execute(conn)
.await?;
Ok::<_, diesel::result::Error>(reaped)
})
.await;
match swept {
Ok(reaped) => stats.rooms_reaped = reaped,
Err(err) => {
tracing::warn!(error = %err, "media rooms: reaper room sweep failed");
return stats;
}
}
stats
})
}
}