use bytes::BytesMut;
use tracing::{trace, warn};
use crate::codec::decode::parse_response_utf8;
use crate::error::Error;
use crate::types::Response;
use super::ImapStream;
pub(crate) struct WireReader {
stream: ImapStream,
buf: BytesMut,
}
impl WireReader {
pub(super) fn new(stream: ImapStream) -> Self {
Self {
stream,
buf: BytesMut::with_capacity(8 * 1024),
}
}
pub(crate) async fn read_one(&mut self, utf8_mode: bool) -> Result<Response, Error> {
loop {
if let Some(resp) = self.try_parse_response_inner(utf8_mode)? {
return Ok(resp);
}
let n = self.stream.read_buf(&mut self.buf).await?;
if n == 0 {
return Err(Error::Closed);
}
}
}
pub(crate) async fn read_greeting(&mut self) -> Result<Response, Error> {
loop {
if let Some(resp) = self.try_parse_greeting_inner()? {
return Ok(resp);
}
let n = self.stream.read_buf(&mut self.buf).await?;
if n == 0 {
return Err(Error::Closed);
}
}
}
pub(crate) async fn write_all(&mut self, bytes: &[u8]) -> Result<(), Error> {
self.stream.write_all(bytes).await?;
self.stream.flush().await?;
Ok(())
}
pub(crate) fn buffer_is_empty(&self) -> bool {
self.buf.is_empty()
}
pub(in crate::connection) fn set_keepalive(
&self,
ka: &super::TcpKeepalive,
) -> Result<(), Error> {
self.stream.set_keepalive(ka)
}
pub(super) fn into_stream(self) -> ImapStream {
self.stream
}
pub(super) fn take_buffer(&mut self) -> bytes::BytesMut {
self.buf.split_off(0)
}
fn try_parse_greeting_inner(&mut self) -> Result<Option<Response>, Error> {
use crate::codec::decode::parse_greeting;
if self.buf.is_empty() {
return Ok(Option::None);
}
if !buffer_may_contain_complete_response(&self.buf) {
return Ok(Option::None);
}
match parse_greeting(&self.buf) {
Ok((remaining, resp)) => {
let consumed = self.buf.len() - remaining.len();
let _ = self.buf.split_to(consumed);
Ok(Some(resp))
}
Err(nom::Err::Incomplete(_)) => Ok(Option::None),
Err(e) => {
if !self.buf.ends_with(b"\r\n") {
trace!(
"greeting parse Error on buffer not ending with CRLF — \
treating as incomplete ({} bytes buffered)",
self.buf.len()
);
return Ok(Option::None);
}
warn!("greeting parse error: {e}");
Err(Error::Parse(format!("greeting parse error: {e}")))
}
}
}
fn try_parse_response_inner(&mut self, utf8_mode: bool) -> Result<Option<Response>, Error> {
if self.buf.is_empty() {
return Ok(Option::None);
}
if !buffer_may_contain_complete_response(&self.buf) {
return Ok(Option::None);
}
match parse_response_utf8(&self.buf, utf8_mode) {
Ok((remaining, resp)) => {
let consumed = self.buf.len() - remaining.len();
let _ = self.buf.split_to(consumed);
Ok(Some(resp))
}
Err(nom::Err::Incomplete(_)) => Ok(Option::None),
Err(e) => {
if !self.buf.ends_with(b"\r\n") {
trace!(
"parse Error on buffer not ending with CRLF — treating as incomplete \
({} bytes buffered)",
self.buf.len()
);
return Ok(Option::None);
}
warn!("response parse error: {e}");
Err(Error::Parse(format!("response parse error: {e}")))
}
}
}
}
pub(super) fn buffer_may_contain_complete_response(buf: &[u8]) -> bool {
let Some(first_crlf) = buf.windows(2).position(|w| w == b"\r\n") else {
return false;
};
let may_contain_literal = if buf.starts_with(b"* ") && buf.len() >= 5 {
let after_star = &buf[2..];
let is_status = after_star
.get(..3)
.is_some_and(|s| s.eq_ignore_ascii_case(b"OK ") || s.eq_ignore_ascii_case(b"NO "))
|| after_star.get(..4).is_some_and(|s| {
s.eq_ignore_ascii_case(b"BAD ") || s.eq_ignore_ascii_case(b"BYE ")
});
!is_status
} else if buf.starts_with(b"+") {
false
} else {
if let Some(sp) = buf.iter().position(|&b| b == b' ') {
let after_tag = &buf[sp + 1..];
let is_tagged_status = after_tag
.get(..3)
.is_some_and(|s| s.eq_ignore_ascii_case(b"OK ") || s.eq_ignore_ascii_case(b"NO "))
|| after_tag
.get(..4)
.is_some_and(|s| s.eq_ignore_ascii_case(b"BAD "));
let is_tagged_status = is_tagged_status
|| after_tag
.get(..4)
.is_some_and(|s| s.eq_ignore_ascii_case(b"OK\r\n"))
|| after_tag
.get(..4)
.is_some_and(|s| s.eq_ignore_ascii_case(b"NO\r\n"))
|| after_tag
.get(..5)
.is_some_and(|s| s.eq_ignore_ascii_case(b"BAD\r\n"));
!is_tagged_status
} else {
true
}
};
if first_crlf >= 2 && may_contain_literal {
if let Some(literal_len) = try_parse_literal_marker(buf, first_crlf) {
let mut pos = first_crlf + 2 + literal_len;
loop {
if pos > buf.len() {
return false;
}
let remaining = &buf[pos..];
let Some(next_crlf) = remaining.windows(2).position(|w| w == b"\r\n") else {
return false;
};
if let Some(next_literal_len) = try_parse_literal_marker(remaining, next_crlf) {
pos += next_crlf + 2 + next_literal_len;
continue;
}
return true;
}
}
}
true
}
fn try_parse_literal_marker(buf: &[u8], crlf_pos: usize) -> Option<usize> {
let before_crlf = &buf[..crlf_pos];
let brace_pos = before_crlf.iter().rposition(|&b| b == b'{')?;
let close_offset = buf[brace_pos + 1..crlf_pos]
.iter()
.position(|&b| b == b'}')?;
let between = &buf[brace_pos + 1..brace_pos + 1 + close_offset];
let digits = if between.last() == Some(&b'+') {
&between[..between.len() - 1]
} else {
between
};
if digits.is_empty() || !digits.iter().all(u8::is_ascii_digit) {
return None;
}
std::str::from_utf8(digits)
.ok()
.and_then(|s| s.parse::<usize>().ok())
}