From 1e2a6d5656167f35d1ebb51b25362f2e097d8443 Mon Sep 17 00:00:00 2001
From: Nikita Vilunov <nikita@vilunov.me>
Date: Thu, 4 Apr 2024 17:49:03 +0000
Subject: [PATCH] xmpp: make xml-headers optional in the c2s stream

---
 crates/projection-xmpp/src/lib.rs | 31 +++----------------------------
 crates/proto-xmpp/src/stream.rs   | 12 +++++++++++-
 2 files changed, 14 insertions(+), 29 deletions(-)

diff --git a/crates/projection-xmpp/src/lib.rs b/crates/projection-xmpp/src/lib.rs
index a2a0a5b..3be3d43 100644
--- a/crates/projection-xmpp/src/lib.rs
+++ b/crates/projection-xmpp/src/lib.rs
@@ -229,7 +229,6 @@ async fn socket_force_tls(
     use proto_xmpp::tls::*;
     let xml_reader = &mut NsReader::from_reader(reader);
     let xml_writer = &mut Writer::new(writer);
-    read_xml_header(xml_reader, reader_buf).await?;
     let _ = ClientStreamStart::parse(xml_reader, reader_buf).await?;
 
     let event = Event::Decl(BytesDecl::new("1.0", None, None));
@@ -261,7 +260,6 @@ async fn socket_auth(
     reader_buf: &mut Vec<u8>,
     storage: &mut Storage,
 ) -> Result<Authenticated> {
-    read_xml_header(xml_reader, reader_buf).await?;
     let _ = ClientStreamStart::parse(xml_reader, reader_buf).await?;
 
     xml_writer.write_event_async(Event::Decl(BytesDecl::new("1.0", None, None))).await?;
@@ -327,7 +325,6 @@ async fn socket_final(
     user_handle: &mut PlayerConnection,
     rooms: &RoomRegistry,
 ) -> Result<()> {
-    read_xml_header(xml_reader, reader_buf).await?;
     let _ = ClientStreamStart::parse(xml_reader, reader_buf).await?;
 
     xml_writer.write_event_async(Event::Decl(BytesDecl::new("1.0", None, None))).await?;
@@ -419,7 +416,7 @@ struct XmppConnection<'a> {
 impl<'a> XmppConnection<'a> {
     async fn handle_packet(&mut self, output: &mut Vec<Event<'static>>, packet: ClientPacket) -> Result<bool> {
         let res = match packet {
-            proto::ClientPacket::Iq(iq) => {
+            ClientPacket::Iq(iq) => {
                 self.handle_iq(output, iq).await;
                 false
             }
@@ -427,11 +424,11 @@ impl<'a> XmppConnection<'a> {
                 self.handle_message(output, m).await?;
                 false
             }
-            proto::ClientPacket::Presence(p) => {
+            ClientPacket::Presence(p) => {
                 self.handle_presence(output, p).await?;
                 false
             }
-            proto::ClientPacket::StreamEnd => {
+            ClientPacket::StreamEnd => {
                 ServerStreamEnd.serialize(output);
                 true
             }
@@ -439,25 +436,3 @@ impl<'a> XmppConnection<'a> {
         Ok(res)
     }
 }
-
-async fn read_xml_header(
-    xml_reader: &mut NsReader<(impl AsyncBufRead + Unpin)>,
-    reader_buf: &mut Vec<u8>,
-) -> Result<()> {
-    if let Event::Decl(bytes) = xml_reader.read_event_into_async(reader_buf).await? {
-        // this is <?xml ...> header
-        if let Some(encoding) = bytes.encoding() {
-            let encoding = encoding?;
-            if &*encoding == b"UTF-8" {
-                Ok(())
-            } else {
-                Err(anyhow!("Unsupported encoding: {encoding:?}"))
-            }
-        } else {
-            // Err(fail("No XML encoding provided"))
-            Ok(())
-        }
-    } else {
-        Err(anyhow!("Expected XML header"))
-    }
-}
diff --git a/crates/proto-xmpp/src/stream.rs b/crates/proto-xmpp/src/stream.rs
index d85df07..8f46f31 100644
--- a/crates/proto-xmpp/src/stream.rs
+++ b/crates/proto-xmpp/src/stream.rs
@@ -24,7 +24,17 @@ impl ClientStreamStart {
         reader: &mut NsReader<impl AsyncBufRead + Unpin>,
         buf: &mut Vec<u8>,
     ) -> Result<ClientStreamStart> {
-        let incoming = skip_text!(reader, buf);
+        let mut incoming = skip_text!(reader, buf);
+        if let Event::Decl(bytes) = incoming {
+            // this is <?xml ...> header
+            if let Some(encoding) = bytes.encoding() {
+                let encoding = encoding?;
+                if &*encoding != b"UTF-8" {
+                    return Err(anyhow!("Unsupported encoding: {encoding:?}"));
+                }
+            }
+            incoming = skip_text!(reader, buf);
+        }
         if let Event::Start(e) = incoming {
             let (ns, local) = reader.resolve_element(e.name());
             if ns != ResolveResult::Bound(Namespace(XMLNS.as_bytes())) {