Skip to content

Commit

Permalink
Fix issue where writing up to exactly the size of the chunked buffer …
Browse files Browse the repository at this point in the history
…size causes data loss

This fixes an issue where a call to write() with a buffer length exactly matching the remaining size of the remainder of the current chunk causes the entire chunk to be discarded
  • Loading branch information
rmja committed Jul 2, 2024
1 parent 908d0e2 commit e0c31dc
Showing 1 changed file with 24 additions and 3 deletions.
27 changes: 24 additions & 3 deletions src/body_writer/buffering_chunked.rs
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,10 @@ where
self.pos = self.header_pos + self.allocated_header;
}

fn current_chunk_is_full(&self) -> bool {
self.pos + NEWLINE.len() == self.buf.len()
}

async fn emit_buffered(&mut self) -> Result<(), C::Error> {
self.conn.write_all(&self.buf[..self.header_pos]).await?;
self.header_pos = 0;
Expand Down Expand Up @@ -149,7 +153,7 @@ where
self.emit_buffered().await.map_err(|e| e.kind())?;
written = self.append_current_chunk(buf);
}
if written < buf.len() {
if self.current_chunk_is_full() {
self.finish_current_chunk();
self.emit_buffered().await.map_err(|e| e.kind())?;
}
Expand Down Expand Up @@ -330,19 +334,36 @@ mod tests {
}

#[tokio::test]
async fn current_chunk_is_emitted_before_empty_chunk_is_emitted() {
async fn written_bytes_fit_exactly() {
// Given
let mut conn = Vec::new();
let mut buf = [0; 14];
buf[..5].copy_from_slice(b"HELLO");

// When
let mut writer = BufferingChunkedBodyWriter::new_with_data(&mut conn, &mut buf, 5);
writer.write_all(b"BODY").await.unwrap(); // Can fit exactly
writer.write_all(b"BODY").await.unwrap(); // Can fit
writer.terminate().await.unwrap(); // Can fit

// Then
assert_eq!(b"HELLO4\r\nBODY\r\n4\r\nBODY\r\n0\r\n\r\n", conn.as_slice());
}

#[tokio::test]
async fn current_chunk_is_emitted_on_termination_before_empty_chunk_is_emitted() {
// Given
let mut conn = Vec::new();
let mut buf = [0; 14];
buf[..5].copy_from_slice(b"HELLO");

// When
let mut writer = BufferingChunkedBodyWriter::new_with_data(&mut conn, &mut buf, 5);
writer.write_all(b"BOD").await.unwrap(); // Can fit
writer.terminate().await.unwrap(); // Cannot fit

// Then
assert_eq!(b"HELLO4\r\nBODY\r\n0\r\n\r\n", conn.as_slice());
assert_eq!(b"HELLO3\r\nBOD\r\n0\r\n\r\n", conn.as_slice());
}

#[tokio::test]
Expand Down

0 comments on commit e0c31dc

Please sign in to comment.