diff --git a/capnp/Cargo.toml b/capnp/Cargo.toml index 36629f1ba..1512c040f 100644 --- a/capnp/Cargo.toml +++ b/capnp/Cargo.toml @@ -37,3 +37,7 @@ unaligned = [] # with the Rust standard library. std = [] +# If enabled, message reader will be `Sync`. +# To be replaced by #[cfg(target_has_atomic = ...)] once it lands. +# See: https://github.com/rust-lang/rust/issues/32976 +sync_reader = [] \ No newline at end of file diff --git a/capnp/src/private/arena.rs b/capnp/src/private/arena.rs index 73a0ece5e..2071ec30c 100644 --- a/capnp/src/private/arena.rs +++ b/capnp/src/private/arena.rs @@ -19,38 +19,18 @@ // THE SOFTWARE. use alloc::vec::Vec; -use core::cell::{Cell, RefCell}; +use core::cell::RefCell; use core::slice; use core::u64; use crate::private::units::*; +use crate::private::read_limiter::ReadLimiter; use crate::message; use crate::message::{Allocator, ReaderSegments}; use crate::{Error, OutputSegments, Result}; pub type SegmentId = u32; -pub struct ReadLimiter { - pub limit: Cell, -} - -impl ReadLimiter { - pub fn new(limit: u64) -> ReadLimiter { - ReadLimiter { limit: Cell::new(limit) } - } - - #[inline] - pub fn can_read(&self, amount: u64) -> Result<()> { - let current = self.limit.get(); - if amount > current { - Err(Error::failed(format!("read limit exceeded"))) - } else { - self.limit.set(current - amount); - Ok(()) - } - } -} - pub trait ReaderArena { // return pointer to start of segment, and number of words in that segment fn get_segment(&self, id: u32) -> Result<(*const u8, u32)>; @@ -129,12 +109,12 @@ impl ReaderArena for ReaderArenaImpl where S: ReaderSegments { if !(start >= this_start && start - this_start + size <= this_size) { Err(Error::failed(format!("message contained out-of-bounds pointer"))) } else { - self.read_limiter.can_read(size_in_words as u64) + self.read_limiter.can_read(size_in_words) } } fn amplified_read(&self, virtual_amount: u64) -> Result<()> { - self.read_limiter.can_read(virtual_amount) + self.read_limiter.can_read(virtual_amount as usize) } } diff --git a/capnp/src/private/mod.rs b/capnp/src/private/mod.rs index 8d6465031..252fc19d1 100644 --- a/capnp/src/private/mod.rs +++ b/capnp/src/private/mod.rs @@ -29,6 +29,7 @@ mod primitive; pub mod layout; mod mask; pub mod units; +mod read_limiter; mod zero; #[cfg(test)] diff --git a/capnp/src/private/read_limiter.rs b/capnp/src/private/read_limiter.rs new file mode 100644 index 000000000..7522869a5 --- /dev/null +++ b/capnp/src/private/read_limiter.rs @@ -0,0 +1,96 @@ +// Copyright (c) 2013-2015 Sandstorm Development Group, Inc. and contributors +// Licensed under the MIT License: +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in +// all copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +// THE SOFTWARE. + +#[cfg(feature = "sync_reader")] +pub use sync::ReadLimiter; + +#[cfg(feature = "sync_reader")] +mod sync { + use crate::{Error, Result}; + use core::sync::atomic::{AtomicUsize, Ordering}; + + pub struct ReadLimiter { + pub limit: AtomicUsize, + } + + impl ReadLimiter { + pub fn new(limit: u64) -> ReadLimiter { + if limit > core::usize::MAX as u64 { + panic!("traversal_limit_in_words cannot be bigger than core::usize::MAX") + } + + ReadLimiter { + limit: AtomicUsize::new(limit as usize), + } + } + + #[inline] + pub fn can_read(&self, amount: usize) -> Result<()> { + let cur_limit = self.limit.load(Ordering::Relaxed); + if cur_limit < amount { + return Err(Error::failed(format!("read limit exceeded"))); + } + + let prev_limit = self.limit.fetch_sub(amount, Ordering::Relaxed); + if prev_limit < amount { + // if the previous limit was lower than the amount we read, the limit has underflowed + // and wrapped around so we need to reset it to 0 for next reader to fail + self.limit.store(0, Ordering::Relaxed); + return Err(Error::failed(format!("read limit exceeded"))); + } + + Ok(()) + } + } +} + +#[cfg(not(feature = "sync_reader"))] +pub use unsync::ReadLimiter; + +#[cfg(not(feature = "sync_reader"))] +mod unsync { + use crate::{Error, Result}; + use core::cell::Cell; + + pub struct ReadLimiter { + pub limit: Cell, + } + + impl ReadLimiter { + pub fn new(limit: u64) -> ReadLimiter { + ReadLimiter { + limit: Cell::new(limit), + } + } + + #[inline] + pub fn can_read(&self, amount: usize) -> Result<()> { + let amount = amount as u64; + let current = self.limit.get(); + if amount > current { + Err(Error::failed(format!("read limit exceeded"))) + } else { + self.limit.set(current - amount); + Ok(()) + } + } + } +}