Refactor STOW

docs/topics/configuration.md

@@ -48,7 +48,8 @@ server:
   http:
     interface: 0.0.0.0
     port: 8080
-    max-upload-size: 50000000
+    max-upload-size: 5000000000
+    max-part-size: 50000000
     request-timeout: 60000
     graceful-shutdown: true
   dimse:
@@ -121,7 +122,8 @@ server:
   http:
     interface: 0.0.0.0
     port: 8080
-    max-upload-size: 50000000
+    max-upload-size: 5000000000
+    max-part-size: 50000000
     request-timeout: 60000
     graceful-shutdown: true
     base-path: /
@@ -135,8 +137,14 @@ server:
         The port for the HTTP server. Uses <code>8080</code> by default.
     </def>
     <def title="server.http.max-upload-size" id="server.http.max-upload-size">
-        The maximum allowed request body size in bytes. 
-        This setting can be used to limit the upload size of STOW-RS requests.
+        The maximum allowed request body size in bytes. Uses <code>5000000000</code> (5 GB) by default.
+        This setting can be used to limit the total upload size of STOW-RS requests.
+    </def>
+    <def title="server.http.max-part-size" id="server.http.max-part-size">
+        The maximum allowed size in bytes of a single part (i.e. a single DICOM instance) of a STOW-RS request.
+        Uses <code>50000000</code> (50 MB) by default.
+        Instances are received while the previous instance is being stored, so a STOW-RS request holds
+        at most a few instances in memory at once. This setting therefore bounds the memory used per request.
     </def>
     <def title="server.http.request-timeout" id="server.http.request.timeout">
         The maximum allowed (total) time for a HTTP request in milliseconds before a timeout occurs. This applies to all endpoints.

src/api/stow/routes.rs

@@ -2,6 +2,7 @@ use crate::api::stow::{StoreError, StoreRequest};
 use crate::backend::ServiceProvider;
 use crate::utils::multipart::DicomMultipart;
 use crate::AppState;
+use async_stream::try_stream;
 use axum::body::Body;
 use axum::http::header::CONTENT_TYPE;
 use axum::http::StatusCode;
@@ -11,7 +12,9 @@ use axum::Router;
 use bytes::Buf;
 use dicom::object::{FileDicomObject, InMemDicomObject};
 use dicom_json::DicomJson;
-use tracing::instrument;
+use futures::channel::mpsc;
+use futures::{FutureExt, StreamExt};
+use tracing::{instrument, Instrument};
 
 /// HTTP Router for the Store Transaction
 /// <https://dicom.nema.org/medical/dicom/current/output/html/part18.html#sect_10.5>
@@ -34,15 +37,31 @@ async fn studies(
 			.into_response());
 	};
 
-	let mut instances = Vec::new();
+	// Parts are read lazily instead of buffering the whole multipart body in memory.
+	let parts = try_stream! {
+		while let Some(field) = multipart.next_field().await? {
+			let data = field.bytes().await?;
+			yield FileDicomObject::from_reader(data.reader())?;
+		}
+	};
 
-	while let Some(field) = multipart.next_field().await? {
-		let data = field.bytes().await?;
-		let file = FileDicomObject::from_reader(data.reader())?;
-		instances.push(file);
-	}
+	// Parts are read by a separate task, so that the next instance is received while the backend
+	// is still storing the previous one. The bounded channel limits how far the reader can get
+	// ahead of the backend, which keeps the memory usage per request bounded.
+	// Note that the capacity of a futures channel is `buffer + 1` per sender.
+	let (sender, receiver) = mpsc::channel(0);
+	// The reader is cancelled as soon as the handle is dropped, i.e. when the backend is done
+	// or the request is cancelled.
+	let (reader, _reader_handle) = parts
+		.map(Ok)
+		.forward(sender)
+		.in_current_span()
+		.remote_handle();
+	tokio::spawn(reader);
 
-	let request = StoreRequest { instances };
+	let request = StoreRequest {
+		instances: receiver.boxed(),
+	};
 	let response = stow.store(request).await?;
 	let json = DicomJson::from(InMemDicomObject::from(response));
 

src/api/stow/service.rs

@@ -9,10 +9,13 @@ use dicom::dicom_value;
 use dicom::dictionary_std::tags;
 use dicom::object::mem::InMemElement;
 use dicom::object::{FileDicomObject, InMemDicomObject};
+use futures::stream::BoxStream;
 use thiserror::Error;
 
 pub struct StoreRequest {
-	pub instances: Vec<FileDicomObject<InMemDicomObject>>,
+	/// The instances to store. They are read lazily from the request body, so that
+	/// backends can store each instance without holding the whole upload in memory.
+	pub instances: BoxStream<'static, Result<FileDicomObject<InMemDicomObject>, StoreError>>,
 }
 
 /// <https://dicom.nema.org/medical/dicom/current/output/html/part03.html#table_10-11>
@@ -91,8 +94,10 @@ pub trait StowService: Sync + Send {
 
 #[derive(Debug, Error)]
 pub enum StoreError {
-	#[error("The file exceeds the configured upload size limit")]
+	#[error("The request exceeds the configured upload size limit")]
 	UploadLimitExceeded,
+	#[error("A part of the request exceeds the configured part size limit of {limit} bytes")]
+	PartLimitExceeded { limit: u64 },
 	#[error(transparent)]
 	Stream(multer::Error),
 	#[error(transparent)]
@@ -106,6 +111,12 @@ impl IntoResponse for StoreError {
 				.status(StatusCode::PAYLOAD_TOO_LARGE)
 				.body(Body::from("Upload limit exceeded"))
 				.unwrap(),
+			Self::PartLimitExceeded { limit } => Response::builder()
+				.status(StatusCode::PAYLOAD_TOO_LARGE)
+				.body(Body::from(format!(
+					"Part limit exceeded: each instance must not exceed {limit} bytes"
+				)))
+				.unwrap(),
 			Self::Stream(err) => Response::builder()
 				.status(StatusCode::BAD_REQUEST)
 				.body(Body::from(format!(
@@ -122,6 +133,10 @@ impl IntoResponse for StoreError {
 
 impl From<multer::Error> for StoreError {
 	fn from(error: multer::Error) -> Self {
+		if let multer::Error::FieldSizeExceeded { limit, .. } = error {
+			return Self::PartLimitExceeded { limit };
+		}
+
 		if let multer::Error::StreamReadFailed(stream_error) = &error {
 			let is_limit_exceeded = stream_error
 				.downcast_ref::<axum::Error>()

src/backend/dimse/stow.rs

@@ -4,6 +4,7 @@ use crate::backend::dimse::cstore::storescu::StoreServiceClassUser;
 use crate::types::UI;
 use association::pool::AssociationPool;
 use async_trait::async_trait;
+use futures::TryStreamExt;
 use std::time::Duration;
 use tracing::info;
 
@@ -20,11 +21,11 @@ impl DimseStowService {
 
 #[async_trait]
 impl StowService for DimseStowService {
-	async fn store(&self, request: StoreRequest) -> Result<StoreResponse, StoreError> {
+	async fn store(&self, mut request: StoreRequest) -> Result<StoreResponse, StoreError> {
 		let mut referenced_sequence = Vec::new();
 		let mut failed_sequence = Vec::new();
 
-		for instance in request.instances {
+		while let Some(instance) = request.instances.try_next().await? {
 			let sop_instance_uid = UI::from(instance.meta().media_storage_sop_instance_uid());
 			let sop_class_uid = UI::from(instance.meta().media_storage_sop_class_uid());
 

src/config/defaults.yaml

@@ -6,7 +6,8 @@ server:
   http:
     interface: 0.0.0.0
     port: 8080
-    max-upload-size: 50000000
+    max-upload-size: 5000000000
+    max-part-size: 50000000
     request-timeout: 60000
     graceful-shutdown: true
     base-path: /

src/config/mod.rs

@@ -291,6 +291,7 @@ pub struct HttpServerConfig {
 	pub interface: IpAddr,
 	pub port: u16,
 	pub max_upload_size: usize,
+	pub max_part_size: usize,
 	pub request_timeout: u64,
 	pub graceful_shutdown: bool,
 	pub base_path: String,
@@ -322,8 +323,9 @@ impl Default for HttpServerConfig {
 			interface: IpAddr::from([0, 0, 0, 0]),
 			port: 8080,
 			graceful_shutdown: true,
-			max_upload_size: 50_000_000, // 50 MB
-			request_timeout: 60_000,     // 1 min
+			max_upload_size: 5_000_000_000, // 5 GB
+			max_part_size: 50_000_000,      // 50 MB
+			request_timeout: 60_000,        // 1 min
 			base_path: String::from("/"),
 		}
 	}

src/utils/multipart.rs

@@ -1,4 +1,5 @@
-use axum::extract::{FromRequest, Request};
+use crate::AppState;
+use axum::extract::{FromRef, FromRequest, Request};
 use axum::http::header::CONTENT_TYPE;
 use axum::http::{HeaderValue, StatusCode};
 use axum::response::{IntoResponse, Response};
@@ -51,22 +52,30 @@
 
 impl<S> FromRequest<S> for DicomMultipart<'_>
 where
+	AppState: FromRef<S>,
 	S: Send + Sync,
 {
 	type Rejection = DicomMultipartRejection;
 
-	async fn from_request(request: Request, _state: &S) -> Result<Self, Self::Rejection> {
+	async fn from_request(request: Request, state: &S) -> Result<Self, Self::Rejection> {
 		let boundary = request
 			.headers()
 			.get(CONTENT_TYPE)
 			.map(HeaderValue::to_str)
 			.and_then(Result::ok)
 			.map(DicomMultipart::parse_boundary)
 			.and_then(Result::ok)
 			.ok_or(Self::Rejection::InvalidBoundary)?;
 
+		// The size of the whole request body is limited by `DefaultBodyLimit` (max-upload-size),
+		// the size of each part (i.e. each DICOM instance) is limited here (max-part-size).
+		let max_part_size = AppState::from_ref(state).config.server.http.max_part_size;
+		let constraints = multer::Constraints::new()
+			.size_limit(multer::SizeLimit::new().per_field(max_part_size as u64));
+
 		let stream = request.with_limited_body().into_body();
-		let multipart = multer::Multipart::new(stream.into_data_stream(), boundary);
+		let multipart =
+			multer::Multipart::with_constraints(stream.into_data_stream(), boundary, constraints);
 		Ok(Self(multipart))
 	}
 }

tests/stow.rs

@@ -172,6 +172,44 @@ async fn returns_413_if_max_upload_size_is_exceeded() -> anyhow::Result<()> {
 	.await
 }
 
+#[tokio::test]
+async fn returns_413_if_max_part_size_is_exceeded() -> anyhow::Result<()> {
+	let config = "
+        server:
+          http:
+            port: 0
+            max-part-size: 1
+          dimse:
+            - aet: DICOM-RST
+              interface: 0.0.0.0
+              port: 0
+        aets:
+          - aet: ORTHANC
+            host: 127.0.0.1
+            port: ${ORTHANC_PORT}
+            backend: DIMSE
+	";
+
+	with_test_environment(config, async |client| {
+		let instance = open_file(dicom_test_files::path("pydicom/CT_small.dcm").unwrap()).unwrap();
+		let result = client
+			.store_instances()
+			.with_instances(futures::stream::iter([instance]))
+			.run()
+			.await;
+
+		assert!(matches!(
+			result,
+			Err(DicomWebError::HttpStatusFailure {
+				status_code: StatusCode::PAYLOAD_TOO_LARGE
+			})
+		));
+
+		Ok(())
+	})
+	.await
+}
+
 // Associations must be reused for subsequent instances instead of being rebuilt for every
 // single instance. Validating a pooled association must not require a DIMSE round trip, as
 // the association only negotiates a presentation context for the storage SOP class of the