lavina/crates/proto-xmpp/src/bind.rs

use std::fmt::Display;

use anyhow::{anyhow, Result};
use quick_xml::events::{BytesEnd, BytesStart, BytesText, Event};
use quick_xml::name::{Namespace, ResolveResult};

use crate::prelude::*;
use crate::xml::*;

pub const XMLNS: &'static str = "urn:ietf:params:xml:ns:xmpp-bind";

// TODO remove `pub` in newtypes, introduce validation

/// Name (node identifier) of an XMPP entity. Placed before the `@` in a JID.
#[derive(PartialEq, Eq, Debug, Clone)]
pub struct Name(pub Str);

/// Server name of an XMPP entity. Placed after the `@` and before the `/` in a JID.
#[derive(PartialEq, Eq, Debug, Clone)]
pub struct Server(pub Str);

/// Resource of an XMPP entity. Placed after the `/` in a JID.
#[derive(PartialEq, Eq, Debug, Clone)]
pub struct Resource(pub Str);

#[derive(PartialEq, Eq, Debug, Clone)]
pub struct Jid {
    pub name: Option<Name>,
    pub server: Server,
    pub resource: Option<Resource>,
}

impl Display for Jid {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if let Some(name) = &self.name {
            write!(f, "{}@", &name.0)?;
        }
        write!(f, "{}", &self.server.0)?;
        if let Some(resource) = &self.resource {
            write!(f, "/{}", &resource.0)?;
        }
        Ok(())
    }
}

impl Jid {
    pub fn from_string(i: &str) -> Result<Jid> {
        use lazy_static::lazy_static;
        use regex::Regex;
        lazy_static! {
            static ref RE: Regex = Regex::new(r"^(([a-zA-Z]+)@)?([a-zA-Z.]+)(/([a-zA-Z\-]+))?$").unwrap();
        }
        let m = RE.captures(i).ok_or(anyhow!("Incorrectly format jid: {i}"))?;

        let name = m.get(2).map(|name| Name(name.as_str().into()));
        let server = m.get(3).unwrap();
        let server = Server(server.as_str().into());
        let resource = m.get(5).map(|resource| Resource(resource.as_str().into()));

        Ok(Jid { name, server, resource })
    }
}

/// Request to bind to a resource.
///
/// Example:
/// ```xml
/// <bind xmlns="urn:ietf:params:xml:ns:xmpp-bind">
///   <resource>mobile</resource>
/// </bind>
/// ```
///
#[derive(PartialEq, Eq, Debug)]
pub struct BindRequest(pub Resource);

impl FromXmlTag for BindRequest {
    const NS: &'static str = XMLNS;
    const NAME: &'static str = "bind";
}

impl FromXml for BindRequest {
    type P = impl Parser<Output = Result<Self>>;

    fn parse() -> Self::P {
        |(namespace, event): (ResolveResult<'static>, &'static Event<'static>)| -> Result<Self> {
            let mut resource: Option<Str> = None;
            let Event::Start(bytes) = event else {
                return Err(anyhow!("Unexpected XML event: {event:?}"));
            };
            if bytes.name().0 != BindRequest::NAME.as_bytes() {
                return Err(anyhow!("Unexpected XML tag: {:?}", bytes.name()));
            }
            let ResolveResult::Bound(Namespace(ns)) = namespace else {
                return Err(anyhow!("No namespace provided"));
            };
            if ns != XMLNS.as_bytes() {
                return Err(anyhow!("Incorrect namespace"));
            }
            loop {
                let (namespace, event) = yield;
                match event {
                    Event::Start(bytes) if bytes.name().0 == b"resource" => {
                        let (namespace, event) = yield;
                        let Event::Text(text) = event else {
                            return Err(anyhow!("Unexpected XML event: {event:?}"));
                        };
                        resource = Some(std::str::from_utf8(&*text)?.into());
                        let (namespace, event) = yield;
                        let Event::End(bytes) = event else {
                            return Err(anyhow!("Unexpected XML event: {event:?}"));
                        };
                        if bytes.name().0 != b"resource" {
                            return Err(anyhow!("Unexpected XML tag: {:?}", bytes.name()));
                        }
                    }
                    Event::End(bytes) if bytes.name().0 == BindRequest::NAME.as_bytes() => {
                        break;
                    }
                    _ => return Err(anyhow!("Unexpected XML event: {event:?}")),
                }
            }
            let Some(resource) = resource else {
                return Err(anyhow!("No resource was provided"));
            };
            Ok(BindRequest(Resource(resource)))
        }
    }
}

pub struct BindResponse(pub Jid);

impl ToXml for BindResponse {
    fn serialize(&self, events: &mut Vec<Event<'static>>) {
        events.extend_from_slice(&[
            Event::Start(BytesStart::new(r#"bind xmlns="urn:ietf:params:xml:ns:xmpp-bind""#)),
            Event::Start(BytesStart::new(r#"jid"#)),
            Event::Text(BytesText::new(self.0.to_string().as_str()).into_owned()),
            Event::End(BytesEnd::new("jid")),
            Event::End(BytesEnd::new("bind")),
        ]);
    }
}

#[cfg(test)]
mod tests {
    use quick_xml::NsReader;

    use super::*;

    #[tokio::test]
    async fn parse_message() {
        let input = r#"<bind xmlns="urn:ietf:params:xml:ns:xmpp-bind"><resource>mobile</resource></bind>"#;
        let mut reader = NsReader::from_reader(input.as_bytes());
        let mut buf = vec![];
        let (ns, event) = reader.read_resolved_event_into_async(&mut buf).await.unwrap();
        let mut parser = BindRequest::parse().consume(ns, &event);
        let result = loop {
            match parser {
                Continuation::Final(res) => break res,
                Continuation::Continue(next) => {
                    let (ns, event) = reader.read_resolved_event_into_async(&mut buf).await.unwrap();
                    parser = next.consume(ns, &event);
                }
            }
        }
        .unwrap();
        assert_eq!(result, BindRequest(Resource("mobile".into())),)
    }

    #[test]
    fn jid_parse_full() {
        let input = "chelik@server.example/kek";
        let expected = Jid {
            name: Some(Name("chelik".into())),
            server: Server("server.example".into()),
            resource: Some(Resource("kek".into())),
        };
        let res = Jid::from_string(input).unwrap();
        assert_eq!(res, expected);
    }

    #[test]
    fn jid_parse_user() {
        let input = "chelik@server.example";
        let expected = Jid {
            name: Some(Name("chelik".into())),
            server: Server("server.example".into()),
            resource: None,
        };
        let res = Jid::from_string(input).unwrap();
        assert_eq!(res, expected);
    }

    #[test]
    fn jid_parse_server() {
        let input = "server.example";
        let expected = Jid {
            name: None,
            server: Server("server.example".into()),
            resource: None,
        };
        let res = Jid::from_string(input).unwrap();
        assert_eq!(res, expected);
    }

    #[test]
    fn jid_parse_server_resource() {
        let input = "server.example/kek";
        let expected = Jid {
            name: None,
            server: Server("server.example".into()),
            resource: Some(Resource("kek".into())),
        };
        let res = Jid::from_string(input).unwrap();
        assert_eq!(res, expected);
    }
}