feat(xmpp): refactor the projection

src/projections/xmpp/mod.rs

@@ -11,7 +11,7 @@ use quick_xml::events::{BytesDecl, Event};
 use quick_xml::{NsReader, Writer};
 use rustls_pemfile::{certs, rsa_private_keys};
 use serde::Deserialize;
-use tokio::io::{AsyncWriteExt, BufReader, BufWriter};
+use tokio::io::{AsyncBufRead, AsyncWrite, AsyncWriteExt, BufReader, BufWriter};
 use tokio::net::{TcpListener, TcpStream};
 use tokio::sync::mpsc::channel;
 use tokio_rustls::rustls::{Certificate, PrivateKey};
@@ -21,6 +21,7 @@ use crate::core::player::PlayerRegistry;
 use crate::core::room::RoomRegistry;
 use crate::prelude::*;
 use crate::protos::xmpp;
+use crate::protos::xmpp::stream::*;
 use crate::util::Terminator;
 
 #[derive(Deserialize, Debug, Clone)]
@@ -129,48 +130,13 @@ async fn handle_socket(
     rooms: RoomRegistry,
     termination: Deferred<()>, // TODO use it to stop the connection gracefully
 ) -> Result<()> {
-    use xmpp::stream::*;
     log::debug!("Received an XMPP connection from {socket_addr}");
     let mut reader_buf = vec![];
     let (reader, writer) = stream.split();
     let mut buf_reader = BufReader::new(reader);
     let mut buf_writer = BufWriter::new(writer);
 
-    {
+    socket_force_tls(&mut buf_reader, &mut buf_writer, &mut reader_buf).await?;
-        let mut xml_reader = NsReader::from_reader(&mut buf_reader);
-        let mut xml_writer = Writer::new(&mut buf_writer);
-        let aaa = xml_reader.read_event_into_async(&mut reader_buf).await?;
-        if let Event::Decl(_) = aaa {
-            // this is <?xml ...> header
-        } else {
-            return Err(fail("expected XML header"));
-        }
-        let _ = ClientStreamStart::parse(&mut xml_reader, &mut reader_buf).await?;
-
-        xml_writer
-            .write_event_async(Event::Decl(BytesDecl::new("1.0", None, None)))
-            .await?;
-        xmpp::stream::ServerStreamStart {
-            from: "localhost".into(),
-            lang: "en".into(),
-            version: "1.0".into(),
-        }
-        .write_xml(&mut xml_writer)
-        .await?;
-        xmpp::stream::Features {
-            start_tls: true,
-            mechanisms: false,
-            bind: false,
-        }
-        .write_xml(&mut xml_writer)
-        .await?;
-        xml_writer.get_mut().flush().await?;
-
-        let StartTLS = StartTLS::parse(&mut xml_reader, &mut reader_buf).await?;
-        // TODO read <starttls xmlns="urn:ietf:params:xml:ns:xmpp-tls"/>
-        xmpp::stream::Proceed.write_xml(&mut xml_writer).await?;
-        xml_writer.get_mut().flush().await?;
-    }
 
     let config = tokio_rustls::rustls::ServerConfig::builder()
         .with_safe_defaults()
@@ -182,68 +148,12 @@ async fn handle_socket(
     log::debug!("TLS connection established");
 
     let (a, b) = tokio::io::split(new_stream);
-    let buf_reader = BufReader::new(a);
+    let mut xml_reader = NsReader::from_reader(BufReader::new(a));
-
-    let mut xml_reader = NsReader::from_reader(buf_reader);
     let mut xml_writer = Writer::new(b);
 
-    {
+    socket_auth(&mut xml_reader, &mut xml_writer, &mut reader_buf).await?;
-        if let Event::Decl(_) = xml_reader.read_event_into_async(&mut reader_buf).await? {
+    socket_final(&mut xml_reader, &mut xml_writer, &mut reader_buf).await?;
-            // this is <?xml ...> header
-        } else {
-            return Err(fail("expected XML header"));
-        }
-        let _ = ClientStreamStart::parse(&mut xml_reader, &mut reader_buf).await?;
 
-        xml_writer
-            .write_event_async(Event::Decl(BytesDecl::new("1.0", None, None)))
-            .await?;
-        xmpp::stream::ServerStreamStart {
-            from: "localhost".into(),
-            lang: "en".into(),
-            version: "1.0".into(),
-        }
-        .write_xml(&mut xml_writer)
-        .await?;
-        xmpp::stream::Features {
-            start_tls: false,
-            mechanisms: true,
-            bind: false,
-        }
-        .write_xml(&mut xml_writer)
-        .await?;
-        xml_writer.get_mut().flush().await?;
-
-        let _ = xmpp::sasl::Auth::parse(&mut xml_reader, &mut reader_buf).await?;
-        xmpp::sasl::Success.write_xml(&mut xml_writer).await?;
-    }
-    {
-        if let Event::Decl(_) = xml_reader.read_event_into_async(&mut reader_buf).await? {
-            // this is <?xml ...> header
-        } else {
-            return Err(fail("expected XML header"));
-        }
-        let _ = ClientStreamStart::parse(&mut xml_reader, &mut reader_buf).await?;
-
-        xml_writer
-            .write_event_async(Event::Decl(BytesDecl::new("1.0", None, None)))
-            .await?;
-        xmpp::stream::ServerStreamStart {
-            from: "localhost".into(),
-            lang: "en".into(),
-            version: "1.0".into(),
-        }
-        .write_xml(&mut xml_writer)
-        .await?;
-        xmpp::stream::Features {
-            start_tls: false,
-            mechanisms: false,
-            bind: true,
-        }
-        .write_xml(&mut xml_writer)
-        .await?;
-        xml_writer.get_mut().flush().await?;
-    }
     loop {
         let event = xml_reader.read_event_into_async(&mut reader_buf).await?;
         println!("EVENT: {event:?}");
@@ -257,3 +167,118 @@ async fn handle_socket(
     a.unsplit(b).shutdown().await?;
     Ok(())
 }
+
+async fn socket_force_tls(
+    reader: &mut (impl AsyncBufRead + Unpin),
+    writer: &mut (impl AsyncWrite + Unpin),
+    reader_buf: &mut Vec<u8>,
+) -> Result<()> {
+    let mut xml_reader = &mut NsReader::from_reader(reader);
+    let mut 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));
+    xml_writer.write_event_async(event).await?;
+    let msg = ServerStreamStart {
+        from: "localhost".into(),
+        lang: "en".into(),
+        version: "1.0".into(),
+    };
+    msg.write_xml(xml_writer).await?;
+    let msg = Features {
+        start_tls: true,
+        mechanisms: false,
+        bind: false,
+    };
+    msg.write_xml(xml_writer).await?;
+    xml_writer.get_mut().flush().await?;
+
+    let StartTLS = StartTLS::parse(xml_reader, reader_buf).await?;
+    Proceed.write_xml(xml_writer).await?;
+    xml_writer.get_mut().flush().await?;
+    Ok(())
+}
+
+async fn socket_auth(
+    xml_reader: &mut NsReader<(impl AsyncBufRead + Unpin)>,
+    xml_writer: &mut Writer<(impl AsyncWrite + Unpin)>,
+    reader_buf: &mut Vec<u8>,
+) -> 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?;
+    ServerStreamStart {
+        from: "localhost".into(),
+        lang: "en".into(),
+        version: "1.0".into(),
+    }
+    .write_xml(xml_writer)
+    .await?;
+    Features {
+        start_tls: false,
+        mechanisms: true,
+        bind: false,
+    }
+    .write_xml(xml_writer)
+    .await?;
+    xml_writer.get_mut().flush().await?;
+
+    let _ = xmpp::sasl::Auth::parse(xml_reader, reader_buf).await?;
+    xmpp::sasl::Success.write_xml(xml_writer).await?;
+    Ok(())
+}
+
+async fn socket_final(
+    xml_reader: &mut NsReader<(impl AsyncBufRead + Unpin)>,
+    xml_writer: &mut Writer<(impl AsyncWrite + Unpin)>,
+    reader_buf: &mut Vec<u8>,
+) -> 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?;
+    ServerStreamStart {
+        from: "localhost".into(),
+        lang: "en".into(),
+        version: "1.0".into(),
+    }
+    .write_xml(xml_writer)
+    .await?;
+    Features {
+        start_tls: false,
+        mechanisms: false,
+        bind: true,
+    }
+    .write_xml(xml_writer)
+    .await?;
+    xml_writer.get_mut().flush().await?;
+    Ok(())
+}
+
+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(fail(format!("Unsupported encoding: {encoding:?}").as_str()))
+            }
+        } else {
+            // Err(fail("No XML encoding provided"))
+            Ok(())
+        }
+    } else {
+        Err(fail("Expected XML header"))
+    }
+}