fix(server): Finish decoding compressed requests correctly [INGEST-619]

CHANGELOG.md

@@ -9,7 +9,7 @@
 
 **Bug Fixes**:
 
-- Avoid unbounded decompression of encoded requests. A particular request crafted to inflate to large amounts of memory, such as a zip bomb, could put Relay out of memory. ([#1117](https://github.com/getsentry/relay/pull/1117), [#1122](https://github.com/getsentry/relay/pull/1122))
+- Avoid unbounded decompression of encoded requests. A particular request crafted to inflate to large amounts of memory, such as a zip bomb, could put Relay out of memory. ([#1117](https://github.com/getsentry/relay/pull/1117), [#1122](https://github.com/getsentry/relay/pull/1122), [#1123](https://github.com/getsentry/relay/pull/1123))
 - Avoid unbounded decompression of UE4 crash reports. Some crash reports could inflate to large amounts of memory before being checked for size, which could put Relay out of memory. ([#1121](https://github.com/getsentry/relay/pull/1121))
 
 **Internal**:

relay-server/src/body/peek_line.rs

@@ -43,8 +43,8 @@ impl PeekLine {
         }
     }
 
-    fn finish(&mut self, overflow: bool) -> Async<Option<Bytes>> {
-        let buffer = self.decoder.take();
+    fn finish(&mut self, overflow: bool) -> std::io::Result<Option<Bytes>> {
+        let buffer = self.decoder.finish()?;
 
         let line = match buffer.iter().position(|b| *b == b'\n') {
             Some(pos) => Some(buffer.slice_to(pos)),
@@ -53,7 +53,7 @@ impl PeekLine {
         };
 
         self.revert_chunks();
-        Async::Ready(line.filter(|line| !line.is_empty()))
+        Ok(line.filter(|line| !line.is_empty()))
     }
 }
 
@@ -63,16 +63,15 @@ impl Future for PeekLine {
 
     fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
         loop {
-            let chunk = match self.payload.poll() {
-                Ok(Async::Ready(Some(chunk))) => chunk,
-                Ok(Async::Ready(None)) => return Ok(self.finish(false)),
-                Ok(Async::NotReady) => return Ok(Async::NotReady),
-                Err(error) => return Err(error),
+            let chunk = match self.payload.poll()? {
+                Async::Ready(Some(chunk)) => chunk,
+                Async::Ready(None) => return Ok(Async::Ready(self.finish(false)?)),
+                Async::NotReady => return Ok(Async::NotReady),
             };
 
             self.chunks.push(chunk.clone());
             if self.decoder.decode(chunk)? {
-                return Ok(self.finish(true));
+                return Ok(Async::Ready(self.finish(true)?));
             }
         }
     }

relay-server/src/body/request_body.rs

@@ -41,17 +41,16 @@ impl Future for RequestBody {
 
         if let Some((ref mut payload, ref mut decoder)) = self.stream {
             loop {
-                return match payload.poll() {
-                    Ok(Async::Ready(Some(encoded))) => {
+                return match payload.poll()? {
+                    Async::Ready(Some(encoded)) => {
                         if decoder.decode(encoded)? {
                             Err(PayloadError::Overflow)
                         } else {
                             continue;
                         }
                     }
-                    Ok(Async::Ready(None)) => Ok(Async::Ready(decoder.take())),
-                    Ok(Async::NotReady) => Ok(Async::NotReady),
-                    Err(error) => Err(error),
+                    Async::Ready(None) => Ok(Async::Ready(decoder.finish()?)),
+                    Async::NotReady => Ok(Async::NotReady),
                 };
             }
         }

relay-server/src/extractors/decoder.rs

@@ -76,6 +76,7 @@ fn write_overflowing<W: Write>(write: &mut W, slice: &[u8]) -> io::Result<bool>
     match write.write_all(slice).and_then(|()| write.flush()) {
         Ok(()) => Ok(false),
         Err(e) => match e.kind() {
+            io::ErrorKind::WouldBlock => Ok(false), // effectively continue
             io::ErrorKind::WriteZero => Ok(true),
             _ => Err(e),
         },
@@ -146,11 +147,30 @@ impl Decoder {
         }
     }
 
+    /// Finish decoding the output stream and validate checksums, returning the final bytes.
+    ///
+    /// This may only be called a single time at the end of decoding. Attempts to write data to this
+    /// decoder may result in a panic after this function is called.
+    pub fn finish(&mut self) -> io::Result<Bytes> {
+        Ok(match &mut self.inner {
+            DecoderInner::Identity(inner) => inner.take(),
+            DecoderInner::Br(inner) => inner.finish()?.take(),
+            DecoderInner::Gzip(inner) => {
+                inner.try_finish()?;
+                inner.get_mut().take()
+            }
+            DecoderInner::Deflate(inner) => {
+                inner.try_finish()?;
+                inner.get_mut().take()
+            }
+        })
+    }
+
     /// Returns decoded bytes from the Decoder's buffer.
     ///
     /// This can be called at any time during decoding. However, the limit stated during
     /// initialization remains in effect across the entire decoded payload.
-    pub fn take(&mut self) -> Bytes {
+    fn take(&mut self) -> Bytes {
         match &mut self.inner {
             DecoderInner::Identity(inner) => inner.take(),
             DecoderInner::Br(inner) => inner.get_mut().take(),
@@ -195,22 +215,26 @@ impl Stream for DecodingPayload {
     type Error = <SharedPayload as Stream>::Error;
 
     fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
-        let encoded = match self.payload.poll() {
-            Ok(Async::Ready(Some(encoded))) => encoded,
-            Ok(Async::Ready(None)) => return Ok(Async::Ready(None)),
-            Ok(Async::NotReady) => return Ok(Async::NotReady),
-            Err(error) => return Err(error),
-        };
+        loop {
+            let encoded = match self.payload.poll()? {
+                Async::Ready(Some(encoded)) => encoded,
+                Async::Ready(None) => {
+                    let chunk = self.decoder.finish()?;
+                    return Ok(Async::Ready(Some(chunk).filter(|c| !c.is_empty())));
+                }
+                Async::NotReady => return Ok(Async::NotReady),
+            };
+
+            if self.decoder.decode(encoded)? {
+                return Err(PayloadError::Overflow);
+            }
 
-        if self.decoder.decode(encoded)? {
-            Err(PayloadError::Overflow)
-        } else {
             let chunk = self.decoder.take();
-            Ok(Async::Ready(if chunk.is_empty() {
-                None
-            } else {
-                Some(chunk)
-            }))
+            if !chunk.is_empty() {
+                return Ok(Async::Ready(Some(chunk)));
+            }
+
+            // loop until the input stream is exhausted
         }
     }
 }