1//! XML stream parser for XMPP
2
3use crate::{ParseError, ParserError};
4use bytes::{BufMut, BytesMut};
5use xmpp_parsers::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 fn to_io_err<E: Into<Box<dyn std::error::Error + Send + Sync>>>(e: E) -> io::Error {
281 io::Error::new(io::ErrorKind::InvalidInput, e)
282 }
283
284 match item {
285 Packet::StreamStart(start_attrs) => {
286 let mut buf = String::new();
287 write!(buf, "<stream:stream")
288 .map_err(to_io_err)?;
289 for (name, value) in start_attrs {
290 write!(buf, " {}=\"{}\"", escape(&name), escape(&value))
291 .map_err(to_io_err)?;
292 if name == "xmlns" {
293 self.ns = Some(value);
294 }
295 }
296 write!(buf, ">\n")
297 .map_err(to_io_err)?;
298
299 // print!(">> {}", buf);
300 write!(dst, "{}", buf)
301 .map_err(to_io_err)
302 }
303 Packet::Stanza(stanza) => {
304 stanza
305 .write_to_inner(&mut EventWriter::new(WriteBytes::new(dst)))
306 .and_then(|_| {
307 // println!(">> {:?}", dst);
308 Ok(())
309 })
310 .map_err(|e| to_io_err(format!("{}", e)))
311 }
312 Packet::Text(text) => {
313 write_text(&text, dst)
314 .and_then(|_| {
315 // println!(">> {:?}", dst);
316 Ok(())
317 })
318 .map_err(to_io_err)
319 }
320 Packet::StreamEnd => {
321 write!(dst, "</stream:stream>\n")
322 .map_err(to_io_err)
323 }
324 }
325 }
326}
327
328/// Write XML-escaped text string
329pub fn write_text<W: Write>(text: &str, writer: &mut W) -> Result<(), std::fmt::Error> {
330 write!(writer, "{}", escape(text))
331}
332
333/// Copied from `RustyXML` for now
334pub fn escape(input: &str) -> String {
335 let mut result = String::with_capacity(input.len());
336
337 for c in input.chars() {
338 match c {
339 '&' => result.push_str("&"),
340 '<' => result.push_str("<"),
341 '>' => result.push_str(">"),
342 '\'' => result.push_str("'"),
343 '"' => result.push_str("""),
344 o => result.push(o),
345 }
346 }
347 result
348}
349
350/// BytesMut impl only std::fmt::Write but not std::io::Write. The
351/// latter trait is required for minidom's
352/// `Element::write_to_inner()`.
353struct WriteBytes<'a> {
354 dst: &'a mut BytesMut,
355}
356
357impl<'a> WriteBytes<'a> {
358 fn new(dst: &'a mut BytesMut) -> Self {
359 WriteBytes { dst }
360 }
361}
362
363impl<'a> std::io::Write for WriteBytes<'a> {
364 fn write(&mut self, buf: &[u8]) -> std::result::Result<usize, std::io::Error> {
365 self.dst.put_slice(buf);
366 Ok(buf.len())
367 }
368
369 fn flush(&mut self) -> std::result::Result<(), std::io::Error> {
370 Ok(())
371 }
372}
373
374#[cfg(test)]
375mod tests {
376 use super::*;
377 use bytes::BytesMut;
378
379 #[test]
380 fn test_stream_start() {
381 let mut c = XMPPCodec::new();
382 let mut b = BytesMut::with_capacity(1024);
383 b.put(r"<?xml version='1.0'?><stream:stream xmlns:stream='http://etherx.jabber.org/streams' version='1.0' xmlns='jabber:client'>");
384 let r = c.decode(&mut b);
385 assert!(match r {
386 Ok(Some(Packet::StreamStart(_))) => true,
387 _ => false,
388 });
389 }
390
391 #[test]
392 fn test_stream_end() {
393 let mut c = XMPPCodec::new();
394 let mut b = BytesMut::with_capacity(1024);
395 b.put(r"<?xml version='1.0'?><stream:stream xmlns:stream='http://etherx.jabber.org/streams' version='1.0' xmlns='jabber:client'>");
396 let r = c.decode(&mut b);
397 assert!(match r {
398 Ok(Some(Packet::StreamStart(_))) => true,
399 _ => false,
400 });
401 b.clear();
402 b.put(r"</stream:stream>");
403 let r = c.decode(&mut b);
404 assert!(match r {
405 Ok(Some(Packet::StreamEnd)) => true,
406 _ => false,
407 });
408 }
409
410 #[test]
411 fn test_truncated_stanza() {
412 let mut c = XMPPCodec::new();
413 let mut b = BytesMut::with_capacity(1024);
414 b.put(r"<?xml version='1.0'?><stream:stream xmlns:stream='http://etherx.jabber.org/streams' version='1.0' xmlns='jabber:client'>");
415 let r = c.decode(&mut b);
416 assert!(match r {
417 Ok(Some(Packet::StreamStart(_))) => true,
418 _ => false,
419 });
420
421 b.clear();
422 b.put(r"<test>ß</test");
423 let r = c.decode(&mut b);
424 assert!(match r {
425 Ok(None) => true,
426 _ => false,
427 });
428
429 b.clear();
430 b.put(r">");
431 let r = c.decode(&mut b);
432 assert!(match r {
433 Ok(Some(Packet::Stanza(ref el))) if el.name() == "test" && el.text() == "ß" => true,
434 _ => false,
435 });
436 }
437
438 #[test]
439 fn test_truncated_utf8() {
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(&b"<test>\xc3"[..]);
451 let r = c.decode(&mut b);
452 assert!(match r {
453 Ok(None) => true,
454 _ => false,
455 });
456
457 b.clear();
458 b.put(&b"\x9f</test>"[..]);
459 let r = c.decode(&mut b);
460 assert!(match r {
461 Ok(Some(Packet::Stanza(ref el))) if el.name() == "test" && el.text() == "ß" => true,
462 _ => false,
463 });
464 }
465
466 /// test case for https://gitlab.com/xmpp-rs/tokio-xmpp/issues/3
467 #[test]
468 fn test_atrribute_prefix() {
469 let mut c = XMPPCodec::new();
470 let mut b = BytesMut::with_capacity(1024);
471 b.put(r"<?xml version='1.0'?><stream:stream xmlns:stream='http://etherx.jabber.org/streams' version='1.0' xmlns='jabber:client'>");
472 let r = c.decode(&mut b);
473 assert!(match r {
474 Ok(Some(Packet::StreamStart(_))) => true,
475 _ => false,
476 });
477
478 b.clear();
479 b.put(r"<status xml:lang='en'>Test status</status>");
480 let r = c.decode(&mut b);
481 assert!(match r {
482 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,
483 _ => false,
484 });
485
486 }
487
488 /// By default, encode() only get's a BytesMut that has 8kb space reserved.
489 #[test]
490 fn test_large_stanza() {
491 use futures::{Future, Sink};
492 use std::io::Cursor;
493 use tokio_codec::FramedWrite;
494 let framed = FramedWrite::new(Cursor::new(vec![]), XMPPCodec::new());
495 let mut text = "".to_owned();
496 for _ in 0..2usize.pow(15) {
497 text = text + "A";
498 }
499 let stanza = Element::builder("message")
500 .append(Element::builder("body").append(&text).build())
501 .build();
502 let framed = framed.send(Packet::Stanza(stanza)).wait().expect("send");
503 assert_eq!(
504 framed.get_ref().get_ref(),
505 &("<message><body>".to_owned() + &text + "</body></message>").as_bytes()
506 );
507 }
508}