lavina/crates/projection-xmpp/src/iq.rs

//! Handling of all client2server iq stanzas

use quick_xml::events::Event;

use lavina_core::room::{RoomId, RoomRegistry};
use proto_xmpp::bind::{BindResponse, Jid, Name, Server};
use proto_xmpp::client::{Iq, IqError, IqErrorType, IqType, Message, MessageType};
use proto_xmpp::disco::{Feature, Identity, InfoQuery, Item, ItemQuery};
use proto_xmpp::mam::{Fin, Set};
use proto_xmpp::roster::RosterQuery;
use proto_xmpp::session::Session;
use proto_xmpp::xml::ToXml;

use crate::proto::IqClientBody;
use crate::XmppConnection;

impl<'a> XmppConnection<'a> {
    pub async fn handle_iq(&self, output: &mut Vec<Event<'static>>, iq: Iq<IqClientBody>) {
        match iq.body {
            IqClientBody::Bind(_) => {
                let req = Iq {
                    from: None,
                    id: iq.id,
                    to: None,
                    r#type: IqType::Result,
                    body: BindResponse(Jid {
                        name: Some(self.user.xmpp_name.clone()),
                        server: Server(self.hostname.clone()),
                        resource: Some(self.user.xmpp_resource.clone()),
                    }),
                };
                req.serialize(output);
            }
            IqClientBody::Session(_) => {
                let req = Iq {
                    from: None,
                    id: iq.id,
                    to: None,
                    r#type: IqType::Result,
                    body: Session,
                };
                req.serialize(output);
            }
            IqClientBody::Roster(_) => {
                let req = Iq {
                    from: None,
                    id: iq.id,
                    to: None,
                    r#type: IqType::Result,
                    body: RosterQuery,
                };
                req.serialize(output);
            }
            IqClientBody::DiscoInfo(info) => {
                let response = self.disco_info(iq.to.as_ref(), &info).await;
                match response {
                    Ok(response) => {
                        let req = Iq {
                            from: iq.to,
                            id: iq.id,
                            to: None,
                            r#type: IqType::Result,
                            body: response,
                        };
                        req.serialize(output);
                    }
                    Err(response) => {
                        let req = Iq {
                            from: iq.to,
                            id: iq.id,
                            to: None,
                            r#type: IqType::Error,
                            body: response,
                        };
                        req.serialize(output);
                    }
                }
            }
            IqClientBody::DiscoItem(item) => {
                let response = self.disco_items(iq.to.as_ref(), &item, self.rooms).await;
                let req = Iq {
                    from: iq.to,
                    id: iq.id,
                    to: None,
                    r#type: IqType::Result,
                    body: response,
                };
                req.serialize(output);
            }
            IqClientBody::MessageArchiveRequest(_) => {
                let response = Iq {
                    from: iq.to,
                    id: iq.id,
                    to: None,
                    r#type: IqType::Result,
                    body: Fin {
                        set: Set { count: Some(0) },
                    },
                };
                response.serialize(output);
            }
            _ => {
                let req = Iq {
                    from: None,
                    id: iq.id,
                    to: None,
                    r#type: IqType::Error,
                    body: IqError {
                        r#type: IqErrorType::Cancel,
                    },
                };
                req.serialize(output);
            }
        }
    }

    async fn disco_info(&self, to: Option<&Jid>, req: &InfoQuery) -> Result<InfoQuery, IqError> {
        let identity;
        let feature;

        match to {
            Some(Jid {
                name: None,
                server,
                resource: None,
            }) if server.0 == self.hostname => {
                identity = vec![Identity {
                    category: "server".into(),
                    name: None,
                    r#type: "im".into(),
                }];
                feature = vec![
                    Feature::new("http://jabber.org/protocol/disco#info"),
                    Feature::new("http://jabber.org/protocol/disco#items"),
                    Feature::new("iq"),
                    Feature::new("presence"),
                ]
            }
            Some(Jid {
                name: None,
                server,
                resource: None,
            }) if server.0 == self.hostname_rooms => {
                identity = vec![Identity {
                    category: "conference".into(),
                    name: Some("Chat rooms".into()),
                    r#type: "text".into(),
                }];
                feature = vec![
                    Feature::new("http://jabber.org/protocol/disco#info"),
                    Feature::new("http://jabber.org/protocol/disco#items"),
                    Feature::new("http://jabber.org/protocol/muc"),
                ]
            }
            Some(Jid {
                name: Some(room_name),
                server,
                resource: None,
            }) if server.0 == self.hostname_rooms => {
                let room_id = RoomId::from(room_name.0.clone()).unwrap();
                let Some(_) = self.rooms.get_room(&room_id).await else {
                    // TODO should return item-not-found
                    // example:
                    // <error type="cancel">
                    // <item-not-found xmlns="urn:ietf:params:xml:ns:xmpp-stanzas"/>
                    // <text xmlns="urn:ietf:params:xml:ns:xmpp-stanzas" xml:lang="en">Conference room does not exist</text>
                    // </error>
                    return Err(IqError {
                        r#type: IqErrorType::Cancel,
                    });
                };
                identity = vec![Identity {
                    category: "conference".into(),
                    name: Some(room_id.into_inner().to_string()),
                    r#type: "text".into(),
                }];
                feature = vec![
                    Feature::new("http://jabber.org/protocol/disco#info"),
                    Feature::new("http://jabber.org/protocol/disco#items"),
                    Feature::new("http://jabber.org/protocol/muc"),
                ]
            }
            _ => {
                identity = vec![];
                feature = vec![];
            }
        };
        Ok(InfoQuery {
            node: None,
            identity,
            feature,
        })
    }

    async fn disco_items(&self, to: Option<&Jid>, req: &ItemQuery, rooms: &RoomRegistry) -> ItemQuery {
        let item = match to {
            Some(Jid {
                name: None,
                server,
                resource: None,
            }) if server.0 == self.hostname => {
                vec![Item {
                    jid: Jid {
                        name: None,
                        server: Server(self.hostname_rooms.clone()),
                        resource: None,
                    },
                    name: None,
                    node: None,
                }]
            }
            Some(Jid {
                name: None,
                server,
                resource: None,
            }) if server.0 == self.hostname_rooms => {
                let room_list = rooms.get_all_rooms().await;
                room_list
                    .into_iter()
                    .map(|room_info| Item {
                        jid: Jid {
                            name: Some(Name(room_info.id.into_inner())),
                            server: Server(self.hostname_rooms.clone()),
                            resource: None,
                        },
                        name: None,
                        node: None,
                    })
                    .collect()
            }
            _ => vec![],
        };
        ItemQuery { item }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{launch, Authenticated, RunningServer, ServerConfig};
    use lavina_core::player::PlayerId;
    use lavina_core::repo::{Storage, StorageConfig};
    use lavina_core::LavinaCore;
    use prometheus::Registry as MetricsRegistry;
    use proto_xmpp::bind::{BindRequest, Resource};
    use quick_xml::{Reader, Writer};
    use std::collections::HashMap;
    use std::io::Cursor;

    struct TestServer {
        metrics: MetricsRegistry,
        storage: Storage,
        core: LavinaCore,
    }

    impl TestServer {
        async fn start() -> anyhow::Result<TestServer> {
            let _ = tracing_subscriber::fmt::try_init();
            let metrics = MetricsRegistry::new();
            let storage = Storage::open(StorageConfig {
                db_path: ":memory:".into(),
            })
            .await?;
            let core = LavinaCore::new(metrics.clone(), storage.clone()).await?;
            Ok(TestServer { metrics, storage, core })
        }

        async fn shutdown(self) -> anyhow::Result<()> {
            self.core.shutdown().await?;
            self.storage.close().await?;
            Ok(())
        }
    }

    #[tokio::test]
    async fn test_handle_iq_bind() {
        let server = TestServer::start().await.unwrap();

        server.storage.create_user("tester").await.unwrap();
        let mut output = Vec::new();
        let player_id = PlayerId::from("tester").unwrap();
        let mut conn = server.core.players.connect_to_player(&player_id).await;
        let user = Authenticated {
            player_id,
            xmpp_name: Name("tester".into()),
            xmpp_resource: Resource("tester".into()),
            xmpp_muc_name: Resource("tester".into()),
        };
        let conn = XmppConnection {
            user: &user,
            user_handle: &mut conn,
            rooms: &server.core.rooms,
            hostname: "localhost".into(),
            hostname_rooms: "rooms.localhost".into(),
        };
        let iq = Iq {
            from: None,
            id: "test".into(),
            to: None,
            r#type: IqType::Result,
            body: IqClientBody::Bind(BindRequest(Resource("whatever".into()))),
        };
        conn.handle_iq(&mut output, iq).await;
        let expected = "<iq type=\"result\" id=\"test\" xmlns=\"jabber:client\"><bind xmlns=\"urn:ietf:params:xml:ns:xmpp-bind\"><jid>tester@localhost/tester</jid></bind></iq>";
        let mut reader = Reader::from_str(expected);
        let mut result_iter = output.into_iter();
        loop {
            let result = result_iter.next();
            let expected = reader.read_event().unwrap();
            match (&result, &expected) {
                (None, Event::Eof) => {
                    break;
                }
                (Some(result), expected) => match (result, expected) {
                    (Event::Text(result), Event::Text(expected)) => {
                        assert_eq!(result.as_ref(), expected.as_ref());
                    }
                    (Event::Start(result), Event::Start(expected)) => {
                        assert_eq!(result.name(), expected.name());
                        let result: HashMap<String, String> =
                            HashMap::from_iter(result.attributes().into_iter().map(|attr| attr.unwrap()).map(|attr| {
                                (
                                    std::str::from_utf8(attr.key.as_ref()).unwrap().to_owned(),
                                    std::str::from_utf8(attr.value.as_ref()).unwrap().to_owned(),
                                )
                            }));
                        let expected: HashMap<String, String> = HashMap::from_iter(
                            expected.attributes().into_iter().map(|attr| attr.unwrap()).map(|attr| {
                                (
                                    std::str::from_utf8(attr.key.as_ref()).unwrap().to_owned(),
                                    std::str::from_utf8(attr.value.as_ref()).unwrap().to_owned(),
                                )
                            }),
                        );
                        assert_eq!(result, expected);
                    }
                    (Event::End(result), Event::End(expected)) => {
                        assert_eq!(result.as_ref(), expected.as_ref());
                    }
                    (Event::Empty(result), Event::Empty(expected)) => {}
                    _ => {
                        panic!("Unexpected result: {:?}, expected: {:?}", result, expected);
                    }
                },
                _ => {
                    panic!("Unexpected result: {:?}, expected: {:?}", result, expected);
                }
            }
        }

        server.shutdown().await.unwrap();
    }
}