1//! XML stream parser for XMPP
2
3use crate::{ParseError, ParserError};
4use bytes::{BufMut, BytesMut};
5use minidom::Element;
6use quick_xml::Writer as EventWriter;
7use std;
8use std::cell::RefCell;
9use std::collections::vec_deque::VecDeque;
10use std::collections::HashMap;
11use std::default::Default;
12use std::fmt::Write;
13use std::io;
14use std::iter::FromIterator;
15use std::rc::Rc;
16use std::str::from_utf8;
17use std::borrow::Cow;
18use tokio_codec::{Decoder, Encoder};
19use xml5ever::interface::Attribute;
20use xml5ever::tokenizer::{Tag, TagKind, Token, TokenSink, XmlTokenizer};
21
22/// Anything that can be sent or received on an XMPP/XML stream
23#[derive(Debug, Clone, PartialEq, Eq)]
24pub enum Packet {
25 /// `<stream:stream>` start tag
26 StreamStart(HashMap<String, String>),
27 /// A complete stanza or nonza
28 Stanza(Element),
29 /// Plain text (think whitespace keep-alive)
30 Text(String),
31 /// `</stream:stream>` closing tag
32 StreamEnd,
33}
34
35type QueueItem = Result<Packet, ParserError>;
36
37/// Parser state
38struct ParserSink {
39 // Ready stanzas, shared with XMPPCodec
40 queue: Rc<RefCell<VecDeque<QueueItem>>>,
41 // Parsing stack
42 stack: Vec<Element>,
43 ns_stack: Vec<HashMap<Option<String>, String>>,
44}
45
46impl ParserSink {
47 pub fn new(queue: Rc<RefCell<VecDeque<QueueItem>>>) -> Self {
48 ParserSink {
49 queue,
50 stack: vec![],
51 ns_stack: vec![],
52 }
53 }
54
55 fn push_queue(&self, pkt: Packet) {
56 self.queue.borrow_mut().push_back(Ok(pkt));
57 }
58
59 fn push_queue_error(&self, e: ParserError) {
60 self.queue.borrow_mut().push_back(Err(e));
61 }
62
63 /// Lookup XML namespace declaration for given prefix (or no prefix)
64 fn lookup_ns(&self, prefix: &Option<String>) -> Option<&str> {
65 for nss in self.ns_stack.iter().rev() {
66 if let Some(ns) = nss.get(prefix) {
67 return Some(ns);
68 }
69 }
70
71 None
72 }
73
74 fn handle_start_tag(&mut self, tag: Tag) {
75 let mut nss = HashMap::new();
76 let is_prefix_xmlns = |attr: &Attribute| {
77 attr.name
78 .prefix
79 .as_ref()
80 .map(|prefix| prefix.eq_str_ignore_ascii_case("xmlns"))
81 .unwrap_or(false)
82 };
83 for attr in &tag.attrs {
84 match attr.name.local.as_ref() {
85 "xmlns" => {
86 nss.insert(None, attr.value.as_ref().to_owned());
87 }
88 prefix if is_prefix_xmlns(attr) => {
89 nss.insert(Some(prefix.to_owned()), attr.value.as_ref().to_owned());
90 }
91 _ => (),
92 }
93 }
94 self.ns_stack.push(nss);
95
96 let el = {
97 let mut el_builder = Element::builder(tag.name.local.as_ref());
98 if let Some(el_ns) =
99 self.lookup_ns(&tag.name.prefix.map(|prefix| prefix.as_ref().to_owned()))
100 {
101 el_builder = el_builder.ns(el_ns);
102 }
103 for attr in &tag.attrs {
104 match attr.name.local.as_ref() {
105 "xmlns" => (),
106 _ if is_prefix_xmlns(attr) => (),
107 _ => {
108 let attr_name = if let Some(ref prefix) = attr.name.prefix {
109 Cow::Owned(format!("{}:{}", prefix, attr.name.local))
110 } else {
111 Cow::Borrowed(attr.name.local.as_ref())
112 };
113 el_builder = el_builder.attr(attr_name, attr.value.as_ref());
114 }
115 }
116 }
117 el_builder.build()
118 };
119
120 if self.stack.is_empty() {
121 let attrs = HashMap::from_iter(tag.attrs.iter().map(|attr| {
122 (
123 attr.name.local.as_ref().to_owned(),
124 attr.value.as_ref().to_owned(),
125 )
126 }));
127 self.push_queue(Packet::StreamStart(attrs));
128 }
129
130 self.stack.push(el);
131 }
132
133 fn handle_end_tag(&mut self) {
134 let el = self.stack.pop().unwrap();
135 self.ns_stack.pop();
136
137 match self.stack.len() {
138 // </stream:stream>
139 0 => self.push_queue(Packet::StreamEnd),
140 // </stanza>
141 1 => self.push_queue(Packet::Stanza(el)),
142 len => {
143 let parent = &mut self.stack[len - 1];
144 parent.append_child(el);
145 }
146 }
147 }
148}
149
150impl TokenSink for ParserSink {
151 fn process_token(&mut self, token: Token) {
152 match token {
153 Token::TagToken(tag) => match tag.kind {
154 TagKind::StartTag => self.handle_start_tag(tag),
155 TagKind::EndTag => self.handle_end_tag(),
156 TagKind::EmptyTag => {
157 self.handle_start_tag(tag);
158 self.handle_end_tag();
159 }
160 TagKind::ShortTag => self.push_queue_error(ParserError::ShortTag),
161 },
162 Token::CharacterTokens(tendril) => match self.stack.len() {
163 0 | 1 => self.push_queue(Packet::Text(tendril.into())),
164 len => {
165 let el = &mut self.stack[len - 1];
166 el.append_text_node(tendril);
167 }
168 },
169 Token::EOFToken => self.push_queue(Packet::StreamEnd),
170 Token::ParseError(s) => {
171 // println!("ParseError: {:?}", s);
172 self.push_queue_error(ParserError::Parse(ParseError(s)));
173 }
174 _ => (),
175 }
176 }
177
178 // fn end(&mut self) {
179 // }
180}
181
182/// Stateful encoder/decoder for a bytestream from/to XMPP `Packet`
183pub struct XMPPCodec {
184 /// Outgoing
185 ns: Option<String>,
186 /// Incoming
187 parser: XmlTokenizer<ParserSink>,
188 /// For handling incoming truncated utf8
189 // TODO: optimize using tendrils?
190 buf: Vec<u8>,
191 /// Shared with ParserSink
192 queue: Rc<RefCell<VecDeque<QueueItem>>>,
193}
194
195impl XMPPCodec {
196 /// Constructor
197 pub fn new() -> Self {
198 let queue = Rc::new(RefCell::new(VecDeque::new()));
199 let sink = ParserSink::new(queue.clone());
200 // TODO: configure parser?
201 let parser = XmlTokenizer::new(sink, Default::default());
202 XMPPCodec {
203 ns: None,
204 parser,
205 queue,
206 buf: vec![],
207 }
208 }
209}
210
211impl Default for XMPPCodec {
212 fn default() -> Self {
213 Self::new()
214 }
215}
216
217impl Decoder for XMPPCodec {
218 type Item = Packet;
219 type Error = ParserError;
220
221 fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
222 let buf1: Box<AsRef<[u8]>> = if !self.buf.is_empty() && !buf.is_empty() {
223 let mut prefix = std::mem::replace(&mut self.buf, vec![]);
224 prefix.extend_from_slice(buf.take().as_ref());
225 Box::new(prefix)
226 } else {
227 Box::new(buf.take())
228 };
229 let buf1 = buf1.as_ref().as_ref();
230 match from_utf8(buf1) {
231 Ok(s) => {
232 if !s.is_empty() {
233 // println!("<< {}", s);
234 let tendril = FromIterator::from_iter(s.chars());
235 self.parser.feed(tendril);
236 }
237 }
238 // Remedies for truncated utf8
239 Err(e) if e.valid_up_to() >= buf1.len() - 3 => {
240 // Prepare all the valid data
241 let mut b = BytesMut::with_capacity(e.valid_up_to());
242 b.put(&buf1[0..e.valid_up_to()]);
243
244 // Retry
245 let result = self.decode(&mut b);
246
247 // Keep the tail back in
248 self.buf.extend_from_slice(&buf1[e.valid_up_to()..]);
249
250 return result;
251 }
252 Err(e) => {
253 // println!("error {} at {}/{} in {:?}", e, e.valid_up_to(), buf1.len(), buf1);
254 return Err(ParserError::Utf8(e));
255 }
256 }
257
258 match self.queue.borrow_mut().pop_front() {
259 None => Ok(None),
260 Some(result) => result.map(|pkt| Some(pkt)),
261 }
262 }
263
264 fn decode_eof(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
265 self.decode(buf)
266 }
267}
268
269impl Encoder for XMPPCodec {
270 type Item = Packet;
271 type Error = io::Error;
272
273 fn encode(&mut self, item: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error> {
274 let remaining = dst.capacity() - dst.len();
275 let max_stanza_size: usize = 2usize.pow(16);
276 if remaining < max_stanza_size {
277 dst.reserve(max_stanza_size - remaining);
278 }
279
280 match item {
281 Packet::StreamStart(start_attrs) => {
282 let mut buf = String::new();
283 write!(buf, "<stream:stream").unwrap();
284 for (name, value) in start_attrs {
285 write!(buf, " {}=\"{}\"", escape(&name), escape(&value)).unwrap();
286 if name == "xmlns" {
287 self.ns = Some(value);
288 }
289 }
290 write!(buf, ">\n").unwrap();
291
292 print!(">> {}", buf);
293 write!(dst, "{}", buf).map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))
294 }
295 Packet::Stanza(stanza) => {
296 stanza
297 .write_to_inner(&mut EventWriter::new(WriteBytes::new(dst)))
298 .and_then(|_| {
299 // println!(">> {:?}", dst);
300 Ok(())
301 })
302 .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, format!("{}", e)))
303 }
304 Packet::Text(text) => {
305 write_text(&text, dst)
306 .and_then(|_| {
307 // println!(">> {:?}", dst);
308 Ok(())
309 })
310 .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, format!("{}", e)))
311 }
312 // TODO: Implement all
313 _ => Ok(()),
314 }
315 }
316}
317
318/// Write XML-escaped text string
319pub fn write_text<W: Write>(text: &str, writer: &mut W) -> Result<(), std::fmt::Error> {
320 write!(writer, "{}", escape(text))
321}
322
323/// Copied from `RustyXML` for now
324pub fn escape(input: &str) -> String {
325 let mut result = String::with_capacity(input.len());
326
327 for c in input.chars() {
328 match c {
329 '&' => result.push_str("&"),
330 '<' => result.push_str("<"),
331 '>' => result.push_str(">"),
332 '\'' => result.push_str("'"),
333 '"' => result.push_str("""),
334 o => result.push(o),
335 }
336 }
337 result
338}
339
340/// BytesMut impl only std::fmt::Write but not std::io::Write. The
341/// latter trait is required for minidom's
342/// `Element::write_to_inner()`.
343struct WriteBytes<'a> {
344 dst: &'a mut BytesMut,
345}
346
347impl<'a> WriteBytes<'a> {
348 fn new(dst: &'a mut BytesMut) -> Self {
349 WriteBytes { dst }
350 }
351}
352
353impl<'a> std::io::Write for WriteBytes<'a> {
354 fn write(&mut self, buf: &[u8]) -> std::result::Result<usize, std::io::Error> {
355 self.dst.put_slice(buf);
356 Ok(buf.len())
357 }
358
359 fn flush(&mut self) -> std::result::Result<(), std::io::Error> {
360 Ok(())
361 }
362}
363
364#[cfg(test)]
365mod tests {
366 use super::*;
367 use bytes::BytesMut;
368
369 #[test]
370 fn test_stream_start() {
371 let mut c = XMPPCodec::new();
372 let mut b = BytesMut::with_capacity(1024);
373 b.put(r"<?xml version='1.0'?><stream:stream xmlns:stream='http://etherx.jabber.org/streams' version='1.0' xmlns='jabber:client'>");
374 let r = c.decode(&mut b);
375 assert!(match r {
376 Ok(Some(Packet::StreamStart(_))) => true,
377 _ => false,
378 });
379 }
380
381 #[test]
382 fn test_truncated_stanza() {
383 let mut c = XMPPCodec::new();
384 let mut b = BytesMut::with_capacity(1024);
385 b.put(r"<?xml version='1.0'?><stream:stream xmlns:stream='http://etherx.jabber.org/streams' version='1.0' xmlns='jabber:client'>");
386 let r = c.decode(&mut b);
387 assert!(match r {
388 Ok(Some(Packet::StreamStart(_))) => true,
389 _ => false,
390 });
391
392 b.clear();
393 b.put(r"<test>ß</test");
394 let r = c.decode(&mut b);
395 assert!(match r {
396 Ok(None) => true,
397 _ => false,
398 });
399
400 b.clear();
401 b.put(r">");
402 let r = c.decode(&mut b);
403 assert!(match r {
404 Ok(Some(Packet::Stanza(ref el))) if el.name() == "test" && el.text() == "ß" => true,
405 _ => false,
406 });
407 }
408
409 #[test]
410 fn test_truncated_utf8() {
411 let mut c = XMPPCodec::new();
412 let mut b = BytesMut::with_capacity(1024);
413 b.put(r"<?xml version='1.0'?><stream:stream xmlns:stream='http://etherx.jabber.org/streams' version='1.0' xmlns='jabber:client'>");
414 let r = c.decode(&mut b);
415 assert!(match r {
416 Ok(Some(Packet::StreamStart(_))) => true,
417 _ => false,
418 });
419
420 b.clear();
421 b.put(&b"<test>\xc3"[..]);
422 let r = c.decode(&mut b);
423 assert!(match r {
424 Ok(None) => true,
425 _ => false,
426 });
427
428 b.clear();
429 b.put(&b"\x9f</test>"[..]);
430 let r = c.decode(&mut b);
431 assert!(match r {
432 Ok(Some(Packet::Stanza(ref el))) if el.name() == "test" && el.text() == "ß" => true,
433 _ => false,
434 });
435 }
436
437 /// test case for https://gitlab.com/xmpp-rs/tokio-xmpp/issues/3
438 #[test]
439 fn test_atrribute_prefix() {
440 let mut c = XMPPCodec::new();
441 let mut b = BytesMut::with_capacity(1024);
442 b.put(r"<?xml version='1.0'?><stream:stream xmlns:stream='http://etherx.jabber.org/streams' version='1.0' xmlns='jabber:client'>");
443 let r = c.decode(&mut b);
444 assert!(match r {
445 Ok(Some(Packet::StreamStart(_))) => true,
446 _ => false,
447 });
448
449 b.clear();
450 b.put(r"<status xml:lang='en'>Test status</status>");
451 let r = c.decode(&mut b);
452 assert!(match r {
453 Ok(Some(Packet::Stanza(ref el))) if el.name() == "status" && el.text() == "Test status" && el.attr("xml:lang").map_or(false, |a| a == "en") => true,
454 _ => false,
455 });
456
457 }
458
459 /// By default, encode() only get's a BytesMut that has 8kb space reserved.
460 #[test]
461 fn test_large_stanza() {
462 use futures::{Future, Sink};
463 use std::io::Cursor;
464 use tokio_codec::FramedWrite;
465 let framed = FramedWrite::new(Cursor::new(vec![]), XMPPCodec::new());
466 let mut text = "".to_owned();
467 for _ in 0..2usize.pow(15) {
468 text = text + "A";
469 }
470 let stanza = Element::builder("message")
471 .append(Element::builder("body").append(&text).build())
472 .build();
473 let framed = framed.send(Packet::Stanza(stanza)).wait().expect("send");
474 assert_eq!(
475 framed.get_ref().get_ref(),
476 &("<message><body>".to_owned() + &text + "</body></message>").as_bytes()
477 );
478 }
479}