lavina/crates/projection-xmpp/tests/lib.rs

use std::io::ErrorKind;
use std::sync::Arc;
use std::time::Duration;

use anyhow::Result;
use assert_matches::*;
use prometheus::Registry as MetricsRegistry;
use quick_xml::events::Event;
use quick_xml::NsReader;
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
use tokio::io::{ReadHalf as GenericReadHalf, WriteHalf as GenericWriteHalf};
use tokio::net::tcp::{ReadHalf, WriteHalf};
use tokio::net::TcpStream;
use tokio_rustls::client::TlsStream;
use tokio_rustls::rustls::client::ServerCertVerifier;
use tokio_rustls::rustls::{ClientConfig, ServerName};
use tokio_rustls::TlsConnector;

use lavina_core::player::PlayerRegistry;
use lavina_core::repo::{Storage, StorageConfig};
use lavina_core::room::RoomRegistry;
use projection_xmpp::{launch, RunningServer, ServerConfig};
use proto_xmpp::xml::{Continuation, FromXml, Parser};

pub async fn read_irc_message(reader: &mut BufReader<ReadHalf<'_>>, buf: &mut Vec<u8>) -> Result<usize> {
    let mut size = 0;
    let res = reader.read_until(b'\n', buf).await?;
    size += res;
    return Ok(size);
}

struct TestScope<'a> {
    reader: NsReader<BufReader<ReadHalf<'a>>>,
    writer: WriteHalf<'a>,
    buffer: Vec<u8>,
}

impl<'a> TestScope<'a> {
    fn new(stream: &mut TcpStream) -> TestScope<'_> {
        let (reader, writer) = stream.split();
        let reader = NsReader::from_reader(BufReader::new(reader));
        let buffer = vec![];
        TestScope { reader, writer, buffer }
    }

    async fn send(&mut self, str: &str) -> Result<()> {
        self.writer.write_all(str.as_bytes()).await?;
        self.writer.flush().await?;
        Ok(())
    }

    async fn next_xml_event(&mut self) -> Result<Event<'_>> {
        self.buffer.clear();
        let event = self.reader.read_event_into_async(&mut self.buffer).await?;
        Ok(event)
    }

    async fn read<T: FromXml>(&mut self) -> Result<T> {
        self.buffer.clear();
        let (ns, event) = self.reader.read_resolved_event_into_async(&mut self.buffer).await?;
        let mut parser: Continuation<_, std::result::Result<T, anyhow::Error>> = T::parse().consume(ns, &event);
        loop {
            match parser {
                Continuation::Final(res) => return Ok(res?),
                Continuation::Continue(next) => {
                    let (ns, event) = self.reader.read_resolved_event_into_async(&mut self.buffer).await?;
                    parser = next.consume(ns, &event);
                }
            }
        }
    }
}

struct TestScopeTls<'a> {
    reader: NsReader<BufReader<GenericReadHalf<&'a mut TlsStream<TcpStream>>>>,
    writer: GenericWriteHalf<&'a mut TlsStream<TcpStream>>,
    buffer: Vec<u8>,
    pub timeout: Duration,
}

impl<'a> TestScopeTls<'a> {
    fn new(stream: &'a mut TlsStream<TcpStream>, buffer: Vec<u8>) -> TestScopeTls<'a> {
        let (reader, writer) = tokio::io::split(stream);
        let reader = NsReader::from_reader(BufReader::new(reader));
        let timeout = Duration::from_millis(100);

        TestScopeTls {
            reader,
            writer,
            buffer,
            timeout,
        }
    }

    async fn send(&mut self, str: &str) -> Result<()> {
        self.writer.write_all(str.as_bytes()).await?;
        self.writer.flush().await?;
        Ok(())
    }

    async fn next_xml_event(&mut self) -> Result<Event<'_>> {
        self.buffer.clear();
        let event = self.reader.read_event_into_async(&mut self.buffer);
        let event = tokio::time::timeout(self.timeout, event).await??;
        Ok(event)
    }
}

struct IgnoreCertVerification;
impl ServerCertVerifier for IgnoreCertVerification {
    fn verify_server_cert(
        &self,
        _end_entity: &tokio_rustls::rustls::Certificate,
        _intermediates: &[tokio_rustls::rustls::Certificate],
        _server_name: &ServerName,
        _scts: &mut dyn Iterator<Item = &[u8]>,
        _ocsp_response: &[u8],
        _now: std::time::SystemTime,
    ) -> std::result::Result<tokio_rustls::rustls::client::ServerCertVerified, tokio_rustls::rustls::Error> {
        Ok(tokio_rustls::rustls::client::ServerCertVerified::assertion())
    }
}

struct TestServer {
    metrics: MetricsRegistry,
    storage: Storage,
    rooms: RoomRegistry,
    players: PlayerRegistry,
    server: RunningServer,
}
impl TestServer {
    async fn start() -> Result<TestServer> {
        let _ = tracing_subscriber::fmt::try_init();
        let config = ServerConfig {
            listen_on: "127.0.0.1:0".parse().unwrap(),
            cert: "tests/certs/xmpp.pem".parse().unwrap(),
            key: "tests/certs/xmpp.key".parse().unwrap(),
            hostname: "localhost".into(),
        };
        let mut metrics = MetricsRegistry::new();
        let mut storage = Storage::open(StorageConfig {
            db_path: ":memory:".into(),
        })
        .await?;
        let rooms = RoomRegistry::new(&mut metrics, storage.clone()).unwrap();
        let players = PlayerRegistry::empty(rooms.clone(), &mut metrics).unwrap();
        let server = launch(config, players.clone(), rooms.clone(), metrics.clone(), storage.clone()).await.unwrap();
        Ok(TestServer {
            metrics,
            storage,
            rooms,
            players,
            server,
        })
    }
}

#[tokio::test]
async fn scenario_basic() -> Result<()> {
    let mut server = TestServer::start().await?;

    // test scenario

    server.storage.create_user("tester").await?;
    server.storage.set_password("tester", "password").await?;

    let mut stream = TcpStream::connect(server.server.addr).await?;
    let mut s = TestScope::new(&mut stream);
    tracing::info!("TCP connection established");

    s.send(r#"<?xml version="1.0"?>"#).await?;
    s.send(r#"<stream:stream xmlns:stream="http://etherx.jabber.org/streams" to="127.0.0.1" xml:lang="en" xmlns:xml="http://www.w3.org/XML/1998/namespace" xmlns="jabber:client" version="1.0">"#).await?;
    assert_matches!(s.next_xml_event().await?, Event::Decl(_) => {});
    assert_matches!(s.next_xml_event().await?, Event::Start(b) => assert_eq!(b.local_name().into_inner(), b"stream"));
    assert_matches!(s.next_xml_event().await?, Event::Start(b) => assert_eq!(b.local_name().into_inner(), b"features"));
    assert_matches!(s.next_xml_event().await?, Event::Start(b) => assert_eq!(b.local_name().into_inner(), b"starttls"));
    assert_matches!(s.next_xml_event().await?, Event::Empty(b) => assert_eq!(b.local_name().into_inner(), b"required"));
    assert_matches!(s.next_xml_event().await?, Event::End(b) => assert_eq!(b.local_name().into_inner(), b"starttls"));
    assert_matches!(s.next_xml_event().await?, Event::End(b) => assert_eq!(b.local_name().into_inner(), b"features"));
    s.send(r#"<starttls/>"#).await?;
    assert_matches!(s.next_xml_event().await?, Event::Empty(b) => assert_eq!(b.local_name().into_inner(), b"proceed"));
    let buffer = s.buffer;
    tracing::info!("TLS feature negotiation complete");

    let connector = TlsConnector::from(Arc::new(
        ClientConfig::builder()
            .with_safe_defaults()
            .with_custom_certificate_verifier(Arc::new(IgnoreCertVerification))
            .with_no_client_auth(),
    ));
    tracing::info!("Initiating TLS connection...");
    let mut stream = connector.connect(ServerName::IpAddress(server.server.addr.ip()), stream).await?;
    tracing::info!("TLS connection established");

    let mut s = TestScopeTls::new(&mut stream, buffer);

    s.send(r#"<?xml version="1.0"?>"#).await?;
    s.send(r#"<stream:stream xmlns:stream="http://etherx.jabber.org/streams" to="127.0.0.1" xml:lang="en" xmlns:xml="http://www.w3.org/XML/1998/namespace" xmlns="jabber:client" version="1.0">"#).await?;
    assert_matches!(s.next_xml_event().await?, Event::Decl(_) => {});
    assert_matches!(s.next_xml_event().await?, Event::Start(b) => assert_eq!(b.local_name().into_inner(), b"stream"));

    stream.shutdown().await?;

    // wrap up

    server.server.terminate().await?;
    Ok(())
}

#[tokio::test]
async fn scenario_basic_without_headers() -> Result<()> {
    let mut server = TestServer::start().await?;

    // test scenario

    server.storage.create_user("tester").await?;
    server.storage.set_password("tester", "password").await?;

    let mut stream = TcpStream::connect(server.server.addr).await?;
    let mut s = TestScope::new(&mut stream);
    tracing::info!("TCP connection established");

    s.send(r#"<stream:stream xmlns:stream="http://etherx.jabber.org/streams" to="127.0.0.1" xml:lang="en" xmlns:xml="http://www.w3.org/XML/1998/namespace" xmlns="jabber:client" version="1.0">"#).await?;
    assert_matches!(s.next_xml_event().await?, Event::Decl(_) => {});
    assert_matches!(s.next_xml_event().await?, Event::Start(b) => assert_eq!(b.local_name().into_inner(), b"stream"));
    assert_matches!(s.next_xml_event().await?, Event::Start(b) => assert_eq!(b.local_name().into_inner(), b"features"));
    assert_matches!(s.next_xml_event().await?, Event::Start(b) => assert_eq!(b.local_name().into_inner(), b"starttls"));
    assert_matches!(s.next_xml_event().await?, Event::Empty(b) => assert_eq!(b.local_name().into_inner(), b"required"));
    assert_matches!(s.next_xml_event().await?, Event::End(b) => assert_eq!(b.local_name().into_inner(), b"starttls"));
    assert_matches!(s.next_xml_event().await?, Event::End(b) => assert_eq!(b.local_name().into_inner(), b"features"));
    s.send(r#"<starttls/>"#).await?;
    assert_matches!(s.next_xml_event().await?, Event::Empty(b) => assert_eq!(b.local_name().into_inner(), b"proceed"));
    let buffer = s.buffer;
    tracing::info!("TLS feature negotiation complete");

    let connector = TlsConnector::from(Arc::new(
        ClientConfig::builder()
            .with_safe_defaults()
            .with_custom_certificate_verifier(Arc::new(IgnoreCertVerification))
            .with_no_client_auth(),
    ));
    tracing::info!("Initiating TLS connection...");
    let mut stream = connector.connect(ServerName::IpAddress(server.server.addr.ip()), stream).await?;
    tracing::info!("TLS connection established");

    let mut s = TestScopeTls::new(&mut stream, buffer);

    s.send(r#"<stream:stream xmlns:stream="http://etherx.jabber.org/streams" to="127.0.0.1" xml:lang="en" xmlns:xml="http://www.w3.org/XML/1998/namespace" xmlns="jabber:client" version="1.0">"#).await?;
    assert_matches!(s.next_xml_event().await?, Event::Decl(_) => {});
    assert_matches!(s.next_xml_event().await?, Event::Start(b) => assert_eq!(b.local_name().into_inner(), b"stream"));

    stream.shutdown().await?;

    // wrap up

    server.server.terminate().await?;
    Ok(())
}

#[tokio::test]
async fn terminate_socket() -> Result<()> {
    let mut server = TestServer::start().await?;

    // test scenario

    server.storage.create_user("tester").await?;
    server.storage.set_password("tester", "password").await?;

    let mut stream = TcpStream::connect(server.server.addr).await?;
    let mut s = TestScope::new(&mut stream);
    tracing::info!("TCP connection established");

    s.send(r#"<?xml version="1.0"?>"#).await?;

    s.send(r#"<stream:stream xmlns:stream="http://etherx.jabber.org/streams" to="127.0.0.1" xml:lang="en" xmlns:xml="http://www.w3.org/XML/1998/namespace" xmlns="jabber:client" version="1.0">"#).await?;
    assert_matches!(s.next_xml_event().await?, Event::Decl(_) => {});
    assert_matches!(s.next_xml_event().await?, Event::Start(b) => assert_eq!(b.local_name().into_inner(), b"stream"));
    assert_matches!(s.next_xml_event().await?, Event::Start(b) => assert_eq!(b.local_name().into_inner(), b"features"));
    assert_matches!(s.next_xml_event().await?, Event::Start(b) => assert_eq!(b.local_name().into_inner(), b"starttls"));
    assert_matches!(s.next_xml_event().await?, Event::Empty(b) => assert_eq!(b.local_name().into_inner(), b"required"));
    assert_matches!(s.next_xml_event().await?, Event::End(b) => assert_eq!(b.local_name().into_inner(), b"starttls"));
    assert_matches!(s.next_xml_event().await?, Event::End(b) => assert_eq!(b.local_name().into_inner(), b"features"));
    s.send(r#"<starttls/>"#).await?;
    assert_matches!(s.next_xml_event().await?, Event::Empty(b) => assert_eq!(b.local_name().into_inner(), b"proceed"));

    let connector = TlsConnector::from(Arc::new(
        ClientConfig::builder()
            .with_safe_defaults()
            .with_custom_certificate_verifier(Arc::new(IgnoreCertVerification))
            .with_no_client_auth(),
    ));

    tracing::info!("Initiating TLS connection...");
    let mut stream = connector.connect(ServerName::IpAddress(server.server.addr.ip()), stream).await?;
    tracing::info!("TLS connection established");

    server.server.terminate().await?;

    assert_eq!(stream.read_u8().await.unwrap_err().kind(), ErrorKind::UnexpectedEof);

    Ok(())
}