Files
Stalwart/crates/jmap/src/services/state.rs

469 lines
19 KiB
Rust

/*
* Copyright (c) 2023 Stalwart Labs Ltd.
*
* This file is part of Stalwart Mail Server.
*
* 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.
* in the LICENSE file at the top-level directory of this distribution.
* You should have received a copy of the GNU Affero General Public License
* along with this program. If not, see <http://www.gnu.org/licenses/>.
*
* You can be released from the requirements of the AGPLv3 license by
* purchasing a commercial license. Please contact licensing@stalw.art
* for more details.
*/
use std::{
sync::Arc,
time::{Duration, Instant, SystemTime},
};
use jmap_proto::types::{id::Id, state::StateChange, type_state::DataType};
use store::ahash::AHashMap;
use tokio::sync::mpsc;
use utils::{config::Config, map::bitmap::Bitmap};
use crate::{
push::{manager::spawn_push_manager, UpdateSubscription},
JMAP,
};
use super::IPC_CHANNEL_BUFFER;
#[derive(Debug)]
pub enum Event {
Subscribe {
account_id: u32,
types: Bitmap<DataType>,
tx: mpsc::Sender<StateChange>,
},
Publish {
state_change: StateChange,
},
UpdateSharedAccounts {
account_id: u32,
},
UpdateSubscriptions {
account_id: u32,
subscriptions: Vec<UpdateSubscription>,
},
Stop,
}
#[derive(Debug)]
struct Subscriber {
types: Bitmap<DataType>,
subscription: SubscriberType,
}
#[derive(Debug)]
pub enum SubscriberType {
Ipc { tx: mpsc::Sender<StateChange> },
Push { expires: u64 },
}
impl Subscriber {
fn is_valid(&self, current_time: u64) -> bool {
match &self.subscription {
SubscriberType::Ipc { tx } => !tx.is_closed(),
SubscriberType::Push { expires } => expires > &current_time,
}
}
}
const PURGE_EVERY: Duration = Duration::from_secs(3600);
const SEND_TIMEOUT: Duration = Duration::from_millis(500);
pub fn init_state_manager() -> (mpsc::Sender<Event>, mpsc::Receiver<Event>) {
mpsc::channel::<Event>(IPC_CHANNEL_BUFFER)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum SubscriberId {
Ipc(u32),
Push(u32),
}
#[allow(clippy::unwrap_or_default)]
pub fn spawn_state_manager(
core: Arc<JMAP>,
settings: &Config,
mut change_rx: mpsc::Receiver<Event>,
) {
let push_tx = spawn_push_manager(settings);
tokio::spawn(async move {
let mut subscribers: AHashMap<u32, AHashMap<SubscriberId, Subscriber>> =
AHashMap::default();
let mut shared_accounts: AHashMap<u32, Vec<u32>> = AHashMap::default();
let mut shared_accounts_map: AHashMap<u32, AHashMap<u32, Bitmap<DataType>>> =
AHashMap::default();
let mut last_purge = Instant::now();
while let Some(event) = change_rx.recv().await {
let mut purge_needed = last_purge.elapsed() >= PURGE_EVERY;
match event {
Event::Stop => {
if let Err(err) = push_tx.send(crate::push::Event::Reset).await {
tracing::debug!("Error sending push reset: {}", err);
}
break;
}
Event::UpdateSharedAccounts { account_id } => {
// Obtain account membership and shared mailboxes
let acl = match core.get_access_token(account_id).await {
Some(result) => result,
None => {
continue;
}
};
// Delete any removed sharings
if let Some(shared_account_ids) = shared_accounts.get(&account_id) {
for shared_account_id in shared_account_ids {
if *shared_account_id != acl.primary_id
&& !acl.member_of.contains(shared_account_id)
&& !acl
.access_to
.iter()
.any(|(id, _)| *id == *shared_account_id)
{
if let Some(shared_list) =
shared_accounts_map.get_mut(shared_account_id)
{
shared_list.remove(&account_id);
if shared_list.is_empty() {
shared_accounts_map.remove(shared_account_id);
}
}
}
}
}
// Update lists
let mut shared_account_ids =
Vec::with_capacity(acl.member_of.len() + 1 + acl.access_to.len());
for member_id in [acl.primary_id].iter().chain(acl.member_of.iter()) {
shared_account_ids.push(*member_id);
shared_accounts_map
.entry(*member_id)
.or_insert_with(AHashMap::new)
.insert(account_id, Bitmap::all());
}
for (shared_account_id, shared_collections) in acl.access_to.iter() {
let mut types: Bitmap<DataType> = Bitmap::new();
for collection in *shared_collections {
if let Ok(type_state) = DataType::try_from(collection) {
types.insert(type_state);
if type_state == DataType::Email {
types.insert(DataType::EmailDelivery);
types.insert(DataType::Thread);
}
}
}
if !types.is_empty() {
shared_account_ids.push(*shared_account_id);
shared_accounts_map
.entry(*shared_account_id)
.or_insert_with(AHashMap::new)
.insert(account_id, types);
}
}
shared_accounts.insert(account_id, shared_account_ids);
}
Event::Subscribe {
account_id,
types,
tx,
} => {
subscribers
.entry(account_id)
.or_insert_with(AHashMap::default)
.insert(
SubscriberId::Ipc(rand::random()),
Subscriber {
types,
subscription: SubscriberType::Ipc { tx },
},
);
}
Event::Publish { state_change } => {
if let Some(shared_accounts) = shared_accounts_map.get(&state_change.account_id)
{
let current_time = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0);
let mut push_ids = Vec::new();
for (owner_account_id, allowed_types) in shared_accounts {
if let Some(subscribers) = subscribers.get(owner_account_id) {
for (subscriber_id, subscriber) in subscribers {
let mut types = Vec::with_capacity(state_change.types.len());
for (state_type, change_id) in &state_change.types {
if subscriber.types.contains(*state_type)
&& allowed_types.contains(*state_type)
{
types.push((*state_type, *change_id));
}
}
if !types.is_empty() {
match &subscriber.subscription {
SubscriberType::Ipc { tx } if !tx.is_closed() => {
let subscriber_tx = tx.clone();
let state_change = state_change.clone();
tokio::spawn(async move {
// Timeout after 500ms in case there is a blocked client
if let Err(err) = subscriber_tx
.send_timeout(
StateChange {
account_id: state_change.account_id,
types,
},
SEND_TIMEOUT,
)
.await
{
tracing::debug!(
"Error sending state change to subscriber: {}",
err
);
}
});
}
SubscriberType::Push { expires }
if expires > &current_time =>
{
push_ids.push(Id::from_parts(
*owner_account_id,
(*subscriber_id).into(),
));
}
_ => {
purge_needed = true;
}
}
}
}
}
}
if !push_ids.is_empty() {
if let Err(err) = push_tx
.send(crate::push::Event::Push {
ids: push_ids,
state_change,
})
.await
{
tracing::debug!("Error sending push updates: {}", err);
}
}
}
}
Event::UpdateSubscriptions {
account_id,
subscriptions,
} => {
let mut updated_ids = Vec::with_capacity(subscriptions.len());
let mut push_updates = Vec::with_capacity(subscriptions.len());
if let Some(subscribers) = subscribers.get_mut(&account_id) {
let mut remove_ids = Vec::new();
for subscriber_id in subscribers.keys() {
if let SubscriberId::Push(push_id) = subscriber_id {
if !subscriptions.iter().any(|s| {
matches!(s, UpdateSubscription::Verified(
crate::push::PushSubscription { id, .. }
) if id == push_id)
}) {
remove_ids.push(*subscriber_id);
}
}
}
for remove_id in remove_ids {
push_updates.push(crate::push::PushUpdate::Unregister {
id: Id::from_parts(account_id, remove_id.into()),
});
subscribers.remove(&remove_id);
}
}
for subscription in subscriptions {
match subscription {
UpdateSubscription::Unverified {
id,
url,
code,
keys,
} => {
push_updates.push(crate::push::PushUpdate::Verify {
id,
account_id,
url,
code,
keys,
});
}
UpdateSubscription::Verified(verified) => {
updated_ids.push(verified.id);
subscribers
.entry(account_id)
.or_insert_with(AHashMap::default)
.insert(
SubscriberId::Push(verified.id),
Subscriber {
types: verified.types,
subscription: SubscriberType::Push {
expires: verified.expires,
},
},
);
push_updates.push(crate::push::PushUpdate::Register {
id: Id::from_parts(account_id, verified.id),
url: verified.url,
keys: verified.keys,
});
}
}
}
if !push_updates.is_empty() {
if let Err(err) = push_tx
.send(crate::push::Event::Update {
updates: push_updates,
})
.await
{
tracing::debug!("Error sending push updates: {}", err);
}
}
}
}
if purge_needed {
let mut remove_account_ids = Vec::new();
let current_time = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0);
for (account_id, subscriber_map) in &mut subscribers {
let mut remove_subscription_ids = Vec::new();
for (id, subscriber) in subscriber_map.iter() {
if !subscriber.is_valid(current_time) {
remove_subscription_ids.push(*id);
}
}
if !remove_subscription_ids.is_empty() {
if remove_subscription_ids.len() < subscriber_map.len() {
for remove_subscription_id in remove_subscription_ids {
subscriber_map.remove(&remove_subscription_id);
}
} else {
remove_account_ids.push(*account_id);
}
}
}
for remove_account_id in remove_account_ids {
subscribers.remove(&remove_account_id);
}
last_purge = Instant::now();
}
}
});
}
impl JMAP {
pub async fn subscribe_state_manager(
&self,
account_id: u32,
types: Bitmap<DataType>,
) -> Option<mpsc::Receiver<StateChange>> {
let (change_tx, change_rx) = mpsc::channel::<StateChange>(IPC_CHANNEL_BUFFER);
let state_tx = self.state_tx.clone();
for event in [
Event::UpdateSharedAccounts { account_id },
Event::Subscribe {
account_id,
types,
tx: change_tx,
},
] {
if let Err(err) = state_tx.send(event).await {
tracing::error!(
"Channel failure while subscribing to state manager: {}",
err
);
return None;
}
}
change_rx.into()
}
pub async fn broadcast_state_change(&self, state_change: StateChange) -> bool {
match self
.state_tx
.clone()
.send(Event::Publish { state_change })
.await
{
Ok(_) => true,
Err(err) => {
tracing::error!("Channel failure while publishing state change: {}", err);
false
}
}
}
pub async fn update_push_subscriptions(&self, account_id: u32) -> bool {
let push_subs = match self.fetch_push_subscriptions(account_id).await {
Ok(push_subs) => push_subs,
Err(err) => {
tracing::error!(context = "update_push_subscriptions",
event = "error",
reason = %err,
"Error fetching push subscriptions.");
return false;
}
};
let state_tx = self.state_tx.clone();
for event in [Event::UpdateSharedAccounts { account_id }, push_subs] {
if let Err(err) = state_tx.send(event).await {
tracing::error!("Channel failure while publishing state change: {}", err);
return false;
}
}
true
}
}
impl From<SubscriberId> for u32 {
fn from(subscriber_id: SubscriberId) -> u32 {
match subscriber_id {
SubscriberId::Ipc(id) => id,
SubscriberId::Push(id) => id,
}
}
}