more lib fix

lib/vey-http/src/body/chunked_decoder.rs

@@ -51,6 +51,96 @@ impl ChunkedDataDecodeReaderInternal {
         self.chunk_header.is_empty() && !self.poll_chunk_end_r && !self.poll_chunk_end_n
     }
 
+    /// Consume a chunk-ending CRLF that was not available when the payload was read.
+    ///
+    /// Stops at the next chunk-size line. Does not read that line.
+    fn poll_finish_chunk_end<R>(
+        &mut self,
+        cx: &mut Context<'_>,
+        mut reader: Pin<&mut R>,
+    ) -> Poll<io::Result<()>>
+    where
+        R: AsyncBufRead + Unpin,
+    {
+        loop {
+            if self.poll_chunk_end {
+                if self.this_chunk_size != 0 {
+                    self.poll_chunk_end = false;
+                    self.poll_chunk_end_n = false;
+                    self.poll_chunk_end_r = false;
+                }
+                return Poll::Ready(Ok(()));
+            } else if self.poll_chunk_end_n {
+                let r_buf = ready!(reader.as_mut().poll_fill_buf(cx))?;
+                if r_buf.is_empty() {
+                    return Poll::Ready(Err(io::Error::new(
+                        io::ErrorKind::UnexpectedEof,
+                        "reader closed while reading chunk line end whitespace",
+                    )));
+                } else if r_buf[0] == b'\n' {
+                    reader.as_mut().consume(1);
+                    self.poll_chunk_end = true;
+                } else {
+                    return Poll::Ready(Err(io::Error::new(
+                        io::ErrorKind::InvalidData,
+                        "no end whitespace found",
+                    )));
+                }
+            } else if self.poll_chunk_end_r {
+                let r_buf = ready!(reader.as_mut().poll_fill_buf(cx))?;
+                match r_buf.len() {
+                    0 => {
+                        return Poll::Ready(Err(io::Error::new(
+                            io::ErrorKind::UnexpectedEof,
+                            "reader closed while reading chunk line end whitespace",
+                        )));
+                    }
+                    1 => match r_buf[0] {
+                        b'\r' => {
+                            reader.as_mut().consume(1);
+                            self.poll_chunk_end_n = true;
+                        }
+                        b'\n' => {
+                            reader.as_mut().consume(1);
+                            self.poll_chunk_end = true;
+                        }
+                        _ => {
+                            return Poll::Ready(Err(io::Error::new(
+                                io::ErrorKind::InvalidData,
+                                "no end whitespace found",
+                            )));
+                        }
+                    },
+                    _ => match r_buf[0] {
+                        b'\r' => {
+                            if r_buf[1] != b'\n' {
+                                return Poll::Ready(Err(io::Error::new(
+                                    io::ErrorKind::InvalidData,
+                                    "invalid end whitespace pair",
+                                )));
+                            } else {
+                                reader.as_mut().consume(2);
+                                self.poll_chunk_end = true;
+                            }
+                        }
+                        b'\n' => {
+                            reader.as_mut().consume(1);
+                            self.poll_chunk_end = true;
+                        }
+                        _ => {
+                            return Poll::Ready(Err(io::Error::new(
+                                io::ErrorKind::InvalidData,
+                                "no end whitespace found",
+                            )));
+                        }
+                    },
+                }
+            } else {
+                return Poll::Ready(Ok(()));
+            }
+        }
+    }
+
     fn poll_decode<R>(
         &mut self,
         cx: &mut Context<'_>,
@@ -248,6 +338,20 @@ impl<'a, R> ChunkedDataDecodeReader<'a, R> {
         self.internal.left_chunk_size()
     }
 
+    /// Finish a chunk-ending CRLF left pending when the payload read returned early.
+    ///
+    /// After this returns, `left_chunk_size` is `Some` unless the body is already finished.
+    pub async fn finish_pending_chunk_end(&mut self) -> io::Result<()>
+    where
+        R: AsyncBufRead + Unpin,
+    {
+        std::future::poll_fn(|cx| {
+            let ChunkedDataDecodeReader { reader, internal } = &mut *self;
+            internal.poll_finish_chunk_end(cx, Pin::new(&mut **reader))
+        })
+        .await
+    }
+
     /**
      * Check whether it's safe to break from a Poll::Pending state
      *
@@ -382,6 +486,60 @@ mod test {
         );
     }
 
+    #[tokio::test]
+    async fn finish_chunk_end_split_from_payload() {
+        let stream = tokio_test::io::Builder::new()
+            .read(b"5\r\nhello")
+            .wait(std::time::Duration::from_millis(1))
+            .read(b"\r\n4\r\nbody\r\n0\r\n\r\n")
+            .build();
+        let mut buf_stream = BufReader::new(stream);
+        let mut decoder = ChunkedDataDecodeReader::new(&mut buf_stream, 1024);
+
+        let mut buf = [0u8; 5];
+        let len = decoder.read(&mut buf).await.unwrap();
+        assert_eq!(&buf[..len], b"hello");
+        assert!(!decoder.finished());
+        assert!(decoder.left_chunk_size().is_none());
+
+        decoder.finish_pending_chunk_end().await.unwrap();
+        assert_eq!(decoder.left_chunk_size(), Some(0));
+        assert!(decoder.pending_cancel_safe());
+
+        let mut decoder = ChunkedDataDecodeReader::resume(decoder.into_reader(), 1024, 0);
+        let mut buf = [0u8; 16];
+        let len = decoder.read(&mut buf).await.unwrap();
+        assert_eq!(&buf[..len], b"body");
+        assert!(decoder.finished());
+    }
+
+    #[tokio::test]
+    async fn finish_chunk_end_split_crlf() {
+        let stream = tokio_test::io::Builder::new()
+            .read(b"5\r\nhello")
+            .wait(std::time::Duration::from_millis(1))
+            .read(b"\r")
+            .wait(std::time::Duration::from_millis(1))
+            .read(b"\n0\r\n\r\n")
+            .build();
+        let mut buf_stream = BufReader::new(stream);
+        let mut decoder = ChunkedDataDecodeReader::new(&mut buf_stream, 1024);
+
+        let mut buf = [0u8; 5];
+        let len = decoder.read(&mut buf).await.unwrap();
+        assert_eq!(&buf[..len], b"hello");
+        assert!(decoder.left_chunk_size().is_none());
+
+        decoder.finish_pending_chunk_end().await.unwrap();
+        assert_eq!(decoder.left_chunk_size(), Some(0));
+
+        let mut decoder = ChunkedDataDecodeReader::resume(decoder.into_reader(), 1024, 0);
+        let mut buf = [0u8; 8];
+        let len = decoder.read(&mut buf).await.unwrap();
+        assert_eq!(len, 0);
+        assert!(decoder.finished());
+    }
+
     #[tokio::test]
     async fn chunk_line_at_max_size_is_accepted() {
         let content = b"0\r\n\r\n";

lib/vey-http/src/client/adaptation.rs

@@ -23,6 +23,8 @@ pub struct HttpAdaptedResponse {
     pub headers: H1HeaderMap,
     pub content_length: Option<u64>,
     connection: ConnectionValue,
+    bytes_read: usize,
+    header_limit: usize,
 }
 
 impl HttpAdaptedResponse {
@@ -34,6 +36,22 @@ impl HttpAdaptedResponse {
             headers: H1HeaderMap::default(),
             content_length: None,
             connection: ConnectionValue::default(),
+            bytes_read: 0,
+            header_limit: 0,
+        }
+    }
+
+    /// The encapsulated header ended on its blank line at `header_size`.
+    pub fn header_section_finished(&self) -> bool {
+        self.bytes_read >= self.header_limit
+    }
+
+    /// A following body can only be read when the header fills the encapsulated section.
+    pub fn require_header_section_finished(&self) -> Result<(), HttpResponseParseError> {
+        if self.header_section_finished() {
+            Ok(())
+        } else {
+            Err(HttpResponseParseError::HeaderEndedEarly(self.header_limit))
         }
     }
 
@@ -118,6 +136,8 @@ impl HttpAdaptedResponse {
             rsp.parse_header_line(&line_buf)?;
         }
 
+        rsp.bytes_read = read_size;
+        rsp.header_limit = header_size;
         Ok(rsp)
     }
 
@@ -394,4 +414,22 @@ mod tests {
         assert_eq!(headers.len(), 1);
         assert_eq!(headers.get("x-a").unwrap(), "1");
     }
+
+    #[tokio::test]
+    async fn blank_line_before_section_end_is_unfinished() {
+        let headers = b"HTTP/1.1 200 OK\r\n\r\n";
+        let mut data = headers.to_vec();
+        data.extend_from_slice(b"LEFTOVER");
+        let mut reader = BufReader::new(std::io::Cursor::new(data));
+        let rsp = HttpAdaptedResponse::parse(&mut reader, headers.len() + 8)
+            .await
+            .unwrap();
+        assert!(!rsp.header_section_finished());
+        let err = rsp.require_header_section_finished().unwrap_err();
+        assert!(matches!(err, HttpResponseParseError::HeaderEndedEarly(_)));
+        let buffered = tokio::io::AsyncBufReadExt::fill_buf(&mut reader)
+            .await
+            .unwrap();
+        assert_eq!(buffered, b"LEFTOVER");
+    }
 }

lib/vey-http/src/client/error.rs

@@ -19,6 +19,8 @@ pub enum HttpResponseParseError {
     RemoteClosed,
     #[error("too large header, should be less than {0}")]
     TooLargeHeader(usize),
+    #[error("http header ended before declared section length {0}")]
+    HeaderEndedEarly(usize),
     #[error("invalid version {0:?}")]
     InvalidVersion(Version),
     #[error("invalid status line: {0}")]

lib/vey-http/src/server/adaptation.rs

@@ -22,6 +22,8 @@ pub struct HttpAdaptedRequest {
     pub version: Version,
     pub headers: H1HeaderMap,
     pub content_length: Option<u64>,
+    bytes_read: usize,
+    header_limit: usize,
 }
 
 impl HttpAdaptedRequest {
@@ -32,6 +34,22 @@ impl HttpAdaptedRequest {
             version,
             headers: H1HeaderMap::default(),
             content_length: None,
+            bytes_read: 0,
+            header_limit: 0,
+        }
+    }
+
+    /// The encapsulated header ended on its blank line at `header_size`.
+    pub fn header_section_finished(&self) -> bool {
+        self.bytes_read >= self.header_limit
+    }
+
+    /// A following body can only be read when the header fills the encapsulated section.
+    pub fn require_header_section_finished(&self) -> Result<(), HttpRequestParseError> {
+        if self.header_section_finished() {
+            Ok(())
+        } else {
+            Err(HttpRequestParseError::HeaderEndedEarly(self.header_limit))
         }
     }
 
@@ -93,6 +111,8 @@ impl HttpAdaptedRequest {
             req.parse_header_line(&line_buf, ignore_via)?;
         }
 
+        req.bytes_read = read_size;
+        req.header_limit = header_size;
         Ok(req)
     }
 
@@ -433,4 +453,22 @@ mod tests {
             "value with spaces"
         );
     }
+
+    #[tokio::test]
+    async fn blank_line_before_section_end_is_unfinished() {
+        let headers = b"GET / HTTP/1.1\r\n\r\n";
+        let mut data = headers.to_vec();
+        data.extend_from_slice(b"LEFTOVER");
+        let mut reader = BufReader::new(std::io::Cursor::new(data));
+        let req = HttpAdaptedRequest::parse(&mut reader, headers.len() + 8, false)
+            .await
+            .unwrap();
+        assert!(!req.header_section_finished());
+        let err = req.require_header_section_finished().unwrap_err();
+        assert!(matches!(err, HttpRequestParseError::HeaderEndedEarly(_)));
+        let buffered = tokio::io::AsyncBufReadExt::fill_buf(&mut reader)
+            .await
+            .unwrap();
+        assert_eq!(buffered, b"LEFTOVER");
+    }
 }

lib/vey-http/src/server/error.rs

@@ -21,6 +21,8 @@ pub enum HttpRequestParseError {
     ClientClosed,
     #[error("too large header, should be less than {0}")]
     TooLargeHeader(usize),
+    #[error("http header ended before declared section length {0}")]
+    HeaderEndedEarly(usize),
     #[error("invalid method line: {0}")]
     InvalidMethodLine(HttpLineParseError),
     #[error("unsupported method: {0}")]