feat(xmpp): handle stream stop, dump tls keys

This commit is contained in:
Nikita Vilunov 2023-04-05 18:57:35 +02:00
parent fb8329a187
commit f71d098420
3 changed files with 53 additions and 27 deletions

View File

@ -148,10 +148,11 @@ async fn handle_socket(
socket_force_tls(&mut buf_reader, &mut buf_writer, &mut reader_buf).await?; socket_force_tls(&mut buf_reader, &mut buf_writer, &mut reader_buf).await?;
let config = tokio_rustls::rustls::ServerConfig::builder() let mut config = tokio_rustls::rustls::ServerConfig::builder()
.with_safe_defaults() .with_safe_defaults()
.with_no_client_auth() .with_no_client_auth()
.with_single_cert(vec![config.cert.clone()], config.key.clone())?; .with_single_cert(vec![config.cert.clone()], config.key.clone())?;
config.key_log = Arc::new(tokio_rustls::rustls::KeyLogFile::new());
let acceptor = TlsAcceptor::from(Arc::new(config)); let acceptor = TlsAcceptor::from(Arc::new(config));
let new_stream = acceptor.accept(stream).await?; let new_stream = acceptor.accept(stream).await?;
@ -164,14 +165,6 @@ async fn handle_socket(
socket_auth(&mut xml_reader, &mut xml_writer, &mut reader_buf).await?; socket_auth(&mut xml_reader, &mut xml_writer, &mut reader_buf).await?;
socket_final(&mut xml_reader, &mut xml_writer, &mut reader_buf).await?; socket_final(&mut xml_reader, &mut xml_writer, &mut reader_buf).await?;
loop {
let event = xml_reader.read_event_into_async(&mut reader_buf).await?;
println!("EVENT: {event:?}");
if event == Event::Eof {
break;
}
}
let a = xml_reader.into_inner().into_inner(); let a = xml_reader.into_inner().into_inner();
let b = xml_writer.into_inner(); let b = xml_writer.into_inner();
a.unsplit(b).shutdown().await?; a.unsplit(b).shutdown().await?;
@ -289,22 +282,29 @@ async fn socket_final(
Continuation::Final(res) => { Continuation::Final(res) => {
let res = res?; let res = res?;
dbg!(&res); dbg!(&res);
handle_packet(&mut events, res); let stop = handle_packet(&mut events, res);
for i in &events { for i in &events {
xml_writer.write_event_async(i).await?; xml_writer.write_event_async(i).await?;
} }
events.clear(); events.clear();
xml_writer.get_mut().flush().await?; xml_writer.get_mut().flush().await?;
if stop {
break;
}
parser = proto::ClientPacket::parse(); parser = proto::ClientPacket::parse();
} }
Continuation::Continue(p) => parser = p, Continuation::Continue(p) => parser = p,
} }
} }
Ok(())
} }
fn handle_packet(output: &mut Vec<Event<'static>>, packet: ClientPacket) { fn handle_packet(output: &mut Vec<Event<'static>>, packet: ClientPacket) -> bool {
match packet { match packet {
proto::ClientPacket::Iq(iq) => handle_iq(output, iq), proto::ClientPacket::Iq(iq) => {
handle_iq(output, iq);
false
}
proto::ClientPacket::Message(_) => todo!(), proto::ClientPacket::Message(_) => todo!(),
proto::ClientPacket::Presence(p) => { proto::ClientPacket::Presence(p) => {
let response = Presence::<()> { let response = Presence::<()> {
@ -313,6 +313,11 @@ fn handle_packet(output: &mut Vec<Event<'static>>, packet: ClientPacket) {
..Default::default() ..Default::default()
}; };
response.serialize(output); response.serialize(output);
false
},
proto::ClientPacket::StreamEnd => {
ServerStreamEnd.serialize(output);
true
} }
} }
} }

View File

@ -51,6 +51,7 @@ pub enum ClientPacket {
Iq(Iq<IqClientBody>), Iq(Iq<IqClientBody>),
Message(Message), Message(Message),
Presence(Presence<Ignore>), Presence(Presence<Ignore>),
StreamEnd,
} }
impl FromXml for ClientPacket { impl FromXml for ClientPacket {
@ -58,22 +59,34 @@ impl FromXml for ClientPacket {
fn parse() -> Self::P { fn parse() -> Self::P {
|(namespace, event): (ResolveResult<'static>, &'static Event<'static>)| -> Result<Self> { |(namespace, event): (ResolveResult<'static>, &'static Event<'static>)| -> Result<Self> {
let Event::Start(bytes) = event else { match event {
return Err(ffail!("Unexpected XML event: {event:?}")); Event::Start(bytes) => {
}; let name = bytes.name();
let name = bytes.name(); match_parser!(name, namespace, event;
match_parser!(name, namespace, event; Iq::<IqClientBody>,
Iq::<IqClientBody>, Presence::<Ignore>,
Presence::<Ignore>, Message,
Message, {
{ Err(ffail!(
Err(ffail!( "Unexpected XML event of name {:?} in namespace {:?}",
"Unexpected XML event of name {:?} in namespace {:?}", name,
name, namespace
namespace ))
)) }
)
} }
) Event::End(bytes) => {
let name = bytes.name();
if name.local_name().as_ref() == b"stream" {
return Ok(ClientPacket::StreamEnd);
} else {
return Err(ffail!("Unexpected XML event: {event:?}"));
}
}
_ => {
return Err(ffail!("Unexpected XML event: {event:?}"));
}
}
} }
} }
} }

View File

@ -6,6 +6,7 @@ use tokio::io::{AsyncBufRead, AsyncWrite};
use super::skip_text; use super::skip_text;
use crate::prelude::*; use crate::prelude::*;
use crate::util::xml::ToXml;
pub static XMLNS: &'static str = "http://etherx.jabber.org/streams"; pub static XMLNS: &'static str = "http://etherx.jabber.org/streams";
pub static PREFIX: &'static str = "stream"; pub static PREFIX: &'static str = "stream";
@ -109,6 +110,13 @@ impl ServerStreamStart {
} }
} }
pub struct ServerStreamEnd;
impl ToXml for ServerStreamEnd {
fn serialize(&self, events: &mut Vec<Event<'static>>) {
events.push(Event::End(BytesEnd::new("stream:stream")));
}
}
pub struct Features { pub struct Features {
pub start_tls: bool, pub start_tls: bool,
pub mechanisms: bool, pub mechanisms: bool,