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 Text(String),
54 StreamEnd,
55}
56
57pub struct XMPPCodec {
58 parser: xml::Parser,
59 root: Option<XMPPRoot>,
60}
61
62impl XMPPCodec {
63 pub fn new() -> Self {
64 XMPPCodec {
65 parser: xml::Parser::new(),
66 root: None,
67 }
68 }
69}
70
71impl Decoder for XMPPCodec {
72 type Item = Packet;
73 type Error = Error;
74
75 fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
76 println!("XMPPCodec.decode {:?}", buf.len());
77 match from_utf8(buf.take().as_ref()) {
78 Ok(s) =>
79 self.parser.feed_str(s),
80 Err(e) =>
81 return Err(Error::new(ErrorKind::InvalidInput, e)),
82 }
83
84 let mut new_root: Option<XMPPRoot> = None;
85 let mut result = None;
86 for event in &mut self.parser {
87 match self.root {
88 None => {
89 // Expecting <stream:stream>
90 match event {
91 Ok(xml::Event::ElementStart(start_tag)) => {
92 let mut attrs: HashMap<String, String> = HashMap::new();
93 for (&(ref name, _), value) in &start_tag.attributes {
94 attrs.insert(name.to_owned(), value.to_owned());
95 }
96 result = Some(Packet::StreamStart(attrs));
97 self.root = Some(XMPPRoot::new(start_tag));
98 break
99 },
100 Err(e) => {
101 result = Some(Packet::Error(Box::new(e)));
102 break
103 },
104 _ =>
105 (),
106 }
107 }
108
109 Some(ref mut root) => {
110 match root.handle_event(event) {
111 None => (),
112 Some(Ok(stanza)) => {
113 println!("stanza: {}", stanza);
114 result = Some(Packet::Stanza(stanza));
115 break
116 },
117 Some(Err(e)) => {
118 result = Some(Packet::Error(Box::new(e)));
119 break
120 }
121 };
122 },
123 }
124
125 match new_root.take() {
126 None => (),
127 Some(root) => self.root = Some(root),
128 }
129 }
130
131 Ok(result)
132 }
133
134 fn decode_eof(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, Error> {
135 self.decode(buf)
136 }
137}
138
139impl Encoder for XMPPCodec {
140 type Item = Packet;
141 type Error = Error;
142
143 fn encode(&mut self, item: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error> {
144 match item {
145 Packet::StreamStart(start_attrs) => {
146 let mut buf = String::new();
147 write!(buf, "<stream:stream").unwrap();
148 for (ref name, ref value) in &start_attrs {
149 write!(buf, " {}=\"{}\"", xml::escape(&name), xml::escape(&value))
150 .unwrap();
151 }
152 write!(buf, ">\n").unwrap();
153
154 println!("Encode start to {}", buf);
155 write!(dst, "{}", buf)
156 },
157 Packet::Stanza(stanza) =>
158 write!(dst, "{}", stanza),
159 Packet::Text(text) =>
160 write!(dst, "{}", xml::escape(&text)),
161 // TODO: Implement all
162 _ => Ok(())
163 }
164 .map_err(|_| Error::from(ErrorKind::InvalidInput))
165 }
166}