Emit warnings when data limits are reached.

Aid in debugging incorrectly configured data limits.
This commit is contained in:
Sergio Benitez 2023-04-03 12:06:10 -07:00
parent 9c4ac79748
commit 887558be60
1 changed files with 14 additions and 1 deletions

View File

@ -65,7 +65,7 @@ enum StreamKind<'r> {
impl<'r> DataStream<'r> { impl<'r> DataStream<'r> {
pub(crate) fn new(buf: Vec<u8>, stream: StreamReader<'r>, limit: u64) -> Self { pub(crate) fn new(buf: Vec<u8>, stream: StreamReader<'r>, limit: u64) -> Self {
let chain = Chain::new(Cursor::new(buf), stream).take(limit); let chain = Chain::new(Cursor::new(buf), stream).take(limit).into();
Self { chain } Self { chain }
} }
@ -73,6 +73,7 @@ impl<'r> DataStream<'r> {
async fn limit_exceeded(&mut self) -> io::Result<bool> { async fn limit_exceeded(&mut self) -> io::Result<bool> {
#[cold] #[cold]
async fn _limit_exceeded(stream: &mut DataStream<'_>) -> io::Result<bool> { async fn _limit_exceeded(stream: &mut DataStream<'_>) -> io::Result<bool> {
// Read one more byte after reaching limit to see if we cut early.
stream.chain.set_limit(1); stream.chain.set_limit(1);
let mut buf = [0u8; 1]; let mut buf = [0u8; 1];
Ok(stream.read(&mut buf).await? != 0) Ok(stream.read(&mut buf).await? != 0)
@ -252,6 +253,18 @@ impl AsyncRead for DataStream<'_> {
cx: &mut Context<'_>, cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>, buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> { ) -> Poll<io::Result<()>> {
if self.chain.limit() == 0 {
let stream: &StreamReader<'_> = &self.chain.get_ref().get_ref().1;
let kind = match stream.inner {
StreamKind::Empty => "an empty stream (vacuous)",
StreamKind::Body(_) => "the request body",
StreamKind::Multipart(_) => "a multipart form field",
};
let msg = yansi::Paint::default(kind).bold();
warn_!("Data limit reached while reading {}.", msg);
}
Pin::new(&mut self.chain).poll_read(cx, buf) Pin::new(&mut self.chain).poll_read(cx, buf)
} }
} }