Change CountingStream to BoundedStream

Author: elramen

relay-server/src/endpoints/upload.rs

@@ -31,7 +31,7 @@ use crate::services::projects::cache::Project;
 use crate::services::upload::Error as ServiceError;
 use crate::services::upstream::UpstreamRequestError;
 use crate::utils::upload::SignedLocation;
-use crate::utils::{CountingStream, find_error_source, tus, upload};
+use crate::utils::{BoundedStream, find_error_source, tus, upload};
 
 #[derive(Debug, thiserror::Error)]
 enum Error {
@@ -137,7 +137,11 @@ async fn handle(
         .into_data_stream()
         .map(|result| result.map_err(io::Error::other))
         .boxed();
-    let stream = CountingStream::new(stream, upload_length);
+    let (lower_bound, upper_bound) = match upload_length {
+        None => (1, config.max_upload_size()),
+        Some(u) => (u, u),
+    };
+    let stream = BoundedStream::new(stream, lower_bound, upper_bound);
     let byte_counter = stream.byte_counter();
 
     let location = upload::Sink::new(&state)

relay-server/src/utils/stream.rs

@@ -16,13 +16,14 @@ use sync_wrapper::SyncWrapper;
 /// This type is `Sync` via [`SyncWrapper`], allowing it to be sent across thread boundaries
 /// as required by the upload service.
 #[derive(Debug)]
-pub struct CountingStream<S> {
+pub struct BoundedStream<S> {
+    pub lower_bound: usize,
+    pub upper_bound: usize,
     inner: Option<SyncWrapper<S>>,
-    expected_length: Option<usize>,
     byte_counter: ByteCounter,
 }
 
-/// A shared counter that can be read after the [`CountingStream`] has been moved.
+/// A shared counter that can be read after the [`BoundedStream`] has been moved.
 #[derive(Clone, Debug)]
 pub struct ByteCounter(Arc<AtomicUsize>);
 
@@ -40,28 +41,24 @@ impl ByteCounter {
     }
 }
 
-impl<S> CountingStream<S> {
-    /// Creates a new [`CountingStream`] wrapping the given stream with the expected total length.
-    pub fn new(stream: S, expected_length: Option<usize>) -> Self {
+impl<S> BoundedStream<S> {
+    /// Creates a new [`BoundedStream`] wrapping the given stream with the expected total length.
+    pub fn new(stream: S, lower_bound: usize, upper_bound: usize) -> Self {
         Self {
             inner: Some(SyncWrapper::new(stream)),
-            expected_length,
+            lower_bound,
+            upper_bound,
             byte_counter: ByteCounter::new(),
         }
     }
 
-    /// Returns the expected total length of the stream.
-    pub fn expected_length(&self) -> Option<usize> {
-        self.expected_length
-    }
-
     /// Returns a shared handle to read the byte count after the stream is consumed.
     pub fn byte_counter(&self) -> ByteCounter {
         self.byte_counter.clone()
     }
 }
 
-impl<S, E> Stream for CountingStream<S>
+impl<S, E> Stream for BoundedStream<S>
 where
     S: Stream<Item = Result<Bytes, E>> + Send + Unpin,
     E: Into<io::Error>,
@@ -78,15 +75,13 @@ where
         match inner.poll_next(cx) {
             Poll::Ready(Some(Ok(bytes))) => {
                 let bytes_received = this.byte_counter.add(bytes.len());
-                if let Some(expected_length) = this.expected_length
-                    && bytes_received > expected_length
-                {
+                if bytes_received > this.upper_bound {
                     this.inner = None;
                     Poll::Ready(Some(Err(io::Error::new(
                         io::ErrorKind::FileTooLarge,
                         format!(
-                            "stream exceeded expected length: received {} > {}",
-                            bytes_received, expected_length
+                            "stream exceeded upper bound: received {} > {}",
+                            bytes_received, this.upper_bound
                         ),
                     ))))
                 } else {
@@ -96,15 +91,13 @@ where
             Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e.into()))),
             Poll::Ready(None) => {
                 let bytes_received = this.byte_counter.get();
-                if let Some(expected_length) = this.expected_length
-                    && bytes_received < expected_length
-                {
+                if bytes_received < this.lower_bound {
                     this.inner = None;
                     Poll::Ready(Some(Err(io::Error::new(
                         io::ErrorKind::UnexpectedEof,
                         format!(
-                            "stream shorter than expected length: received {} < {}",
-                            bytes_received, expected_length
+                            "stream shorter than lower bound: received {} < {}",
+                            bytes_received, this.lower_bound
                         ),
                     ))))
                 } else {

relay-server/src/utils/upload.rs

@@ -24,7 +24,7 @@ use crate::services::upload::{Error as ServiceError, Upload};
 use crate::services::upstream::{
     SendRequest, UpstreamRelay, UpstreamRequest, UpstreamRequestError,
 };
-use crate::utils::{CountingStream, tus};
+use crate::utils::{BoundedStream, tus};
 
 /// An error that occurs during upload.
 #[derive(Debug, thiserror::Error)]
@@ -57,7 +57,7 @@ pub struct Stream {
     /// The organization and project the stream belongs to.
     pub scoping: Scoping,
     /// The body to be uploaded to objectstore, with length validation.
-    pub stream: CountingStream<BoxStream<'static, std::io::Result<Bytes>>>,
+    pub stream: BoundedStream<BoxStream<'static, std::io::Result<Bytes>>>,
 }
 
 /// A dispatcher for uploading large files.
@@ -204,7 +204,7 @@ impl SignedLocation {
 /// An upstream request made to the `/upload` endpoint.
 struct UploadRequest {
     scoping: Scoping,
-    body: Option<CountingStream<BoxStream<'static, std::io::Result<Bytes>>>>,
+    body: Option<BoundedStream<BoxStream<'static, std::io::Result<Bytes>>>>,
     sender: oneshot::Sender<Result<Response, UpstreamRequestError>>,
 }
 
@@ -280,7 +280,12 @@ impl UpstreamRequest for UploadRequest {
 
         let project_key = self.scoping.project_key;
         builder.header("X-Sentry-Auth", format!("Sentry sentry_key={project_key}"));
-        for (key, value) in tus::request_headers(body.expected_length()) {
+        let upload_length = if body.lower_bound == body.upper_bound {
+            Some(body.lower_bound)
+        } else {
+            None
+        };
+        for (key, value) in tus::request_headers(upload_length) {
             let Some(key) = key else { continue };
             builder.header(key, value);
         }