diff --git a/src/connection.rs b/src/connection.rs index 9ef21f2..4d315a6 100644 --- a/src/connection.rs +++ b/src/connection.rs @@ -106,8 +106,8 @@ impl EppConnection { cx: &mut Context<'_>, ) -> Result { match &mut state { - RequestState::Writing { mut start, buf } => { - let wrote = match Pin::new(&mut self.stream).poll_write(cx, &buf[start..]) { + RequestState::Writing { start, buf } => { + let wrote = match Pin::new(&mut self.stream).poll_write(cx, &buf[*start..]) { Poll::Ready(Ok(wrote)) => wrote, Poll::Ready(Err(err)) => return Err(err.into()), Poll::Pending => return Ok(Transition::Pending(state)), @@ -121,7 +121,7 @@ impl EppConnection { .into()); } - start += wrote; + let start = *start + wrote; debug!( "{}: Wrote {} bytes, {} out of {} done", self.registry, @@ -133,7 +133,10 @@ impl EppConnection { // Transition to reading the response's frame header once // we've written the entire request if start < buf.len() { - return Ok(Transition::Next(state)); + return Ok(Transition::Next(RequestState::Writing { + start, + buf: mem::take(buf), + })); } Ok(Transition::Next(RequestState::ReadLength { @@ -141,8 +144,8 @@ impl EppConnection { buf: vec![0; 256], })) } - RequestState::ReadLength { mut read, buf } => { - let mut read_buf = ReadBuf::new(&mut buf[read..]); + RequestState::ReadLength { read, buf } => { + let mut read_buf = ReadBuf::new(&mut buf[*read..]); match Pin::new(&mut self.stream).poll_read(cx, &mut read_buf) { Poll::Ready(Ok(())) => {} Poll::Ready(Err(err)) => return Err(err.into()), @@ -162,26 +165,29 @@ impl EppConnection { // The frame header is a 32-bit (4-byte) big-endian unsigned integer. If we don't // have 4 bytes yet, stay in the `ReadLength` state, otherwise we transition to `Reading`. - read += filled.len(); - if read < 4 { - return Ok(Transition::Next(state)); + let new_read = *read + filled.len(); + if new_read < 4 { + return Ok(Transition::Next(RequestState::ReadLength { + read: new_read, + buf: mem::take(buf), + })); } let expected = u32::from_be_bytes(filled[..4].try_into()?) as usize; debug!("{}: Expected response length: {}", self.registry, expected); buf.resize(expected, 0); Ok(Transition::Next(RequestState::Reading { - read, + read: new_read, buf: mem::take(buf), expected, })) } RequestState::Reading { - mut read, + read, buf, expected, } => { - let mut read_buf = ReadBuf::new(&mut buf[read..]); + let mut read_buf = ReadBuf::new(&mut buf[*read..]); match Pin::new(&mut self.stream).poll_read(cx, &mut read_buf) { Poll::Ready(Ok(())) => {} Poll::Ready(Err(err)) => return Err(err.into()), @@ -197,20 +203,24 @@ impl EppConnection { .into()); } - read += filled.len(); + let new_read = *read + filled.len(); debug!( "{}: Read {} bytes, {} out of {} done", self.registry, filled.len(), - read, + new_read, expected ); // - Ok(if read < *expected { + Ok(if new_read < *expected { // If we haven't received the entire response yet, stick to the `Reading` state. - Transition::Next(state) + Transition::Next(RequestState::Reading { + read: new_read, + buf: mem::take(buf), + expected: *expected, + }) } else if let Some(next) = self.next.take() { // Otherwise, if we were just pushing through this request because it was already // in flight when we started a new one, ignore this response and move to the diff --git a/tests/chunked_read.rs b/tests/chunked_read.rs new file mode 100644 index 0000000..505a7f7 --- /dev/null +++ b/tests/chunked_read.rs @@ -0,0 +1,67 @@ +use std::time::Duration; + +use async_trait::async_trait; +use tokio_test::io::Builder; + +use instant_epp::client::{Connector, EppClient}; +use instant_epp::Error; + +#[tokio::test] +async fn greeting_single_chunk() { + assert!(connect_with_chunks(1).await.is_ok()); +} + +#[tokio::test] +async fn greeting_two_chunks() { + assert!(connect_with_chunks(2).await.is_ok()); +} + +async fn connect_with_chunks(num_chunks: usize) -> Result, Error> { + struct FakeConnector { + num_chunks: usize, + } + + #[async_trait] + impl Connector for FakeConnector { + type Connection = tokio_test::io::Mock; + + async fn connect(&self, _: Duration) -> Result { + let mut builder = Builder::new(); + + builder.read(&((GREETING.len() as u32) + 4).to_be_bytes()); + + let chunk_size = GREETING.len() / self.num_chunks; + for i in 0..self.num_chunks { + let start = i * chunk_size; + let end = if i == self.num_chunks - 1 { + GREETING.len() + } else { + start + chunk_size + }; + builder.read(&GREETING[start..end]); + } + + Ok(builder.build()) + } + } + + EppClient::new( + FakeConnector { num_chunks }, + "test".into(), + Duration::from_secs(5), + ) + .await +} + +const GREETING: &[u8] = br#" + + + Test EPP Server + 2024-01-01T00:00:00Z + + 1.0 + en + urn:ietf:params:xml:ns:domain-1.0 + + +"#;