Add blob size limits. (#2705) (#2728)

* Add blob size limits. (#2705)

* Add a blob size limit.

* Add a bytecode size limit.

* Add unit tests for limits.

* Don't enforce the limit for already published bytecode.

* Simplify LimitedWriter; add unit test.

* Add decompressed_size_at_most.

* Update and copy comment about #2710.

* Don't cache blobs that are too large. (#2726)

* Don't cache blobs that are too large.

* Check size first, then decompress.

Author: afck

linera-base/src/data_types.rs

@@ -37,6 +37,7 @@ use crate::{
         ApplicationId, BlobId, BlobType, BytecodeId, Destination, GenericApplicationId, MessageId,
         UserApplicationId,
     },
+    limited_writer::{LimitedWriter, LimitedWriterError},
     time::{Duration, SystemTime},
 };
 
@@ -861,14 +862,8 @@ impl fmt::Debug for Bytecode {
 #[derive(Error, Debug)]
 pub enum DecompressionError {
     /// Compressed bytecode is invalid, and could not be decompressed.
-    #[cfg(not(target_arch = "wasm32"))]
-    #[error("Bytecode could not be decompressed")]
-    InvalidCompressedBytecode(#[source] io::Error),
-
-    /// Compressed bytecode is invalid, and could not be decompressed.
-    #[cfg(target_arch = "wasm32")]
-    #[error("Bytecode could not be decompressed")]
-    InvalidCompressedBytecode(#[from] ruzstd::frame_decoder::FrameDecoderError),
+    #[error("Bytecode could not be decompressed: {0}")]
+    InvalidCompressedBytecode(#[from] io::Error),
 }
 
 /// A compressed WebAssembly module's bytecode.
@@ -880,20 +875,53 @@ pub struct CompressedBytecode {
     pub compressed_bytes: Vec<u8>,
 }
 
+#[cfg(not(target_arch = "wasm32"))]
 impl CompressedBytecode {
+    /// Returns `true` if the decompressed size does not exceed the limit.
+    pub fn decompressed_size_at_most(&self, limit: u64) -> Result<bool, DecompressionError> {
+        let mut decoder = zstd::stream::Decoder::new(&*self.compressed_bytes)?;
+        let limit = usize::try_from(limit).unwrap_or(usize::MAX);
+        let mut writer = LimitedWriter::new(io::sink(), limit);
+        match io::copy(&mut decoder, &mut writer) {
+            Ok(_) => Ok(true),
+            Err(error) => {
+                error.downcast::<LimitedWriterError>()?;
+                Ok(false)
+            }
+        }
+    }
+
     /// Decompresses a [`CompressedBytecode`] into a [`Bytecode`].
-    #[cfg(not(target_arch = "wasm32"))]
     pub fn decompress(&self) -> Result<Bytecode, DecompressionError> {
         #[cfg(with_metrics)]
         let _decompression_latency = BYTECODE_DECOMPRESSION_LATENCY.measure_latency();
-        let bytes = zstd::stream::decode_all(&*self.compressed_bytes)
-            .map_err(DecompressionError::InvalidCompressedBytecode)?;
+        let bytes = zstd::stream::decode_all(&*self.compressed_bytes)?;
 
         Ok(Bytecode { bytes })
     }
+}
+
+#[cfg(target_arch = "wasm32")]
+impl CompressedBytecode {
+    /// Returns `true` if the decompressed size does not exceed the limit.
+    pub fn decompressed_size_at_most(&self, limit: u64) -> Result<bool, DecompressionError> {
+        let compressed_bytes = &*self.compressed_bytes;
+        let limit = usize::try_from(limit).unwrap_or(usize::MAX);
+        let mut writer = LimitedWriter::new(io::sink(), limit);
+        let mut decoder = ruzstd::streaming_decoder::StreamingDecoder::new(compressed_bytes)
+            .map_err(io::Error::other)?;
+
+        // TODO(#2710): Decode multiple frames, if present
+        match io::copy(&mut decoder, &mut writer) {
+            Ok(_) => Ok(true),
+            Err(error) => {
+                error.downcast::<LimitedWriterError>()?;
+                Ok(false)
+            }
+        }
+    }
 
     /// Decompresses a [`CompressedBytecode`] into a [`Bytecode`].
-    #[cfg(target_arch = "wasm32")]
     pub fn decompress(&self) -> Result<Bytecode, DecompressionError> {
         use ruzstd::{io::Read, streaming_decoder::StreamingDecoder};
 
@@ -902,10 +930,9 @@ impl CompressedBytecode {
 
         let compressed_bytes = &*self.compressed_bytes;
         let mut bytes = Vec::new();
-        let mut decoder = StreamingDecoder::new(compressed_bytes)?;
+        let mut decoder = StreamingDecoder::new(compressed_bytes).map_err(io::Error::other)?;
 
-        // Decode multiple frames, if present
-        // (https://github.com/KillingSpark/zstd-rs/issues/57)
+        // TODO(#2710): Decode multiple frames, if present
         while !decoder.get_ref().is_empty() {
             decoder
                 .read_to_end(&mut bytes)
@@ -1035,6 +1062,17 @@ impl BlobContent {
     pub fn blob_bytes(&self) -> BlobBytes {
         BlobBytes(self.inner_bytes())
     }
+
+    /// Returns the size of the blob content in bytes.
+    pub fn size(&self) -> usize {
+        match self {
+            BlobContent::Data(bytes) => bytes.len(),
+            BlobContent::ContractBytecode(compressed_bytecode)
+            | BlobContent::ServiceBytecode(compressed_bytecode) => {
+                compressed_bytecode.compressed_bytes.len()
+            }
+        }
+    }
 }
 
 impl From<Blob> for BlobContent {

linera-base/src/lib.rs

@@ -18,6 +18,7 @@ pub mod crypto;
 pub mod data_types;
 mod graphql;
 pub mod identifiers;
+mod limited_writer;
 pub mod ownership;
 #[cfg(not(target_arch = "wasm32"))]
 pub mod port;

linera-base/src/limited_writer.rs

@@ -0,0 +1,68 @@
+// Copyright (c) Zefchain Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+use std::io::{self, Write};
+
+use thiserror::Error;
+
+use crate::ensure;
+
+#[derive(Error, Debug)]
+#[error("Writer limit exceeded")]
+pub struct LimitedWriterError;
+
+/// Custom writer that enforces a byte limit.
+pub struct LimitedWriter<W: Write> {
+    inner: W,
+    limit: usize,
+    written: usize,
+}
+
+impl<W: Write> LimitedWriter<W> {
+    pub fn new(inner: W, limit: usize) -> Self {
+        Self {
+            inner,
+            limit,
+            written: 0,
+        }
+    }
+}
+
+impl<W: Write> Write for LimitedWriter<W> {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        // Calculate the number of bytes we can write without exceeding the limit.
+        // Fail if the buffer doesn't fit.
+        ensure!(
+            self.limit
+                .checked_sub(self.written)
+                .is_some_and(|remaining| buf.len() <= remaining),
+            io::Error::other(LimitedWriterError)
+        );
+        // Forward to the inner writer.
+        let n = self.inner.write(buf)?;
+        self.written += n;
+        Ok(n)
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        self.inner.flush()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_limited_writer() {
+        let mut out_buffer = Vec::new();
+        let mut writer = LimitedWriter::new(&mut out_buffer, 5);
+        assert_eq!(writer.write(b"foo").unwrap(), 3);
+        assert_eq!(writer.write(b"ba").unwrap(), 2);
+        assert!(writer
+            .write(b"r")
+            .unwrap_err()
+            .downcast::<LimitedWriterError>()
+            .is_ok());
+    }
+}

linera-core/src/chain_worker/state/temporary_changes.rs

@@ -3,10 +3,10 @@
 
 //! Operations that don't persist any changes to the chain state.
 
-use std::borrow::Cow;
+use std::{borrow::Cow, collections::BTreeSet};
 
 use linera_base::{
-    data_types::{ArithmeticError, Timestamp, UserApplicationDescription},
+    data_types::{ArithmeticError, BlobContent, Timestamp, UserApplicationDescription},
     ensure,
     identifiers::{GenericApplicationId, UserApplicationId},
 };
@@ -28,6 +28,7 @@ use {
 
 use super::{check_block_epoch, ChainWorkerState};
 use crate::{
+    client::{MAXIMUM_BLOB_SIZE, MAXIMUM_BYTECODE_SIZE},
     data_types::{ChainInfo, ChainInfoQuery, ChainInfoResponse},
     worker::WorkerError,
 };
@@ -207,12 +208,20 @@ where
         self.0.chain.remove_bundles_from_inboxes(block).await?;
         // Verify that all required bytecode hashed certificate values and blobs are available, and no
         // unrelated ones provided.
+        let published_blob_ids = block.published_blob_ids();
         self.0
-            .check_no_missing_blobs(block.published_blob_ids(), blobs)
+            .check_no_missing_blobs(published_blob_ids.clone(), blobs)
             .await?;
         for blob in blobs {
+            Self::check_blob_size(blob.content())?;
             self.0.cache_recent_blob(Cow::Borrowed(blob)).await;
         }
+        let checked_blobs = blobs.iter().map(|blob| blob.id()).collect::<BTreeSet<_>>();
+        for blob in self.0.get_blobs(published_blob_ids).await? {
+            if !checked_blobs.contains(&blob.id()) {
+                Self::check_blob_size(blob.content())?;
+            }
+        }
 
         let local_time = self.0.storage.clock().current_time();
         ensure!(
@@ -333,6 +342,26 @@ where
         }
         Ok(ChainInfoResponse::new(info, self.0.config.key_pair()))
     }
+
+    fn check_blob_size(content: &BlobContent) -> Result<(), WorkerError> {
+        ensure!(
+            u64::try_from(content.size())
+                .ok()
+                .is_some_and(|size| size <= MAXIMUM_BLOB_SIZE),
+            WorkerError::BlobTooLarge
+        );
+        match content {
+            BlobContent::ContractBytecode(compressed_bytecode)
+            | BlobContent::ServiceBytecode(compressed_bytecode) => {
+                ensure!(
+                    compressed_bytecode.decompressed_size_at_most(MAXIMUM_BYTECODE_SIZE)?,
+                    WorkerError::BytecodeTooLarge
+                );
+            }
+            BlobContent::Data(_) => {}
+        }
+        Ok(())
+    }
 }
 
 impl<StorageClient> Drop for ChainWorkerStateWithTemporaryChanges<'_, StorageClient>

linera-core/src/client/mod.rs

@@ -83,6 +83,13 @@ mod chain_state;
 #[path = "../unit_tests/client_tests.rs"]
 mod client_tests;
 
+const MEBIBYTE: u64 = 1024 * 1024;
+
+/// The maximum size of a data or bytecode blob, in bytes.
+pub(crate) const MAXIMUM_BLOB_SIZE: u64 = 3 * MEBIBYTE;
+/// The maximum size of decompressed bytecode, in bytes.
+pub(crate) const MAXIMUM_BYTECODE_SIZE: u64 = 30 * MEBIBYTE;
+
 #[cfg(with_metrics)]
 mod metrics {
     use std::sync::LazyLock;

linera-core/src/unit_tests/client_tests.rs

@@ -41,7 +41,7 @@ use crate::test_utils::ServiceStorageBuilder;
 use crate::{
     client::{
         BlanketMessagePolicy, ChainClient, ChainClientError, ClientOutcome, MessageAction,
-        MessagePolicy,
+        MessagePolicy, MAXIMUM_BLOB_SIZE,
     },
     local_node::LocalNodeError,
     node::{
@@ -2458,5 +2458,9 @@ where
     assert_eq!(executed_block.block.incoming_bundles.len(), 1);
     assert_eq!(executed_block.required_blob_ids().len(), 1);
 
+    let large_blob_bytes = vec![0; MAXIMUM_BLOB_SIZE as usize + 1];
+    let result = client1.publish_data_blob(large_blob_bytes).await;
+    assert_matches!(result, Err(ChainClientError::LocalNodeError(_)));
+
     Ok(())
 }

linera-core/src/unit_tests/wasm_client_tests.rs

@@ -41,7 +41,10 @@ use crate::client::client_tests::RocksDbStorageBuilder;
 use crate::client::client_tests::ScyllaDbStorageBuilder;
 #[cfg(feature = "storage-service")]
 use crate::client::client_tests::ServiceStorageBuilder;
-use crate::client::client_tests::{MemoryStorageBuilder, StorageBuilder, TestBuilder};
+use crate::client::{
+    client_tests::{MemoryStorageBuilder, StorageBuilder, TestBuilder},
+    ChainClientError, MAXIMUM_BYTECODE_SIZE,
+};
 
 #[cfg_attr(feature = "wasmer", test_case(WasmRuntime::Wasmer ; "wasmer"))]
 #[cfg_attr(feature = "wasmtime", test_case(WasmRuntime::Wasmtime ; "wasmtime"))]
@@ -148,6 +151,18 @@ where
     let balance_after_init = creator.local_balance().await?;
     assert!(balance_after_init < balance_after_messaging);
 
+    let large_bytecode = Bytecode::new(vec![0; MAXIMUM_BYTECODE_SIZE as usize + 1]);
+    let small_bytecode = Bytecode::new(vec![]);
+    // Publishing bytecode that exceeds the limit fails.
+    let result = publisher
+        .publish_bytecode(large_bytecode.clone(), small_bytecode.clone())
+        .await;
+    assert_matches!(result, Err(ChainClientError::LocalNodeError(_)));
+    let result = publisher
+        .publish_bytecode(small_bytecode, large_bytecode)
+        .await;
+    assert_matches!(result, Err(ChainClientError::LocalNodeError(_)));
+
     Ok(())
 }
 

linera-core/src/worker.rs

@@ -14,7 +14,9 @@ use std::{
 use linera_base::crypto::PublicKey;
 use linera_base::{
     crypto::{CryptoHash, KeyPair},
-    data_types::{ArithmeticError, Blob, BlockHeight, Round, UserApplicationDescription},
+    data_types::{
+        ArithmeticError, Blob, BlockHeight, DecompressionError, Round, UserApplicationDescription,
+    },
     doc_scalar,
     identifiers::{BlobId, ChainId, Owner, UserApplicationId},
     time::timer::{sleep, timeout},
@@ -212,6 +214,12 @@ pub enum WorkerError {
     FullChainWorkerCache,
     #[error("Failed to join spawned worker task")]
     JoinError,
+    #[error("Blob exceeds size limit")]
+    BlobTooLarge,
+    #[error("Bytecode exceeds size limit")]
+    BytecodeTooLarge,
+    #[error(transparent)]
+    Decompression(#[from] DecompressionError),
 }
 
 impl From<linera_chain::ChainError> for WorkerError {