From d95838b6ed183dba5aadb8b938d4f38eb70f1954 Mon Sep 17 00:00:00 2001 From: moparisthebest Date: Wed, 14 Apr 2021 02:09:00 -0400 Subject: [PATCH] Fix ejabberd s2s --- src/main.rs | 51 ++++++++++++++++++++++++----------------- src/slicesubsequence.rs | 36 +++++++++++++++++------------ 2 files changed, 51 insertions(+), 36 deletions(-) diff --git a/src/main.rs b/src/main.rs index 31f5e60..0c5a493 100644 --- a/src/main.rs +++ b/src/main.rs @@ -103,8 +103,10 @@ impl Config { } } -fn to_str(buf: &[u8]) -> &str { - std::str::from_utf8(buf).unwrap_or("[invalid utf-8]") +fn to_str(buf: &[u8]) -> std::borrow::Cow<'_, str> { + //&str { + //std::str::from_utf8(buf).unwrap_or("[invalid utf-8]") + String::from_utf8_lossy(buf) } async fn handle_connection(mut stream: tokio::net::TcpStream, client_addr: SocketAddr, local_addr: SocketAddr, config: CloneableConfig) -> Result<()> { @@ -148,7 +150,11 @@ async fn handle_connection(mut stream: tokio::net::TcpStream, client_addr: Socke if !direct_tls { let mut stream_open = Vec::new(); - while let Ok(n) = stream.read(in_filter.current_buf()).await { + let (in_rd, mut in_wr) = stream.split(); + // we naively read 1 byte at a time, which buffering significantly speeds up + let mut in_rd = tokio::io::BufReader::with_capacity(IN_BUFFER_SIZE, in_rd); + + while let Ok(n) = in_rd.read(in_filter.current_buf()).await { if n == 0 { bail!("stream ended before open"); } @@ -160,32 +166,36 @@ async fn handle_connection(mut stream: tokio::net::TcpStream, client_addr: Socke continue; } else if buf.starts_with(b" {} '{}'", client_addr, to_str(&stream_open)); - stream.write_all(&stream_open).await?; + in_wr.write_all(&stream_open).await?; stream_open.clear(); // gajim seems to REQUIRE an id here... let buf = if buf.contains_seq(b"id=") { - buf.replace(b" id='", b" id='xmpp-proxy") - .replace(br#" id=""#, br#" id="xmpp-proxy"#) - .replace(b" to=", br#" bla toblala="#) - .replace(b" from=", b" to=") - .replace(br#" bla toblala="#, br#" from="#) + buf.replace_first(b" id='", b" id='xmpp-proxy") + .replace_first(br#" id=""#, br#" id="xmpp-proxy"#) + .replace_first(b" to=", br#" bla toblala="#) + .replace_first(b" from=", b" to=") + .replace_first(br#" bla toblala="#, br#" from="#) } else { - buf.replace(b" to=", br#" bla toblala="#) - .replace(b" from=", b" to=") - .replace(br#" bla toblala="#, br#" id='xmpp-proxy' from="#) + buf.replace_first(b" to=", br#" bla toblala="#) + .replace_first(b" from=", b" to=") + .replace_first(br#" bla toblala="#, br#" id='xmpp-proxy' from="#) }; debug!("> {} '{}'", client_addr, to_str(&buf)); - stream.write_all(&buf).await?; + in_wr.write_all(&buf).await?; - stream - .write_all(br###""###) - .await?; - stream.flush().await?; + // ejabberd never sends with the first, only the second? + //let buf = br###""###; + let buf = br###""###; + debug!("> {} '{}'", client_addr, to_str(buf)); + in_wr.write_all(buf).await?; + in_wr.flush().await?; } else if buf.starts_with(b""###).await?; - stream.flush().await?; + let buf = br###""###; + debug!("> {} '{}'", client_addr, to_str(buf)); + in_wr.write_all(buf).await?; + in_wr.flush().await?; break; } else { bail!("bad pre-tls stanza: {}", to_str(&buf)); @@ -197,7 +207,6 @@ async fn handle_connection(mut stream: tokio::net::TcpStream, client_addr: Socke let stream = config.acceptor.accept(stream).await?; let (in_rd, mut in_wr) = tokio::io::split(stream); - // we naively read 1 byte at a time, which buffering significantly speeds up let mut in_rd = tokio::io::BufReader::with_capacity(IN_BUFFER_SIZE, in_rd); @@ -396,7 +405,7 @@ impl StanzaFilter { //println!("b: '{}', cnt: {}, tag_cnt: {}, self.buf.len(): {}", b as char, self.cnt, self.tag_cnt, self.buf.len()); self.cnt += 1; if self.cnt == self.buf_size { - bail!("stanza too big"); + bail!("stanza too big: {}", to_str(&self.buf)); } Ok(None) } diff --git a/src/slicesubsequence.rs b/src/slicesubsequence.rs index d6cb8e8..a047bf6 100644 --- a/src/slicesubsequence.rs +++ b/src/slicesubsequence.rs @@ -3,7 +3,7 @@ use anyhow::{anyhow, Result}; pub trait SliceSubsequence { fn trim_start(&self, needle: &[T]) -> &[T]; fn first_index_of(&self, needle: &[T]) -> Result; - fn replace(self, needle: &[T], replacement: &[T]) -> Vec; + fn replace_first(self, needle: &[T], replacement: &[T]) -> Vec; fn contains_seq(&self, needle: &[T]) -> bool { self.first_index_of(needle).is_ok() @@ -36,8 +36,8 @@ impl SliceSubsequence for &[T] { Err(anyhow!("not found")) } - fn replace(self, needle: &[T], replacement: &[T]) -> Vec { - self.to_vec().replace(needle, replacement) + fn replace_first(self, needle: &[T], replacement: &[T]) -> Vec { + self.to_vec().replace_first(needle, replacement) } } @@ -50,11 +50,17 @@ impl SliceSubsequence for Vec { return (self.as_slice()).first_index_of(needle); } - fn replace(mut self, needle: &[T], replacement: &[T]) -> Vec { - while let Ok(idx) = self.first_index_of(needle) { - self.splice(idx..(idx + needle.len()), replacement.iter().cloned()); + fn replace_first(self, needle: &[T], replacement: &[T]) -> Vec { + if let Ok(idx) = self.first_index_of(needle) { + let second = &self[(idx + needle.len())..]; + let mut ret = Vec::with_capacity(idx + replacement.len() + second.len()); + ret.extend_from_slice(&self[..idx]); + ret.extend_from_slice(replacement); + ret.extend_from_slice(second); + ret + } else { + self } - self } } @@ -74,23 +80,23 @@ mod tests { assert_eq!(buf, b"bla"); } #[test] - fn replace() { - let buf = b"bla to='tsnhaou' bla2".replace(b" to=", b" from="); + fn replace_first() { + let buf = b"bla to='tsnhaou' bla2".replace_first(b" to=", b" from="); assert_eq!(buf, b"bla from='tsnhaou' bla2"); - let buf = buf.replace(b" from=", b" to="); + let buf = buf.replace_first(b" from=", b" to="); assert_eq!(buf, b"bla to='tsnhaou' bla2"); - let buf = buf.replace(b" to=", b" from="); + let buf = buf.replace_first(b" to=", b" from="); assert_eq!(buf, b"bla from='tsnhaou' bla2"); - let buf = b"bla to='tsnhaou' bla2".replace(b"bla", b"boo"); - assert_eq!(buf, b"boo to='tsnhaou' boo2"); + let buf = b"bla to='tsnhaou' bla2".replace_first(b"bla", b"boo"); + assert_eq!(buf, b"boo to='tsnhaou' bla2"); - let buf = buf.replace(b"boo", b"bla"); + let buf = buf.replace_first(b"boo", b"bla"); assert_eq!(buf, b"bla to='tsnhaou' bla2"); - let buf = buf.replace(b" bla2", b""); + let buf = buf.replace_first(b" bla2", b""); assert_eq!(buf, b"bla to='tsnhaou'"); } }