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        println!("root attributes: {:?}", root.attributes);
 24        for (name_ns, value) in root.attributes {
 25            match name_ns {
 26                (ref name, None) if name == "xmlns" =>
 27                    builder.set_default_ns(value),
 28                (ref prefix, Some(ref ns)) if ns == NS_XMLNS =>
 29                    builder.define_prefix(prefix.to_owned(), value),
 30                _ => {
 31                    attributes.insert(name_ns, value);
 32                },
 33            }
 34        }
 35
 36        XMPPRoot {
 37            builder: builder,
 38            attributes: attributes,
 39        }
 40    }
 41
 42    fn handle_event(&mut self, event: Result<xml::Event, xml::ParserError>)
 43                    -> Option<Result<xml::Element, xml::BuilderError>> {
 44        self.builder.handle_event(event)
 45    }
 46}
 47
 48#[derive(Debug)]
 49pub enum Packet {
 50    Error(Box<std::error::Error>),
 51    StreamStart(HashMap<String, String>),
 52    Stanza(xml::Element),
 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        println!("XMPPCodec.decode {:?}", buf.len());
 76        match from_utf8(buf.take().as_ref()) {
 77            Ok(s) =>
 78                self.parser.feed_str(s),
 79            Err(e) =>
 80                return Err(Error::new(ErrorKind::InvalidInput, e)),
 81        }
 82
 83        let mut new_root: Option<XMPPRoot> = None;
 84        let mut result = None;
 85        for event in &mut self.parser {
 86            match self.root {
 87                None => {
 88                    // Expecting <stream:stream>
 89                    match event {
 90                        Ok(xml::Event::ElementStart(start_tag)) => {
 91                            let mut attrs: HashMap<String, String> = HashMap::new();
 92                            for (&(ref name, _), value) in &start_tag.attributes {
 93                                attrs.insert(name.to_owned(), value.to_owned());
 94                            }
 95                            result = Some(Packet::StreamStart(attrs));
 96                            self.root = Some(XMPPRoot::new(start_tag));
 97                            break
 98                        },
 99                        Err(e) => {
100                            result = Some(Packet::Error(Box::new(e)));
101                            break
102                        },
103                        _ =>
104                            (),
105                    }
106                }
107
108                Some(ref mut root) => {
109                    match root.handle_event(event) {
110                        None => (),
111                        Some(Ok(stanza)) => {
112                            println!("stanza: {}", stanza);
113                            result = Some(Packet::Stanza(stanza));
114                            break
115                        },
116                        Some(Err(e)) => {
117                            result = Some(Packet::Error(Box::new(e)));
118                            break
119                        }
120                    };
121                },
122            }
123
124            match new_root.take() {
125                None => (),
126                Some(root) => self.root = Some(root),
127            }
128        }
129
130        Ok(result)
131    }
132
133    fn decode_eof(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, Error> {
134        self.decode(buf)
135    }
136}
137
138impl Encoder for XMPPCodec {
139    type Item = Packet;
140    type Error = Error;
141
142    fn encode(&mut self, item: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error> {
143        match item {
144            Packet::StreamStart(start_attrs) => {
145                let mut buf = String::new();
146                write!(buf, "<stream:stream").unwrap();
147                for (ref name, ref value) in &start_attrs {
148                    write!(buf, " {}=\"{}\"", xml::escape(&name), xml::escape(&value))
149                        .unwrap();
150                }
151                write!(buf, ">\n").unwrap();
152
153                println!("Encode start to {}", buf);
154                write!(dst, "{}", buf)
155            },
156            Packet::Stanza(stanza) =>
157                write!(dst, "{}", stanza),
158            // TODO: Implement all
159            _ => Ok(())
160        }
161        .map_err(|_| Error::from(ErrorKind::InvalidInput))
162    }
163}