diff --git a/libs/remote_storage/src/local_fs.rs b/libs/remote_storage/src/local_fs.rs index 07b04084b9..a65d0887af 100644 --- a/libs/remote_storage/src/local_fs.rs +++ b/libs/remote_storage/src/local_fs.rs @@ -150,8 +150,7 @@ impl RemoteStorage for LocalFs { ); let from_size_bytes = from_size_bytes as u64; - // Require to read 1 byte more than the expected to check later, that the stream and its size match. - let mut buffer_to_read = from.take(from_size_bytes + 1); + let mut buffer_to_read = from.take(from_size_bytes); let bytes_read = io::copy(&mut buffer_to_read, &mut destination) .await @@ -162,17 +161,15 @@ impl RemoteStorage for LocalFs { ) })?; + if bytes_read < from_size_bytes { + bail!("Provided stream was shorter than expected: {bytes_read} vs {from_size_bytes} bytes"); + } + // Check if there is any extra data after the given size. + let mut from = buffer_to_read.into_inner(); + let extra_read = from.read(&mut [1]).await?; ensure!( - bytes_read == from_size_bytes, - "Provided stream has actual size {} fthat is smaller than the given stream size {}", - bytes_read, - from_size_bytes - ); - - ensure!( - buffer_to_read.read(&mut [0]).await? == 0, - "Provided stream has bigger size than the given stream size {}", - from_size_bytes + extra_read == 0, + "Provided stream was larger than expected: expected {from_size_bytes} bytes", ); destination.flush().await.with_context(|| { @@ -609,6 +606,34 @@ mod fs_tests { Ok(()) } + #[tokio::test] + async fn upload_file_negatives() -> anyhow::Result<()> { + let storage = create_storage()?; + + let id = storage.remote_object_id(&storage.working_directory.join("dummy"))?; + let content = std::io::Cursor::new(b"12345"); + + // Check that you get an error if the size parameter doesn't match the actual + // size of the stream. + storage + .upload(content.clone(), 0, &id, None) + .await + .expect_err("upload with zero size succeeded"); + storage + .upload(content.clone(), 4, &id, None) + .await + .expect_err("upload with too short size succeeded"); + storage + .upload(content.clone(), 6, &id, None) + .await + .expect_err("upload with too large size succeeded"); + + // Correct size is 5, this should succeed. + storage.upload(content, 5, &id, None).await?; + + Ok(()) + } + fn create_storage() -> anyhow::Result { LocalFs::new(tempdir()?.path().to_owned(), tempdir()?.path().to_owned()) }