path: root/src/stanza/mod.rs

// use quick_xml::events::BytesDecl;

pub mod bind;
pub mod iq;
pub mod sasl;
pub mod stream;

use std::collections::BTreeMap;
use std::str;

// const DECLARATION: BytesDecl<'_> = BytesDecl::new("1.0", None, None);
use async_recursion::async_recursion;
use quick_xml::events::{BytesEnd, BytesStart, BytesText, Event};
use quick_xml::name::PrefixDeclaration;
use quick_xml::{Reader, Writer};
use tokio::io::{AsyncBufRead, AsyncWrite};

use crate::{JabberError, Result};

// #[derive(Clone, Debug)]
// pub struct EventTree<'e> {
//     pub event: Event<'e>,
//     pub children: Option<Vec<Element<'e>>>,
// }

pub type Prefix<'s> = Option<&'s str>;

#[derive(Clone, Debug)]
/// represents an xml element as a tree of nodes
pub struct Element<'s> {
    /// element prefix
    /// e.g. `foo` in `<foo:bar />`.
    prefix: Option<&'s str>,
    /// element name
    /// e.g. `bar` in `<foo:bar />`.
    localname: &'s str,
    /// qualifying namespace
    /// an element must be qualified by a namespace
    /// e.g. for `<stream:features>` in
    /// ```
    /// <stream:stream xmlns='jabber:client' xmlns:stream='http://etherx.jabber.org/streams'>
    ///     <stream:features>
    ///         <bind xmlns='urn:ietf:params:xml:ns:xmpp-bind'/>
    ///         <compression xmlns='http://jabber.org/features/compress'>
    ///             <method>zlib</method>
    ///             <method>lzw</method>
    ///         </compression>
    ///     </stream:features>
    /// </stream:stream>
    /// ```
    /// would be `"http://etherx.jabber.org/streams"` but for
    /// ```
    /// <stream:stream xmlns='jabber:client' xmlns:stream='http://etherx.jabber.org/streams'>
    ///     <features>
    ///         <bind xmlns='urn:ietf:params:xml:ns:xmpp-bind'/>
    ///         <compression xmlns='http://jabber.org/features/compress'>
    ///             <method>zlib</method>
    ///             <method>lzw</method>
    ///         </compression>
    ///     </features>
    /// </stream:stream>
    /// ```
    /// would be `"jabber:client"`
    namespace: &'s str,
    /// all namespaces applied to element
    /// e.g. for `<bind>` in
    /// ```
    /// <stream:features xmlns='jabber:client' xmlns:stream='http://etherx.jabber.org/streams'>
    ///     <bind xmlns='urn:ietf:params:xml:ns:xmpp-bind'/>
    ///     <compression xmlns='http://jabber.org/features/compress'>
    ///         <method>zlib</method>
    ///         <method>lzw</method>
    ///     </compression>
    /// </stream:features>
    /// ```
    /// would be `[(None, "urn:ietf:params:xml:ns:xmpp-bind")]` despite
    /// `(Some("stream"), "http://etherx.jabber.org/streams")` also being available
    namespace_declarations: Box<BTreeMap<Option<&'s str>, &'s str>>,
    /// element attributes
    attributes: Box<BTreeMap<&'s str, &'s str>>,
    // children elements namespaces contain their parents' namespaces
    children: Box<Vec<Node<'s>>>,
}

#[derive(Clone, Debug)]
pub enum Node<'s> {
    Element(Element<'s>),
    Text(&'s str),
    Unknown,
}

impl<'s> From<&Node<'s>> for Vec<Event<'s>> {
    fn from(node: &Node<'s>) -> Self {
        match node {
            Node::Element(e) => e.into(),
            Node::Text(t) => vec![Event::Text(BytesText::new(t))],
            Unknown => vec![],
        }
    }
}

impl<'s> Element<'s> {
    /// returns the fully qualified name
    /// e.g. `foo:bar` in
    /// `<foo:bar>`.
    pub fn name(&self) -> &str {
        if let Some(prefix) = self.prefix {
            format!("{}:{}", prefix, self.localname).as_str()
        } else {
            self.localname
        }
    }

    /// returns the localname.
    /// e.g. `bar` in `<foo:bar>`
    pub fn localname(&self) -> &str {
        self.localname
    }

    /// returns the prefix.
    /// e.g. `foo` in `<foo:bar>`. returns None if there is
    /// no prefix.
    pub fn prefix(&self) -> Option<&str> {
        self.prefix
    }

    /// returns the namespace which applies to the current element, e.g. for
    /// `<element xmlns='foo' xmlns:bar='bar'>`
    /// it will be `foo` but for
    /// `<bar:element xmlns='foo' xmlns:bar='bar'>`
    /// it will be `bar`.
    pub fn namespace(&self) -> &str {
        self.namespace
    }
}

impl<'s> From<&Element<'s>> for Vec<Event<'s>> {
    fn from(element: &Element<'s>) -> Self {
        let name = element.name();

        let event = BytesStart::new(name);

        // namespace declarations
        let namespace_declarations = element.namespace_declarations.iter().map(|declaration| {
            let (prefix, namespace) = declaration;
            match prefix {
                Some(prefix) => return (format!("xmlns:{}", prefix).as_str(), *namespace),
                None => return ("xmlns", *namespace),
            }
        });
        let event = event.with_attributes(namespace_declarations);

        // attributes
        let event = event.with_attributes(element.attributes.into_iter());

        match element.children.is_empty() {
            true => return vec![Event::Empty(event)],
            false => {
                return {
                    let start: Vec<Event<'s>> = vec![Event::Start(event)];
                    let start_and_content: Vec<Event<'s>> = start
                        .into_iter()
                        .chain({
                            let u = element.children.iter().fold(
                                Vec::new(),
                                |acc: Vec<Event<'s>>, child: &Node<'s>| {
                                    acc.into_iter()
                                        .chain(Into::<Vec<Event<'s>>>::into(child).into_iter())
                                        .collect()
                                },
                            );
                            u
                        })
                        .collect();
                    let full: Vec<Event<'s>> = start_and_content
                        .into_iter()
                        .chain(vec![Event::End(BytesEnd::new(name))])
                        .collect();
                    full
                }
            }
        }
    }
}

impl<'s> Element<'s> {
    /// if there is only one child in the vec of children, will return that element
    pub fn child(&self) -> Result<&Node<'s>> {
        if self.children.len() == 1 {
            Ok(&self.children[0])
        } else if self.children.len() > 1 {
            Err(ElementError::MultipleChildren.into())
        } else {
            Err(ElementError::NoChildren.into())
        }
    }

    /// returns reference to children
    pub fn children(&self) -> Result<&Vec<Node<'s>>> {
        if !self.children.is_empty() {
            Ok(&self.children)
        } else {
            Err(ElementError::NoChildren.into())
        }
    }

    /// returns text content, error if there is none
    pub fn text_content(&self) -> Result<Vec<&str>> {
        let text = Vec::new();
        for node in *self.children {
            match node {
                Node::Text(t) => text.push(t),
                _ => {}
            }
        }
        if text.is_empty() {
            Err(ElementError::NotText)
        }
        Ok(text)
    }

    pub async fn write<W: AsyncWrite + Unpin + Send>(&self, writer: &mut Writer<W>) -> Result<()> {
        let events: Vec<Event> = self.into();
        for event in events {
            writer.write_event_async(event).await?
        }
        Ok(())
    }

    pub async fn write_start<W: AsyncWrite + Unpin + Send>(
        &self,
        writer: &mut Writer<W>,
    ) -> Result<()> {
        let mut event = BytesStart::new(self.name());

        // namespace declarations
        self.namespace_declarations.iter().for_each(|declaration| {
            let (prefix, namespace) = declaration;
            match prefix {
                Some(prefix) => {
                    event.push_attribute((format!("xmlns:{}", prefix).as_str(), *namespace))
                }
                None => event.push_attribute(("xmlns", *namespace)),
            }
        });

        // attributes
        let event =
            event.with_attributes(self.attributes.iter().map(|(attr, value)| (*attr, *value)));

        writer.write_event_async(Event::Start(event)).await?;

        Ok(())
    }

    pub async fn write_end<W: AsyncWrite + Unpin + Send>(
        &self,
        writer: &mut Writer<W>,
    ) -> Result<()> {
        let event = BytesEnd::new(self.name());
        writer.write_event_async(Event::End(event)).await?;
        Ok(())
    }

    #[async_recursion]
    pub async fn read<R: AsyncBufRead + Unpin + Send>(
        reader: &mut Reader<R>,
        local_namespaces: &BTreeMap<Option<&str>, &str>,
    ) -> Result<Self> {
        let node = Node::read_recursive(reader, local_namespaces)
            .await?
            .ok_or(JabberError::UnexpectedEnd)?;
        match node {
            Node::Element(e) => Ok(e),
            Node::Text(_) => Err(JabberError::UnexpectedText),
            Node::Unknown => Err(JabberError::UnexpectedElement),
        }
    }
    pub async fn read_start<R: AsyncBufRead + Unpin + Send>(
        reader: &mut Reader<R>,
        local_namespaces: &BTreeMap<Option<&str>, &str>,
    ) -> Result<Self> {
        let mut buf = Vec::new();
        let event = reader.read_event_into_async(&mut buf).await?;
        match event {
            Event::Start(e) => {
                let prefix = e
                    .name()
                    .prefix()
                    .map(|prefix| str::from_utf8(prefix.into_inner())?);
                let localname = str::from_utf8(e.local_name().into_inner())?;

                let mut namespaces = local_namespaces.clone();
                let mut namespace_declarations = BTreeMap::new();
                let attributes = BTreeMap::new();

                for attribute in e.attributes() {
                    let attribute = attribute?;
                    if let Some(prefix_declaration) = attribute.key.as_namespace_binding() {
                        match prefix_declaration {
                            PrefixDeclaration::Default => {
                                let value = str::from_utf8(attribute.value.as_ref())?;
                                namespace_declarations.try_insert(None, value);
                                namespaces.insert(None, value);
                            }
                            PrefixDeclaration::Named(prefix) => {
                                let key = str::from_utf8(prefix)?;
                                let value = str::from_utf8(attribute.value.as_ref())?;
                                namespace_declarations
                                    .try_insert(Some(key), value)
                                    .map_err(ParseError::DuplicateAttribute)?;
                                namespaces.insert(Some(key), value);
                            }
                        }
                    } else {
                        attributes
                            .try_insert(
                                str::from_utf8(attribute.key.into_inner())?,
                                str::from_utf8(attribute.value.as_ref())?,
                            )
                            .map_err(ParseError::DuplicateAttribute)?;
                    }
                }

                let namespace = *namespaces.get(&prefix).ok_or(ParseError::NoNamespace)?;

                let mut children = Vec::new();

                Ok(Some(Self::Element(Element {
                    prefix,
                    localname,
                    namespace,
                    namespace_declarations,
                    attributes,
                    children,
                })))
            }
            e => Err(ElementError::NotAStart(e)),
        }
    }
}

impl<'s> Node<'s> {
    async fn read_recursive<R: AsyncBufRead + Unpin + Send>(
        reader: &mut Reader<R>,
        local_namespaces: &BTreeMap<Option<&str>, &str>,
    ) -> Result<Option<Node<'s>>> {
        let mut buf = Vec::new();
        let event = reader.read_event_into_async(&mut buf).await?;
        match event {
            Event::Empty(e) => {
                let prefix = e
                    .name()
                    .prefix()
                    .map(|prefix| str::from_utf8(prefix.into_inner())?);
                let localname = str::from_utf8(e.local_name().into_inner())?;

                let mut namespaces = local_namespaces.clone();
                let mut namespace_declarations = BTreeMap::new();
                let attributes = BTreeMap::new();

                for attribute in e.attributes() {
                    let attribute = attribute?;
                    if let Some(prefix_declaration) = attribute.key.as_namespace_binding() {
                        match prefix_declaration {
                            PrefixDeclaration::Default => {
                                let value = str::from_utf8(attribute.value.as_ref())?;
                                namespace_declarations.try_insert(None, value);
                                namespaces.insert(None, value);
                            }
                            PrefixDeclaration::Named(prefix) => {
                                let key = str::from_utf8(prefix)?;
                                let value = str::from_utf8(attribute.value.as_ref())?;
                                namespace_declarations
                                    .try_insert(Some(key), value)
                                    .map_err(ParseError::DuplicateAttribute)?;
                                namespaces.insert(Some(key), value);
                            }
                        }
                    } else {
                        attributes
                            .try_insert(
                                str::from_utf8(attribute.key.into_inner())?,
                                str::from_utf8(attribute.value.as_ref())?,
                            )
                            .map_err(ParseError::DuplicateAttribute)?;
                    }
                }

                let namespace = *namespaces.get(&prefix).ok_or(ParseError::NoNamespace)?;

                let children = Vec::new();

                Ok(Some(Self::Element(Element {
                    prefix,
                    localname,
                    namespace,
                    namespace_declarations,
                    attributes,
                    children,
                })))
            }
            Event::Start(e) => {
                let prefix = e
                    .name()
                    .prefix()
                    .map(|prefix| str::from_utf8(prefix.into_inner())?);
                let localname = str::from_utf8(e.local_name().into_inner())?;

                let mut namespaces = local_namespaces.clone();
                let mut namespace_declarations = BTreeMap::new();
                let attributes = BTreeMap::new();

                for attribute in e.attributes() {
                    let attribute = attribute?;
                    if let Some(prefix_declaration) = attribute.key.as_namespace_binding() {
                        match prefix_declaration {
                            PrefixDeclaration::Default => {
                                let value = str::from_utf8(attribute.value.as_ref())?;
                                namespace_declarations.try_insert(None, value);
                                namespaces.insert(None, value);
                            }
                            PrefixDeclaration::Named(prefix) => {
                                let key = str::from_utf8(prefix)?;
                                let value = str::from_utf8(attribute.value.as_ref())?;
                                namespace_declarations
                                    .try_insert(Some(key), value)
                                    .map_err(ParseError::DuplicateAttribute)?;
                                namespaces.insert(Some(key), value);
                            }
                        }
                    } else {
                        attributes
                            .try_insert(
                                str::from_utf8(attribute.key.into_inner())?,
                                str::from_utf8(attribute.value.as_ref())?,
                            )
                            .map_err(ParseError::DuplicateAttribute)?;
                    }
                }

                let namespace = *namespaces.get(&prefix).ok_or(ParseError::NoNamespace)?;

                let mut children = Vec::new();
                while let Some(child_node) = Node::read_recursive(reader, &namespaces).await? {
                    children.push(child_node)
                }

                Ok(Some(Self::Element(Element {
                    prefix,
                    localname,
                    namespace,
                    namespace_declarations,
                    attributes,
                    children,
                })))
            }
            Event::End(_) => Ok(None),
            Event::Text(e) => Ok(Some(Self::Text(e.unescape()?.as_ref()))),
            e => Ok(Some(Self::Unknown)),
        }
    }
}

// #[async_recursion]
// pub async fn read_start<R: AsyncBufRead + Unpin + Send>(
//     reader: &mut Reader<R>,
// ) -> Result<Self, JabberError> {
//     let mut buf = Vec::new();
//     let event = reader.read_event_into_async(&mut buf).await?;
//     match event {
//         Event::Start(e) => {
//             return Ok(Self {
//                 event: Event::Start(e.into_owned()),
//                 children: None,
//             })
//         }
//         e => Err(ElementError::NotAStart(e.into_owned()).into()),
//     }
// }

// pub trait IntoElement<'e> {
//     fn event(&self) -> Event<'e>;
//     fn children(&self) -> Option<Vec<Element<'e>>>;
// }

// impl<'e, T: IntoElement<'e>> From<T> for Element<'e> {
//     fn from(value: T) -> Self {
//         Element {
//             event: value.event(),
//             children: value.children(),
//         }
//     }
// }

#[derive(Debug)]
pub enum ElementError<'e> {
    NotAStart(Event<'e>),
    NotText,
    NoChildren,
    MultipleChildren,
}

#[derive(Debug)]
pub enum ParseError {
    DuplicateAttribute,
    NoNamespace,
}