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)]
/// represents an xml element as a tree of nodes
pub struct Element {
    /// element prefix
    /// e.g. `foo` in `<foo:bar />`.
    prefix: Option<String>,
    /// element name
    /// e.g. `bar` in `<foo:bar />`.
    localname: String,
    /// 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: String,
    /// 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
    // TODO: maybe not even needed, as can calculate when writing which namespaces need to be declared
    // but then can't have unused namespace on element, confusing.
    namespace_declarations: Box<BTreeMap<Option<String>, String>>,
    /// element attributes
    attributes: Box<BTreeMap<String, String>>,
    // children elements namespaces contain their parents' namespaces
    children: Box<Vec<Node>>,
}

#[derive(Clone, Debug)]
pub enum Node {
    Element(Element),
    Text(String),
    Unknown,
}

impl From<Element> for Node {
    fn from(element: Element) -> Self {
        Self::Element(element)
    }
}

impl<S: ToString> From<S> for Node {
    fn from(text: S) -> Self {
        Self::Text(text.to_string())
    }
}

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

impl Element {
    /// 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> for Vec<Event<'s>> {
    fn from(element: &Element) -> 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 Element {
    /// if there is only one child in the vec of children, will return that element
    pub fn child(&self) -> Result<&Node> {
        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>> {
        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 mut text = Vec::new();
        for node in *self.children {
            match node {
                Node::Text(t) => text.push(t),
                _ => {}
            }
        }
        if text.is_empty() {
            return Err(ElementError::NotText.into());
        }
        Ok(text)
    }

    /// returns whether or not the element is qualified by a namespace, either declared
    /// by a parent, or itself.
    fn namespace_qualified<S: AsRef<str>>(
        &self,
        namespace_context: &BTreeMap<Option<S>, S>,
    ) -> bool {
        // create a new local_namespaces combining that in the context and those declared within the element
        let mut local_namespaces = *namespace_context.clone();
        self.namespace_declarations.iter().for_each(|prefix, declaration| local_namespaces.insert(prefix, declaration));

        if let Some(namespace) = local_namespaces.get(self.prefix) {
            if namespace != self.namespace {
                return false;
            }
        } else {
            return false;
        };

        for child in *self.children {
            if child.namespace_qualified(&local_namespaces) == false {
                return false;
            }
        }

        true
    }


    /// writes an element to a writer. the element's namespace must be qualified by the
    /// context given in `local_namespaces` or the element's internal namespace declarations
    pub async fn write<S: AsRef<str>, W: AsyncWrite + Unpin + Send>(
        &self,
        writer: &mut Writer<W>,
        local_namespaces: &BTreeMap<Option<S>, S>,
    ) -> Result<()> {
        // TODO: instead of having a check if namespace qualified function, have the namespace declarations be added if needed given the context when converting from `Element` to `Event`s
        if self.namespace_qualified(local_namespaces) {
            let events: Vec<Event> = self.into();
            for event in events {
                writer.write_event_async(event).await?
            }
            Ok(())
        } else {
            Err(ElementError::NamespaceNotQualified.into())
        }
    }

    pub async fn write_start<S: AsRef<str>, W: AsyncWrite + Unpin + Send>(
        &self,
        writer: &mut Writer<W>,
        local_namespaces: &BTreeMap<Option<S>, S>,
    ) -> Result<()> {
        if self.namespace_qualified(local_namespaces) {
            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(())
        } else {
            Err(ElementError::NamespaceNotQualified.into())
        }
    }

    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<S: AsRef<str>, R: AsyncBufRead + Unpin + Send>(
        reader: &mut Reader<R>,
        local_namespaces: &BTreeMap<Option<S>, S>,
    ) -> 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<S: AsRef<str>, R: AsyncBufRead + Unpin + Send>(
        reader: &mut Reader<R>,
        local_namespaces: &BTreeMap<Option<S>, S>,
    ) -> Result<Element> {
        let 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| prefix.into_inner());
                let converted_prefix;
                if let Some(raw_prefix) = prefix {
                    converted_prefix = Some(str::from_utf8(raw_prefix)?)
                }
                let prefix = converted_prefix;

                let localname = str::from_utf8(e.local_name().into_inner())?.to_owned();

                let mut local_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())?;
                                if let Some(_) = namespace_declarations.insert(None, value) {
                                    return Err(ElementParseError::DuplicateAttribute.into());
                                };
                                local_namespaces.insert(None, value);
                            }
                            PrefixDeclaration::Named(prefix) => {
                                let key = str::from_utf8(prefix)?;
                                let value = str::from_utf8(attribute.value.as_ref())?;
                                if let Some(_) = namespace_declarations.insert(Some(key), value) {
                                    return Err(ElementParseError::DuplicateAttribute.into());
                                };
                                local_namespaces.insert(Some(key), value);
                            }
                        }
                    } else {
                        if let Some(_) = attributes.insert(
                            str::from_utf8(attribute.key.into_inner())?,
                            str::from_utf8(attribute.value.as_ref())?,
                        ) {
                            return Err(ElementParseError::DuplicateAttribute.into());
                        };
                    }
                }

                let namespace = *local_namespaces
                    .get(&prefix)
                    .ok_or(ElementParseError::NoNamespace)?;

                let mut children = Vec::new();

                Ok(Self {
                    prefix,
                    localname,
                    namespace,
                    namespace_declarations: Box::new(namespace_declarations),
                    attributes: Box::new(attributes),
                    children: Box::new(children),
                })
            }
            e => Err(ElementError::NotAStart(e).into()),
        }
    }
}

impl Node {
    #[async_recursion]
    async fn read_recursive<S: AsRef<str>, R: AsyncBufRead + Unpin + Send>(
        reader: &mut Reader<R>,
        local_namespaces: &BTreeMap<Option<S>, S>,
    ) -> Result<Option<Node>> {
        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| prefix.into_inner());
                let converted_prefix;
                if let Some(raw_prefix) = prefix {
                    converted_prefix = Some(str::from_utf8(raw_prefix)?)
                }
                let prefix = converted_prefix;

                let localname = str::from_utf8(e.local_name().into_inner())?.to_owned();

                let mut local_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())?;
                                if let Some(_) = namespace_declarations.insert(None, value) {
                                    return Err(ElementParseError::DuplicateAttribute.into());
                                };
                                local_namespaces.insert(None, value);
                            }
                            PrefixDeclaration::Named(prefix) => {
                                let key = str::from_utf8(prefix)?;
                                let value = str::from_utf8(attribute.value.as_ref())?;
                                if let Some(_) = namespace_declarations.insert(Some(key), value) {
                                    return Err(ElementParseError::DuplicateAttribute.into());
                                };
                                local_namespaces.insert(Some(key), value);
                            }
                        }
                    } else {
                        if let Some(_) = attributes.insert(
                            str::from_utf8(attribute.key.into_inner())?,
                            str::from_utf8(attribute.value.as_ref())?,
                        ) {
                            return Err(ElementParseError::DuplicateAttribute.into());
                        };
                    }
                }

                let namespace = *local_namespaces
                    .get(&prefix)
                    .ok_or(ElementParseError::NoNamespace)?;

                let mut children = Vec::new();

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

                let localname = str::from_utf8(e.local_name().into_inner())?.to_owned();

                let mut local_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())?;
                                if let Some(_) = namespace_declarations.insert(None, value) {
                                    return Err(ElementParseError::DuplicateAttribute.into());
                                };
                                local_namespaces.insert(None, value);
                            }
                            PrefixDeclaration::Named(prefix) => {
                                let key = str::from_utf8(prefix)?;
                                let value = str::from_utf8(attribute.value.as_ref())?;
                                if let Some(_) = namespace_declarations.insert(Some(key), value) {
                                    return Err(ElementParseError::DuplicateAttribute.into());
                                };
                                local_namespaces.insert(Some(key), value);
                            }
                        }
                    } else {
                        if let Some(_) = attributes.insert(
                            str::from_utf8(attribute.key.into_inner())?,
                            str::from_utf8(attribute.value.as_ref())?,
                        ) {
                            return Err(ElementParseError::DuplicateAttribute.into());
                        };
                    }
                }

                let namespace = *local_namespaces
                    .get(&prefix)
                    .ok_or(ElementParseError::NoNamespace)?;

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

                let mut children = Vec::new();

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

    fn namespace_qualified<S: AsRef<str>>(
        &self,
        namespace_context: &BTreeMap<Option<S>, S>,
    ) -> bool {
        match self {
            Self::Element(e) => e.namespace_qualified(namespace_context),
            _ => true,
        }
    }
}

pub enum NodeBuilder {
    Text(String),
    Element(ElementBuilder),
}

pub struct ElementBuilder {
    localname: String,
    prefix: Option<String>,
    namespace: String,
    namespace_declarations: BTreeMap<Option<String>, String>,
    attributes: BTreeMap<String, String>,
    children: Vec<NodeBuilder>,
}

impl ElementBuilder {
    pub fn new<S: ToString>(localname: S, prefix: Option<S>, namespace: S) -> Self {
        Self {
            prefix,
            localname,
            namespace,
            namespace_declarations: Box::new(BTreeMap::new()),
            attributes: Box::new(BTreeMap::new()),
            children: Box::new(Vec::new()),
        }
    }

    pub fn push_namespace_declaration<S: ToString>(
        &mut self,
        (prefix, namespace): (Option<S>, S),
    ) -> Option<S> {
        self.namespace_declarations.insert(prefix, namespace)
    }

    pub fn push_attribute<S: ToString>(&mut self, (key, value): (S, S)) -> Option<S> {
        self.attributes.insert(key, value)
    }

    pub fn push_child(&mut self, child: Node) {
        self.children.push(child)
    }

    /// checks if there is a namespace conflict within the element being built
    pub fn namespace_conflict<S: AsRef<str>>(
        &self
    ) -> bool {
        self.namespace_conflict_recursive(&BTreeMap::new())
    }

    fn namespace_conflict_recursive<S: AsRef<str>>(
        &self,
        parent_namespaces: &BTreeMap<Option<S>, S>,
    ) -> bool {
        // create a new local_namespaces combining that in the context and those declared within the element
        let mut local_namespaces = *parent_namespaces.clone();
        self.namespace_declarations.iter().for_each(|prefix, declaration| local_namespaces.insert(prefix, declaration));

        if let Some(namespace) = local_namespaces.get(self.prefix) {
            if namespace != self.namespace {
                return false;
            }
        } else {
            return false;
        };

        for child in *self.children {
            if child.namespace_conflict(&local_namespaces) == false {
                return false;
            }
        }

        true
    }

    // check for possible conflicts in namespace
    pub fn build(self) -> Result<Element> {
        for child in self.children {
            match child {
                Node::Element(e) => {
                    if !e.namespace_conflict()
                }
            }
        }
        Element {
            prefix: self.prefix,
            localname: self.localname,
            namespace: self.namespace,
            namespace_declarations: self.namespace_declarations,
            attributes: self.attributes,
            children: self.children,
        }
    }
}

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

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