diff --git a/program-runtime/src/compute_budget.rs b/program-runtime/src/compute_budget.rs index 136363f521..137061f3ae 100644 --- a/program-runtime/src/compute_budget.rs +++ b/program-runtime/src/compute_budget.rs @@ -36,6 +36,8 @@ pub struct ComputeBudget { pub sha256_base_cost: u64, /// Incremental number of units consumed by SHA256 (based on bytes) pub sha256_byte_cost: u64, + /// Maximum number of slices hashed per syscall + pub sha256_max_slices: u64, /// Maximum BPF to BPF call depth pub max_call_depth: usize, /// Size of a stack frame in bytes, must match the size specified in the LLVM BPF backend @@ -82,6 +84,7 @@ impl ComputeBudget { max_invoke_depth: 4, sha256_base_cost: 85, sha256_byte_cost: 1, + sha256_max_slices: 20_000, max_call_depth: 64, stack_frame_size: 4_096, log_pubkey_units: 100, diff --git a/programs/bpf_loader/src/syscalls.rs b/programs/bpf_loader/src/syscalls.rs index b37e7fd9c6..4fe4f004b3 100644 --- a/programs/bpf_loader/src/syscalls.rs +++ b/programs/bpf_loader/src/syscalls.rs @@ -89,6 +89,8 @@ pub enum SyscallError { CopyOverlapping, #[error("Return data too large ({0} > {1})")] ReturnDataTooLarge(u64, u64), + #[error("Hashing too many sequences")] + TooManySlices, } impl From for EbpfError { fn from(error: SyscallError) -> Self { @@ -1039,6 +1041,20 @@ impl<'a, 'b> SyscallObject for SyscallSha256<'a, 'b> { result ); let compute_budget = invoke_context.get_compute_budget(); + if invoke_context + .feature_set + .is_active(&update_syscall_base_costs::id()) + && compute_budget.sha256_max_slices < vals_len + { + ic_msg!( + invoke_context, + "Sha256 hashing {} sequences in one syscall is over the limit {}", + vals_len, + compute_budget.sha256_max_slices, + ); + *result = Err(SyscallError::TooManySlices.into()); + return; + } question_mark!( invoke_context .get_compute_meter() @@ -1072,7 +1088,18 @@ impl<'a, 'b> SyscallObject for SyscallSha256<'a, 'b> { ), result ); - let cost = compute_budget.sha256_byte_cost * (val.len() as u64 / 2); + let cost = if invoke_context + .feature_set + .is_active(&update_syscall_base_costs::id()) + { + compute_budget.mem_op_base_cost.max( + compute_budget + .sha256_byte_cost + .saturating_mul(val.len() as u64 / 2), + ) + } else { + compute_budget.sha256_byte_cost * (val.len() as u64 / 2) + }; question_mark!(invoke_context.get_compute_meter().consume(cost), result); hasher.hash(bytes); } @@ -1268,6 +1295,20 @@ impl<'a, 'b> SyscallObject for SyscallKeccak256<'a, 'b> { result ); let compute_budget = invoke_context.get_compute_budget(); + if invoke_context + .feature_set + .is_active(&update_syscall_base_costs::id()) + && compute_budget.sha256_max_slices < vals_len + { + ic_msg!( + invoke_context, + "Keccak256 hashing {} sequences in one syscall is over the limit {}", + vals_len, + compute_budget.sha256_max_slices, + ); + *result = Err(SyscallError::TooManySlices.into()); + return; + } question_mark!( invoke_context .get_compute_meter() @@ -1306,9 +1347,19 @@ impl<'a, 'b> SyscallObject for SyscallKeccak256<'a, 'b> { ), result ); - let cost = compute_budget.sha256_byte_cost * (val.len() as u64 / 2); + let cost = if invoke_context + .feature_set + .is_active(&update_syscall_base_costs::id()) + { + compute_budget.mem_op_base_cost.max( + compute_budget + .sha256_byte_cost + .saturating_mul(val.len() as u64 / 2), + ) + } else { + compute_budget.sha256_byte_cost * (val.len() as u64 / 2) + }; question_mark!(invoke_context.get_compute_meter().consume(cost), result); - hasher.hash(bytes); } } @@ -1667,6 +1718,20 @@ impl<'a, 'b> SyscallObject for SyscallBlake3<'a, 'b> { result ); let compute_budget = invoke_context.get_compute_budget(); + if invoke_context + .feature_set + .is_active(&update_syscall_base_costs::id()) + && compute_budget.sha256_max_slices < vals_len + { + ic_msg!( + invoke_context, + "Blake3 hashing {} sequences in one syscall is over the limit {}", + vals_len, + compute_budget.sha256_max_slices, + ); + *result = Err(SyscallError::TooManySlices.into()); + return; + } question_mark!( invoke_context .get_compute_meter() @@ -1705,10 +1770,19 @@ impl<'a, 'b> SyscallObject for SyscallBlake3<'a, 'b> { ), result ); - - let cost = compute_budget.sha256_byte_cost * (val.len() as u64 / 2); + let cost = if invoke_context + .feature_set + .is_active(&update_syscall_base_costs::id()) + { + compute_budget.mem_op_base_cost.max( + compute_budget + .sha256_byte_cost + .saturating_mul(val.len() as u64 / 2), + ) + } else { + compute_budget.sha256_byte_cost * (val.len() as u64 / 2) + }; question_mark!(invoke_context.get_compute_meter().consume(cost), result); - hasher.hash(bytes); } } @@ -3687,8 +3761,12 @@ mod tests { .borrow_mut() .mock_set_remaining( (invoke_context.get_compute_budget().sha256_base_cost - + ((bytes1.len() + bytes2.len()) as u64 / 2) - * invoke_context.get_compute_budget().sha256_byte_cost) + + invoke_context.get_compute_budget().mem_op_base_cost.max( + invoke_context + .get_compute_budget() + .sha256_byte_cost + .saturating_mul((bytes1.len() + bytes2.len()) as u64 / 2), + )) * 4, ); invoke_context