muchrooms/src/room.rs

// Copyright (C) 2022-2099 The crate authors.
//
// 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://www.gnu.org/licenses/>.

use crate::component::ComponentTrait;
use crate::error::Error;

use std::collections::BTreeMap;
use std::iter::IntoIterator;

use chrono::{FixedOffset, Utc};
use log::debug;
use xmpp_parsers::{
    date::DateTime,
    delay::Delay,
    message::{Message, MessageType, Subject},
    muc::{
        user::{Affiliation, Item as MucItem, Role, Status as MucStatus},
        MucUser,
    },
    presence::{Presence, Type as PresenceType},
    BareJid, FullJid, Jid,
};

pub type Nick = String;
type Session = FullJid;

#[derive(Debug, PartialEq)]
pub enum BroadcastPresence {
    /// Resource joined the room. It needs to know about all other participants, and other
    /// participants needs to know about it.
    Join,
    /// Resource lost sync. It needs to know about all other participants.
    Resync,
    /// Resource change status (becomes busy, etc.), tell all other participants.
    Update,
    /// Resource leaves. It only needs confirmation that it leaves. Tell all other participants.
    Leave,
}

#[derive(Debug, PartialEq)]
pub struct Room {
    pub jid: BareJid,
    pub occupants: BTreeMap<Nick, Occupant>,
    // TODO: Subject struct.
    // TODO: Store subject lang
    pub subject: Option<(String, Occupant, DateTime)>,
}

impl Room {
    pub fn new(jid: BareJid) -> Self {
        Room {
            jid,
            occupants: BTreeMap::new(),
            subject: None,
        }
    }

    pub async fn broadcast_presence<C: ComponentTrait>(
        &self,
        component: &mut C,
        own_occupant: &Occupant,
        own_session: &Session,
        mode: BroadcastPresence,
    ) -> Result<(), Error> {
        let leave = mode == BroadcastPresence::Leave;

        // All participants to new participant
        let presence_to_new = Presence::new(if leave {
            PresenceType::Unavailable
        } else {
            PresenceType::None
        })
        .with_to(own_session.clone())
        .with_payloads(vec![MucUser {
            status: Vec::new(),
            items: vec![MucItem::new(Affiliation::Owner, Role::Moderator)],
        }
        .into()]);

        // New participant to all other sessions
        let presence_to_old = Presence::new(if leave {
            PresenceType::Unavailable
        } else {
            PresenceType::None
        })
        .with_from(Jid::Full(own_occupant.participant.clone()))
        .with_payloads(vec![MucUser {
            status: Vec::new(),
            items: vec![MucItem::new(
                Affiliation::Owner,
                if leave { Role::None } else { Role::Moderator },
            )],
        }
        .into()]);

        let sync = match mode {
            BroadcastPresence::Join | BroadcastPresence::Resync => true,
            _ => false,
        };
        let update = match mode {
            BroadcastPresence::Join | BroadcastPresence::Update => true,
            _ => false,
        };

        for (_, other) in self.occupants.iter() {
            if own_occupant.nick == other.nick {
                continue;
            }
            if sync {
                // Send presences from others to participant.
                let presence = presence_to_new
                    .clone()
                    .with_from(Jid::Full(other.participant.clone()));
                component.send_stanza(presence).await?;
            }
            if update || leave {
                // Send presence from participant to others.
                for session in other.iter() {
                    // Skip sending if it's us.
                    if session == own_session {
                        continue;
                    }

                    let presence = presence_to_old.clone().with_to(Jid::Full(session.clone()));
                    component.send_stanza(presence).await?;
                }
            }
        }

        // MucItems to be sent to sessions of this occupant
        let self_items = own_occupant
            .iter()
            .map(|session| MucItem {
                affiliation: Affiliation::Owner,
                role: if leave { Role::None } else { Role::Moderator },
                jid: Some(session.clone()),
                nick: None,
                actor: None,
                continue_: None,
                reason: None,
            })
            .collect::<Vec<MucItem>>();

        // Multi-Session Nick: For this occupant, include all sessions all with item@jid discovered
        // so they can identify each other as being the same account under the same nick.
        let session_presence = Presence::new(if leave {
            PresenceType::Unavailable
        } else {
            PresenceType::None
        })
        .with_from(Jid::Full(own_occupant.participant.clone()));

        for session in own_occupant.iter() {
            if session == own_session {
                continue;
            }
            let presence = session_presence
                .clone()
                .with_to(Jid::Full(session.clone()))
                .with_payloads(vec![MucUser {
                    status: vec![],
                    items: self_items.clone(),
                }
                .into()]);
            component.send_stanza(presence).await?;
        }

        // Send self-presence
        if sync || leave {
            // New participant to all other sessions
            let self_presence = Presence::new(if leave {
                PresenceType::Unavailable
            } else {
                PresenceType::None
            })
            .with_from(Jid::Full(own_occupant.participant.clone()))
            .with_to(own_session.clone())
            .with_payloads(vec![MucUser {
                status: if leave {
                    vec![MucStatus::SelfPresence]
                } else {
                    vec![MucStatus::SelfPresence, MucStatus::AssignedNick]
                },
                items: if leave {
                    vec![MucItem::new(Affiliation::Owner, Role::None)]
                } else {
                    self_items
                },
            }
            .into()]);
            component.send_stanza(self_presence).await?;
        }

        Ok(())
    }

    pub async fn send_subject<C: ComponentTrait>(
        &mut self,
        component: &mut C,
        realjid: Session,
        occupant: Occupant,
    ) -> Result<(), Error> {
        debug!("Sending subject!");
        if self.subject.is_none() {
            let subject = String::from("");
            let setter = occupant;
            let stamp = DateTime(Utc::now().with_timezone(&FixedOffset::east(0)));
            self.subject = Some((subject, setter, stamp));
        }

        let mut subject = Message::new(Some(Jid::Full(realjid)));
        subject.from = Some(Jid::Full(
            self.subject.as_ref().unwrap().1.participant.clone(),
        ));
        subject.subjects.insert(
            String::from("en"),
            Subject(self.subject.as_ref().unwrap().0.clone()),
        );
        subject.type_ = MessageType::Groupchat;
        subject.payloads = vec![Delay {
            from: Some(Jid::Bare(self.jid.clone())),
            stamp: self.subject.as_ref().unwrap().2.clone(),
            data: None,
        }
        .into()];
        component.send_stanza(subject).await
    }

    pub async fn add_session<C: ComponentTrait>(
        &mut self,
        component: &mut C,
        realjid: Session,
        new_nick: Nick,
    ) -> Result<(), Error> {
        // Ensure nick isn't already assigned
        let _ = self.occupants.iter().try_for_each(|(nick, occupant)| {
            let new_nick = new_nick.as_str();
            if new_nick == nick && occupant.real != BareJid::from(realjid.clone()) {
                return Err(Error::NickAlreadyAssigned(String::from(new_nick)));
            }
            Ok(())
        })?;

        let mode: Option<BroadcastPresence> = {
            if let Some(occupant) = self.occupants.get_mut(&new_nick) {
                match occupant.add_session(realjid.clone()) {
                    Ok(_) => Some(BroadcastPresence::Join),
                    Err(Error::SessionAlreadyExists(_)) => Some(BroadcastPresence::Resync),
                    Err(err) => return Err(err),
                }
            } else {
                Some(BroadcastPresence::Join)
            }
        };

        if !self.occupants.contains_key(&new_nick) {
            let _ = self.occupants.insert(
                new_nick.clone(),
                Occupant::new(&self, realjid.clone(), new_nick.clone()),
            );
        }
        let occupant = self.occupants.get(&new_nick).unwrap();

        match mode {
            Some(BroadcastPresence::Resync) => {
                self.broadcast_presence(component, &occupant, &realjid, BroadcastPresence::Resync)
                    .await?;
            }
            Some(BroadcastPresence::Join) => {
                debug!("{} is joining {}", realjid, self.jid);

                self.broadcast_presence(component, &occupant, &realjid, BroadcastPresence::Join)
                    .await?;
                self.send_subject(component, realjid, occupant.clone())
                    .await?;
            }
            _ => (),
        }

        Ok(())
    }

    pub async fn remove_session<C: ComponentTrait>(
        &mut self,
        component: &mut C,
        realjid: Session,
        nick: Nick,
    ) -> Result<(), Error> {
        // If occupant doesn't exist, ignore.
        if let Some(mut occupant) = self.occupants.remove(&nick) {
            self.broadcast_presence(component, &occupant, &realjid, BroadcastPresence::Leave)
                .await?;

            occupant.remove_session(realjid)?;
        } else {
            // TODO: Error
        }

        Ok(())
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct Occupant {
    /// Public Jid for the Occupant
    pub real: BareJid,
    pub participant: FullJid,
    pub nick: Nick,
    pub sessions: Vec<FullJid>,
}

impl Occupant {
    fn new(room: &Room, real: FullJid, nick: Nick) -> Occupant {
        Occupant {
            real: BareJid::from(real.clone()),
            participant: room.jid.clone().with_resource(nick.clone()),
            nick,
            sessions: vec![real],
        }
    }

    pub fn add_session(&mut self, real: FullJid) -> Result<(), Error> {
        if BareJid::from(real.clone()) != self.real {
            return Err(Error::MismatchJids(
                Jid::from(self.real.clone()),
                Jid::from(real.clone()),
            ));
        }

        for session in &self.sessions {
            if &real == session {
                return Err(Error::SessionAlreadyExists(real));
            }
        }

        self.sessions.push(real);
        Ok(())
    }

    pub fn remove_session(&mut self, real: FullJid) -> Result<(), Error> {
        if BareJid::from(real.clone()) != self.real {
            return Err(Error::MismatchJids(
                Jid::from(self.real.clone()),
                Jid::from(real.clone()),
            ));
        }

        let len = self.sessions.len();
        self.sessions.retain(|session| session != &real);

        // An item has been removed
        if len != self.sessions.len() {
            Ok(())
        } else {
            Err(Error::NonexistantSession(real))
        }
    }
}

impl IntoIterator for Occupant {
    type Item = FullJid;
    type IntoIter = std::vec::IntoIter<Self::Item>;

    fn into_iter(self) -> Self::IntoIter {
        self.sessions.into_iter()
    }
}

impl Occupant {
    fn iter(&self) -> std::slice::Iter<'_, FullJid> {
        self.sessions.iter()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::component::TestComponent;
    use std::str::FromStr;
    use xmpp_parsers::{
        muc::{
            user::{Affiliation, Item as MucItem, Role, Status as MucStatus},
            MucUser,
        },
        presence::{Presence, Type as PresenceType},
        BareJid, Element,
    };

    #[tokio::test]
    async fn test_broadcast_presence_resync() {
        let roomjid = BareJid::from_str("room@muc").unwrap();
        let realjid1 = FullJid::from_str("foo@bar/qxx").unwrap();
        let participant1 = roomjid.clone().with_resource(String::from("nick1"));
        let realjid2 = FullJid::from_str("qxx@foo/bar").unwrap();
        let participant2 = roomjid.clone().with_resource(String::from("nick2"));
        let realjid3 = FullJid::from_str("bar@qxx/foo").unwrap();
        let participant3 = roomjid.clone().with_resource(String::from("nick3"));

        let mut room = Room::new(roomjid.clone());

        room.occupants = BTreeMap::new();
        room.occupants.insert(
            participant1.resource.clone(),
            Occupant::new(&room, realjid1.clone(), String::from("nick1")),
        );
        room.occupants.insert(
            participant2.resource.clone(),
            Occupant::new(&room, realjid2.clone(), String::from("nick2")),
        );
        let occupant3 = Occupant::new(&room, realjid3.clone(), String::from("nick3"));
        room.occupants
            .insert(participant3.resource.clone(), occupant3.clone());

        let self_item = MucItem::try_from(
            r#"<item xmlns="http://jabber.org/protocol/muc#user" affiliation="owner" role="moderator" jid="bar@qxx/foo"/>"#
                .parse::<Element>()
                .unwrap()
            ).unwrap();

        // BroadcastPresence::Resync
        let mut component = TestComponent::new(vec![]);

        component.expect(
            Presence::new(PresenceType::None)
                .with_from(participant1.clone())
                .with_to(realjid3.clone())
                .with_payloads(vec![MucUser {
                    status: vec![],
                    items: vec![MucItem::new(Affiliation::Owner, Role::Moderator)],
                }
                .into()]),
        );

        component.expect(
            Presence::new(PresenceType::None)
                .with_from(participant2.clone())
                .with_to(realjid3.clone())
                .with_payloads(vec![MucUser {
                    status: vec![],
                    items: vec![MucItem::new(Affiliation::Owner, Role::Moderator)],
                }
                .into()]),
        );

        component.expect(
            Presence::new(PresenceType::None)
                .with_from(participant3.clone())
                .with_to(realjid3.clone())
                .with_payloads(vec![MucUser {
                    status: vec![MucStatus::SelfPresence, MucStatus::AssignedNick],
                    items: vec![self_item.clone()],
                }
                .into()]),
        );

        room.broadcast_presence(
            &mut component,
            &occupant3,
            &realjid3,
            BroadcastPresence::Resync,
        )
        .await
        .unwrap();
    }

    #[tokio::test]
    async fn test_broadcast_presence_update() {
        let roomjid = BareJid::from_str("room@muc").unwrap();
        let realjid1 = FullJid::from_str("foo@bar/qxx").unwrap();
        let participant1 = roomjid.clone().with_resource(String::from("nick1"));
        let realjid2 = FullJid::from_str("qxx@foo/bar").unwrap();
        let participant2 = roomjid.clone().with_resource(String::from("nick2"));
        let realjid3 = FullJid::from_str("bar@qxx/foo").unwrap();
        let participant3 = roomjid.clone().with_resource(String::from("nick3"));

        let mut room = Room::new(roomjid.clone());

        room.occupants = BTreeMap::new();
        room.occupants.insert(
            participant1.resource.clone(),
            Occupant::new(&room, realjid1.clone(), String::from("nick1")),
        );
        room.occupants.insert(
            participant2.resource.clone(),
            Occupant::new(&room, realjid2.clone(), String::from("nick2")),
        );
        let occupant3 = Occupant::new(&room, realjid3.clone(), String::from("nick3"));
        room.occupants
            .insert(participant3.resource.clone(), occupant3.clone());

        // BroadcastPresence::Update
        let mut component = TestComponent::new(vec![]);

        component.expect(
            Presence::new(PresenceType::None)
                .with_from(participant3.clone())
                .with_to(realjid1.clone())
                .with_payloads(vec![MucUser {
                    status: vec![],
                    items: vec![MucItem::new(Affiliation::Owner, Role::Moderator)],
                }
                .into()]),
        );

        component.expect(
            Presence::new(PresenceType::None)
                .with_from(participant3.clone())
                .with_to(realjid2.clone())
                .with_payloads(vec![MucUser {
                    status: vec![],
                    items: vec![MucItem::new(Affiliation::Owner, Role::Moderator)],
                }
                .into()]),
        );

        room.broadcast_presence(
            &mut component,
            &occupant3,
            &realjid3,
            BroadcastPresence::Update,
        )
        .await
        .unwrap();
    }

    #[tokio::test]
    async fn test_broadcast_presence_join() {
        let roomjid = BareJid::from_str("room@muc").unwrap();
        let realjid1 = FullJid::from_str("foo@bar/qxx").unwrap();
        let participant1 = roomjid.clone().with_resource(String::from("nick1"));
        let realjid2 = FullJid::from_str("qxx@foo/bar").unwrap();
        let participant2 = roomjid.clone().with_resource(String::from("nick2"));
        let realjid3 = FullJid::from_str("bar@qxx/foo").unwrap();
        let participant3 = roomjid.clone().with_resource(String::from("nick3"));

        let mut room = Room::new(roomjid.clone());

        room.occupants = BTreeMap::new();
        room.occupants.insert(
            participant1.resource.clone(),
            Occupant::new(&room, realjid1.clone(), String::from("nick1")),
        );
        room.occupants.insert(
            participant2.resource.clone(),
            Occupant::new(&room, realjid2.clone(), String::from("nick2")),
        );
        let occupant3 = Occupant::new(&room, realjid3.clone(), String::from("nick3"));
        room.occupants
            .insert(participant3.resource.clone(), occupant3.clone());

        let self_item = MucItem::try_from(
            r#"<item xmlns="http://jabber.org/protocol/muc#user" affiliation="owner" role="moderator" jid="bar@qxx/foo"/>"#
                .parse::<Element>()
                .unwrap()
            ).unwrap();

        // BroadcastPresence::Join
        let mut component = TestComponent::new(vec![]);

        component.expect(
            Presence::new(PresenceType::None)
                .with_from(participant1.clone())
                .with_to(realjid3.clone())
                .with_payloads(vec![MucUser {
                    status: vec![],
                    items: vec![MucItem::new(Affiliation::Owner, Role::Moderator)],
                }
                .into()]),
        );

        component.expect(
            Presence::new(PresenceType::None)
                .with_from(participant3.clone())
                .with_to(realjid1.clone())
                .with_payloads(vec![MucUser {
                    status: vec![],
                    items: vec![MucItem::new(Affiliation::Owner, Role::Moderator)],
                }
                .into()]),
        );

        component.expect(
            Presence::new(PresenceType::None)
                .with_from(participant2.clone())
                .with_to(realjid3.clone())
                .with_payloads(vec![MucUser {
                    status: vec![],
                    items: vec![MucItem::new(Affiliation::Owner, Role::Moderator)],
                }
                .into()]),
        );

        component.expect(
            Presence::new(PresenceType::None)
                .with_from(participant3.clone())
                .with_to(realjid2.clone())
                .with_payloads(vec![MucUser {
                    status: vec![],
                    items: vec![MucItem::new(Affiliation::Owner, Role::Moderator)],
                }
                .into()]),
        );

        component.expect(
            Presence::new(PresenceType::None)
                .with_from(participant3.clone())
                .with_to(realjid3.clone())
                .with_payloads(vec![MucUser {
                    status: vec![MucStatus::SelfPresence, MucStatus::AssignedNick],
                    items: vec![self_item],
                }
                .into()]),
        );

        room.broadcast_presence(
            &mut component,
            &occupant3,
            &realjid3,
            BroadcastPresence::Join,
        )
        .await
        .unwrap();
    }

    #[tokio::test]
    async fn test_broadcast_presence_leave() {
        let roomjid = BareJid::from_str("room@muc").unwrap();
        let realjid1 = FullJid::from_str("foo@bar/qxx").unwrap();
        let participant1 = roomjid.clone().with_resource(String::from("nick1"));
        let realjid2 = FullJid::from_str("qxx@foo/bar").unwrap();
        let participant2 = roomjid.clone().with_resource(String::from("nick2"));
        let realjid3 = FullJid::from_str("bar@qxx/foo").unwrap();
        let participant3 = roomjid.clone().with_resource(String::from("nick3"));

        let mut room = Room::new(roomjid.clone());

        room.occupants = BTreeMap::new();
        room.occupants.insert(
            participant1.resource.clone(),
            Occupant::new(&room, realjid1.clone(), String::from("nick1")),
        );
        room.occupants.insert(
            participant2.resource.clone(),
            Occupant::new(&room, realjid2.clone(), String::from("nick2")),
        );
        let occupant3 = Occupant::new(&room, realjid3.clone(), String::from("nick3"));
        room.occupants
            .insert(participant3.resource.clone(), occupant3.clone());

        // BroadcastPresence::Leave
        let mut component = TestComponent::new(vec![]);

        component.expect(
            Presence::new(PresenceType::Unavailable)
                .with_from(participant3.clone())
                .with_to(realjid1.clone())
                .with_payloads(vec![MucUser {
                    status: vec![],
                    items: vec![MucItem::new(Affiliation::Owner, Role::None)],
                }
                .into()]),
        );

        component.expect(
            Presence::new(PresenceType::Unavailable)
                .with_from(participant3.clone())
                .with_to(realjid2.clone())
                .with_payloads(vec![MucUser {
                    status: vec![],
                    items: vec![MucItem::new(Affiliation::Owner, Role::None)],
                }
                .into()]),
        );

        component.expect(
            Presence::new(PresenceType::Unavailable)
                .with_from(participant3.clone())
                .with_to(realjid3.clone())
                .with_payloads(vec![MucUser {
                    status: vec![MucStatus::SelfPresence],
                    items: vec![MucItem::new(Affiliation::Owner, Role::None)],
                }
                .into()]),
        );

        room.broadcast_presence(
            &mut component,
            &occupant3,
            &realjid3,
            BroadcastPresence::Leave,
        )
        .await
        .unwrap();
    }
}