use std::collections::{BTreeSet, HashMap};
use mail4agent_api::{Ack, Address, Message, MessageId, ParticipantId, RoomId, SessionCard, SessionId};
use thiserror::Error;
pub type SecretDigest = [u8; 32];
#[derive(Debug, Error)]
#[error("{0}")]
pub struct StoreError(pub String);
impl StoreError {
pub fn new(message: impl Into<String>) -> Self {
Self(message.into())
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ParticipantRecord {
pub label: Option<String>,
pub secret_digest: SecretDigest,
pub may_send: bool,
pub may_read: bool,
pub operator: bool,
pub listener_url: Option<String>,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct RoomRecord {
pub created_at_unix_ms: u64,
pub members: BTreeSet<ParticipantId>,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ParticipantSummary {
pub id: ParticipantId,
pub label: Option<String>,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct RoomSummary {
pub id: RoomId,
pub members: BTreeSet<ParticipantId>,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct SessionRecord {
pub account: ParticipantId,
pub card: SessionCard,
pub last_seen_unix_ms: u64,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum InsertMessageOutcome {
Inserted,
Deduplicated { message_id: MessageId },
}
pub trait MailStore {
fn register_participant(&mut self, id: ParticipantId, record: ParticipantRecord) -> Result<(), StoreError>;
fn deregister_participant(&mut self, id: &ParticipantId) -> Result<(), StoreError>;
fn set_participant_secret_digest(&mut self, id: &ParticipantId, digest: SecretDigest) -> Result<(), StoreError>;
fn set_listener_url(&mut self, id: &ParticipantId, url: Option<String>) -> Result<(), StoreError>;
fn get_participant(&self, id: &ParticipantId) -> Result<Option<ParticipantRecord>, StoreError>;
fn find_participant_by_digest(
&self,
digest: &SecretDigest,
) -> Result<Option<(ParticipantId, ParticipantRecord)>, StoreError>;
fn create_room(&mut self, id: RoomId, created_at_unix_ms: u64) -> Result<(), StoreError>;
fn add_room_member(&mut self, room: &RoomId, participant: ParticipantId) -> Result<(), StoreError>;
fn remove_room_member(&mut self, room: &RoomId, participant: &ParticipantId) -> Result<(), StoreError>;
fn get_room(&self, id: &RoomId) -> Result<Option<RoomRecord>, StoreError>;
fn insert_message(
&mut self,
message: Message,
idempotency: Option<(Address, String)>,
) -> Result<InsertMessageOutcome, StoreError>;
fn get_message(&self, id: &MessageId) -> Result<Option<Message>, StoreError>;
fn messages_to_since(&self, to: &Address, since_unix_ms: u64) -> Result<Vec<Message>, StoreError>;
fn room_messages_since(&self, room: &RoomId, since_unix_ms: u64) -> Result<Vec<Message>, StoreError>;
fn rooms_containing(&self, participant: &ParticipantId) -> Result<Vec<RoomId>, StoreError>;
fn record_ack(&mut self, ack: Ack) -> Result<Ack, StoreError>;
fn get_ack(&self, message_id: &MessageId, reader: &Address) -> Result<Option<Ack>, StoreError>;
fn list_participants(&self) -> Result<Vec<ParticipantSummary>, StoreError>;
fn list_rooms(&self) -> Result<Vec<RoomSummary>, StoreError>;
fn get_session(&self, session: &SessionId) -> Result<Option<SessionRecord>, StoreError>;
fn upsert_session(&mut self, session: SessionId, record: SessionRecord) -> Result<(), StoreError>;
fn sessions_of(&self, account: &ParticipantId) -> Result<Vec<(SessionId, SessionRecord)>, StoreError>;
}
#[derive(Default)]
pub struct InMemoryStore {
participants: HashMap<ParticipantId, ParticipantRecord>,
digest_index: HashMap<SecretDigest, ParticipantId>,
rooms: HashMap<RoomId, RoomRecord>,
messages: HashMap<MessageId, Message>,
idempotency: HashMap<(Address, String), MessageId>,
acks: HashMap<(MessageId, Address), Ack>,
sessions: HashMap<SessionId, SessionRecord>,
}
impl MailStore for InMemoryStore {
fn register_participant(&mut self, id: ParticipantId, record: ParticipantRecord) -> Result<(), StoreError> {
self.digest_index.insert(record.secret_digest, id.clone());
self.participants.insert(id, record);
Ok(())
}
fn deregister_participant(&mut self, id: &ParticipantId) -> Result<(), StoreError> {
if let Some(record) = self.participants.remove(id) {
self.digest_index.remove(&record.secret_digest);
}
Ok(())
}
fn set_participant_secret_digest(&mut self, id: &ParticipantId, digest: SecretDigest) -> Result<(), StoreError> {
if let Some(record) = self.participants.get_mut(id) {
self.digest_index.remove(&record.secret_digest);
record.secret_digest = digest;
self.digest_index.insert(digest, id.clone());
}
Ok(())
}
fn get_participant(&self, id: &ParticipantId) -> Result<Option<ParticipantRecord>, StoreError> {
Ok(self.participants.get(id).cloned())
}
fn set_listener_url(&mut self, id: &ParticipantId, url: Option<String>) -> Result<(), StoreError> {
if let Some(record) = self.participants.get_mut(id) {
record.listener_url = url;
}
Ok(())
}
fn find_participant_by_digest(
&self,
digest: &SecretDigest,
) -> Result<Option<(ParticipantId, ParticipantRecord)>, StoreError> {
let Some(id) = self.digest_index.get(digest) else {
return Ok(None);
};
Ok(self.participants.get(id).map(|record| (id.clone(), record.clone())))
}
fn create_room(&mut self, id: RoomId, created_at_unix_ms: u64) -> Result<(), StoreError> {
self.rooms.insert(id, RoomRecord { created_at_unix_ms, members: BTreeSet::new() });
Ok(())
}
fn add_room_member(&mut self, room: &RoomId, participant: ParticipantId) -> Result<(), StoreError> {
if let Some(record) = self.rooms.get_mut(room) {
record.members.insert(participant);
}
Ok(())
}
fn remove_room_member(&mut self, room: &RoomId, participant: &ParticipantId) -> Result<(), StoreError> {
if let Some(record) = self.rooms.get_mut(room) {
record.members.remove(participant);
}
Ok(())
}
fn get_room(&self, id: &RoomId) -> Result<Option<RoomRecord>, StoreError> {
Ok(self.rooms.get(id).cloned())
}
fn insert_message(
&mut self,
message: Message,
idempotency: Option<(Address, String)>,
) -> Result<InsertMessageOutcome, StoreError> {
if let Some(key) = &idempotency {
if let Some(existing) = self.idempotency.get(key) {
return Ok(InsertMessageOutcome::Deduplicated { message_id: existing.clone() });
}
}
let message_id = message.message_id.clone();
self.messages.insert(message_id.clone(), message);
if let Some(key) = idempotency {
self.idempotency.insert(key, message_id);
}
Ok(InsertMessageOutcome::Inserted)
}
fn get_message(&self, id: &MessageId) -> Result<Option<Message>, StoreError> {
Ok(self.messages.get(id).cloned())
}
fn messages_to_since(&self, to: &Address, since_unix_ms: u64) -> Result<Vec<Message>, StoreError> {
Ok(self
.messages
.values()
.filter(|message| message.created_at_unix_ms >= since_unix_ms && &message.to == to)
.cloned()
.collect())
}
fn room_messages_since(&self, room: &RoomId, since_unix_ms: u64) -> Result<Vec<Message>, StoreError> {
Ok(self
.messages
.values()
.filter(|message| {
message.created_at_unix_ms >= since_unix_ms
&& matches!(&message.to, Address::Room { room: r } if r == room)
})
.cloned()
.collect())
}
fn rooms_containing(&self, participant: &ParticipantId) -> Result<Vec<RoomId>, StoreError> {
Ok(self
.rooms
.iter()
.filter(|(_, record)| record.members.contains(participant))
.map(|(id, _)| id.clone())
.collect())
}
fn record_ack(&mut self, ack: Ack) -> Result<Ack, StoreError> {
Ok(self.acks.entry((ack.message_id.clone(), ack.reader.clone())).or_insert(ack).clone())
}
fn get_ack(&self, message_id: &MessageId, reader: &Address) -> Result<Option<Ack>, StoreError> {
Ok(self.acks.get(&(message_id.clone(), reader.clone())).cloned())
}
fn list_participants(&self) -> Result<Vec<ParticipantSummary>, StoreError> {
Ok(self
.participants
.iter()
.map(|(id, record)| ParticipantSummary { id: id.clone(), label: record.label.clone() })
.collect())
}
fn list_rooms(&self) -> Result<Vec<RoomSummary>, StoreError> {
Ok(self
.rooms
.iter()
.map(|(id, record)| RoomSummary { id: id.clone(), members: record.members.clone() })
.collect())
}
fn get_session(&self, session: &SessionId) -> Result<Option<SessionRecord>, StoreError> {
Ok(self.sessions.get(session).cloned())
}
fn upsert_session(&mut self, session: SessionId, record: SessionRecord) -> Result<(), StoreError> {
self.sessions.insert(session, record);
Ok(())
}
fn sessions_of(&self, account: &ParticipantId) -> Result<Vec<(SessionId, SessionRecord)>, StoreError> {
Ok(self
.sessions
.iter()
.filter(|(_, record)| &record.account == account)
.map(|(id, record)| (id.clone(), record.clone()))
.collect())
}
}