From f09bd6bc887c8370f6f82b4e942b504213eb8164 Mon Sep 17 00:00:00 2001
From: Heikki Linnakangas <heikki@neon.tech>
Date: Tue, 30 Aug 2022 18:44:06 +0300
Subject: [PATCH] Fix size checks in the "local" remote storage implementation.

The code correctly detected too short and too long inputs, but the error
message was bogus for the case the input stream was too long:

    Error: Provided stream has actual size 5 fthat is smaller than the given stream size 4

That check was only supposed to check for too small inputs, but it in
fact caught too long inputs too. That was good, because the check
below that that was supposed to check for too long inputs was in fact
broken, and never did anything. It tried to read input a buffer of
size 0, to check if there is any extra data, but reading to a
zero-sized buffer always returns 0.
---
 libs/remote_storage/src/local_fs.rs | 49 ++++++++++++++++++++++-------
 1 file changed, 37 insertions(+), 12 deletions(-)

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> {
         LocalFs::new(tempdir()?.path().to_owned(), tempdir()?.path().to_owned())
     }