Protection against OOM when reading corrupted bincode

crates/hyperqueue/src/common/serialization.rs

@@ -13,15 +13,17 @@ pub struct DefaultConfig;
 
 impl SerializationConfig for DefaultConfig {
     fn config() -> impl Options {
-        bincode::DefaultOptions::new()
+        bincode::DefaultOptions::new().with_limit(tako::MAX_FRAME_SIZE as u64)
     }
 }
 
 pub struct TrailingAllowedConfig;
 
 impl SerializationConfig for TrailingAllowedConfig {
     fn config() -> impl Options {
-        bincode::DefaultOptions::new().allow_trailing_bytes()
+        bincode::DefaultOptions::new()
+            .allow_trailing_bytes()
+            .with_limit(tako::MAX_FRAME_SIZE as u64)
     }
 }
 

crates/hyperqueue/src/server/event/journal/read.rs

@@ -59,7 +59,7 @@ impl Iterator for &mut JournalReader {
 
     fn next(&mut self) -> Option<Self::Item> {
         self.position = self.source.stream_position().unwrap();
-        if self.position == self.size {
+        if self.position >= self.size {
             return None;
         }
         match EventSerializationConfig::config().deserialize_from(&mut self.source) {

crates/hyperqueue/src/transfer/protocol.rs

@@ -4,5 +4,5 @@ use tokio_util::codec::length_delimited::Builder;
 pub fn make_protocol_builder() -> Builder {
     *LengthDelimitedCodec::builder()
         .little_endian()
-        .max_frame_length(128 * 1024 * 1024)
+        .max_frame_length(tako::MAX_FRAME_SIZE)
 }

crates/tako/src/internal/tests/integration/utils/task.rs

@@ -100,6 +100,7 @@ pub fn build_task_def_from_config(
         cwd,
     };
     let body = bincode::DefaultOptions::new()
+        .with_limit(crate::MAX_FRAME_SIZE as u64)
         .serialize(&program_def)
         .unwrap();
 

crates/tako/src/internal/tests/integration/utils/worker.rs

@@ -266,6 +266,7 @@ impl TaskLauncher for TestTaskLauncher {
                 ctx.body().len(),
             );
             bincode::DefaultOptions::new()
+                .with_limit(crate::MAX_FRAME_SIZE as u64)
                 .deserialize(ctx.body())
                 .map_err(|_| DsError::GenericError("Body deserialization failed".into()))?
         };

crates/tako/src/internal/transfer/auth.rs

@@ -244,6 +244,7 @@ where
     T: serde::Serialize + ?Sized,
 {
     DefaultOptions::new()
+        .with_limit(crate::MAX_FRAME_SIZE as u64)
         .with_fixint_encoding()
         .serialize(value)
         .map_err(|e| format!("Serialization failed: {e:?}").into())
@@ -255,6 +256,7 @@ where
     T: Deserialize<'a>,
 {
     DefaultOptions::new()
+        .with_limit(crate::MAX_FRAME_SIZE as u64)
         .with_fixint_encoding()
         .deserialize(bytes)
         .map_err(|e| format!("Deserialization failed: {e:?}, data {bytes:?}").into())

crates/tako/src/internal/transfer/transport.rs

@@ -3,5 +3,5 @@ use tokio_util::codec::length_delimited::{Builder, LengthDelimitedCodec};
 pub(crate) fn make_protocol_builder() -> Builder {
     *LengthDelimitedCodec::builder()
         .little_endian()
-        .max_frame_length(128 * 1024 * 1024)
+        .max_frame_length(crate::MAX_FRAME_SIZE)
 }

crates/tako/src/lib.rs

@@ -26,6 +26,8 @@ pub type PriorityTuple = (Priority, Priority); // user priority, scheduler prior
 pub type Error = internal::common::error::DsError;
 pub type Result<T> = std::result::Result<T, Error>;
 
+pub const MAX_FRAME_SIZE: usize = 128 * 1024 * 1024;
+
 pub mod resources {
     pub use crate::internal::common::resources::{
         AMD_GPU_RESOURCE_NAME, Allocation, AllocationRequest, CPU_RESOURCE_ID, CPU_RESOURCE_NAME,