1use std;
  2use std::fmt::Write;
  3use std::str::from_utf8;
  4use std::io::{Error, ErrorKind};
  5use std::collections::HashMap;
  6use tokio_io::codec::{Encoder, Decoder};
  7use xml;
  8use bytes::*;
  9
 10const NS_XMLNS: &'static str = "http://www.w3.org/2000/xmlns/";
 11
 12pub type Attributes = HashMap<(String, Option<String>), String>;
 13
 14struct XMPPRoot {
 15    builder: xml::ElementBuilder,
 16    pub attributes: Attributes,
 17}
 18
 19impl XMPPRoot {
 20    fn new(root: xml::StartTag) -> Self {
 21        let mut builder = xml::ElementBuilder::new();
 22        let mut attributes = HashMap::new();
 23        for (name_ns, value) in root.attributes {
 24            match name_ns {
 25                (ref name, None) if name == "xmlns" =>
 26                    builder.set_default_ns(value),
 27                (ref prefix, Some(ref ns)) if ns == NS_XMLNS =>
 28                    builder.define_prefix(prefix.to_owned(), value),
 29                _ => {
 30                    attributes.insert(name_ns, value);
 31                },
 32            }
 33        }
 34
 35        XMPPRoot {
 36            builder: builder,
 37            attributes: attributes,
 38        }
 39    }
 40
 41    fn handle_event(&mut self, event: Result<xml::Event, xml::ParserError>)
 42                    -> Option<Result<xml::Element, xml::BuilderError>> {
 43        self.builder.handle_event(event)
 44    }
 45}
 46
 47#[derive(Debug)]
 48pub enum Packet {
 49    Error(Box<std::error::Error>),
 50    StreamStart(HashMap<String, String>),
 51    Stanza(xml::Element),
 52    Text(String),
 53    StreamEnd,
 54}
 55
 56pub struct XMPPCodec {
 57    parser: xml::Parser,
 58    root: Option<XMPPRoot>,
 59}
 60
 61impl XMPPCodec {
 62    pub fn new() -> Self {
 63        XMPPCodec {
 64            parser: xml::Parser::new(),
 65            root: None,
 66        }
 67    }
 68}
 69
 70impl Decoder for XMPPCodec {
 71    type Item = Packet;
 72    type Error = Error;
 73
 74    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
 75        match from_utf8(buf.take().as_ref()) {
 76            Ok(s) => {
 77                if s.len() > 0 {
 78                    println!("<< {}", s);
 79                    self.parser.feed_str(s);
 80                }
 81            },
 82            Err(e) =>
 83                return Err(Error::new(ErrorKind::InvalidInput, e)),
 84        }
 85
 86        let mut new_root: Option<XMPPRoot> = None;
 87        let mut result = None;
 88        for event in &mut self.parser {
 89            match self.root {
 90                None => {
 91                    // Expecting <stream:stream>
 92                    match event {
 93                        Ok(xml::Event::ElementStart(start_tag)) => {
 94                            let mut attrs: HashMap<String, String> = HashMap::new();
 95                            for (&(ref name, _), value) in &start_tag.attributes {
 96                                attrs.insert(name.to_owned(), value.to_owned());
 97                            }
 98                            result = Some(Packet::StreamStart(attrs));
 99                            self.root = Some(XMPPRoot::new(start_tag));
100                            break
101                        },
102                        Err(e) => {
103                            result = Some(Packet::Error(Box::new(e)));
104                            break
105                        },
106                        _ =>
107                            (),
108                    }
109                }
110
111                Some(ref mut root) => {
112                    match root.handle_event(event) {
113                        None => (),
114                        Some(Ok(stanza)) => {
115                            // Emit the stanza
116                            result = Some(Packet::Stanza(stanza));
117                            break
118                        },
119                        Some(Err(e)) => {
120                            result = Some(Packet::Error(Box::new(e)));
121                            break
122                        }
123                    };
124                },
125            }
126
127            match new_root.take() {
128                None => (),
129                Some(root) => self.root = Some(root),
130            }
131        }
132
133        Ok(result)
134    }
135
136    fn decode_eof(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, Error> {
137        self.decode(buf)
138    }
139}
140
141impl Encoder for XMPPCodec {
142    type Item = Packet;
143    type Error = Error;
144
145    fn encode(&mut self, item: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error> {
146        match item {
147            Packet::StreamStart(start_attrs) => {
148                let mut buf = String::new();
149                write!(buf, "<stream:stream").unwrap();
150                for (ref name, ref value) in &start_attrs {
151                    write!(buf, " {}=\"{}\"", xml::escape(&name), xml::escape(&value))
152                        .unwrap();
153                }
154                write!(buf, ">\n").unwrap();
155
156                print!(">> {}", buf);
157                write!(dst, "{}", buf)
158            },
159            Packet::Stanza(stanza) => {
160                println!(">> {}", stanza);
161                write!(dst, "{}", stanza)
162            },
163            Packet::Text(text) => {
164                let escaped = xml::escape(&text);
165                println!(">> {}", escaped);
166                write!(dst, "{}", escaped)
167            },
168            // TODO: Implement all
169            _ => Ok(())
170        }
171        .map_err(|_| Error::from(ErrorKind::InvalidInput))
172    }
173}