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