diff --git a/src/io.rs b/src/io.rs index 8aeaa3c..188e2d4 100644 --- a/src/io.rs +++ b/src/io.rs @@ -68,7 +68,10 @@ pub trait WriteExt: Write { loop { match FormatBuffer::new(ignore_count).write(args) { Ok(data) => return self.write_all(data).await, - Err(FormatBufferWriteError::FormatError) => return Ok(()), + Err(FormatBufferWriteError::FormatError) => { + log_warn!("Skipping writing due to Format Error"); + return Ok(()); + } Err(FormatBufferWriteError::OutOfSpace(data)) => { self.write_all(data).await?; ignore_count += data.len(); diff --git a/src/response.rs b/src/response.rs index e6fe30b..e5fef40 100644 --- a/src/response.rs +++ b/src/response.rs @@ -40,11 +40,11 @@ pub use sse::EventStream; pub use status::StatusCode; pub use ws::WebSocketUpgrade; -struct MeasureFormatSize(pub usize); +struct MeasureFormatSize<'a>(&'a mut usize); -impl fmt::Write for MeasureFormatSize { +impl<'a> fmt::Write for MeasureFormatSize<'a> { fn write_str(&mut self, s: &str) -> fmt::Result { - self.0 += s.len(); + *self.0 += s.len(); Ok(()) } @@ -360,8 +360,8 @@ impl<'a> Content for fmt::Arguments<'a> { fn content_length(&self) -> usize { use fmt::Write; - let mut size = MeasureFormatSize(0); - write!(&mut size, "{self}").map_or(0, |()| size.0) + let mut size = 0; + write!(MeasureFormatSize(&mut size), "{self}").map_or(0, |()| size) } async fn write_content( diff --git a/src/response/chunked.rs b/src/response/chunked.rs index 9a06aba..9cb948a 100644 --- a/src/response/chunked.rs +++ b/src/response/chunked.rs @@ -9,7 +9,7 @@ pub struct ChunkWriter { } impl ChunkWriter { - /// Write a chunk to the client + /// Write a chunk to the client. pub async fn write_chunk(&mut self, chunk: &[u8]) -> Result<(), W::Error> { use crate::io::WriteExt; @@ -25,12 +25,33 @@ impl ChunkWriter { self.writer.flush().await } - /// Finish writing chunks + /// Finish writing chunks. pub async fn finalize(mut self) -> Result { self.writer.write_all(b"0\r\n\r\n").await?; + self.writer.flush().await?; Ok(ChunksWritten(())) } + + /// Write formatted text as a single chunk. This is typically called using the `write!` macro. + pub async fn write_fmt(&mut self, args: core::fmt::Arguments<'_>) -> Result<(), W::Error> { + use crate::io::WriteExt; + use core::fmt::Write; + + let mut chunk_size = 0; + + if super::MeasureFormatSize(&mut chunk_size) + .write_fmt(args) + .is_err() + { + log_warn!("Skipping writing chunk due to Format Error"); + + return Ok(()); + } + + write!(&mut self.writer, "{chunk_size:x}\r\n{args}\r\n",).await?; + self.writer.flush().await + } } /// A series of chunks forming the response body @@ -64,8 +85,8 @@ impl ChunkedResponse { impl super::Body for Body { async fn write_response_body< - R: embedded_io_async::Read, - W: embedded_io_async::Write, + R: crate::io::Read, + W: crate::io::Write, >( self, _connection: super::Connection<'_, R>, diff --git a/src/response/json.rs b/src/response/json.rs index a57f7d2..ff687a7 100644 --- a/src/response/json.rs +++ b/src/response/json.rs @@ -481,10 +481,12 @@ impl super::Content for JsonBody { match &self.0 { JsonStream::Short { buffer } => buffer.data.len(), JsonStream::Long { buffer: _, value } => { - let mut content_length = super::MeasureFormatSize(0); + let mut content_length = 0; value - .serialize(Serializer(&mut content_length)) - .map_or(0, |()| content_length.0) + .serialize(Serializer(&mut super::MeasureFormatSize( + &mut content_length, + ))) + .map_or(0, |()| content_length) } } }