Sync ReadLimiter

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 = []

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<u64>,
-}
-
-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 <S> ReaderArena for ReaderArenaImpl<S> 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)
     }
 }
 

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)]

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<u64>,
+    }
+
+    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(())
+            }
+        }
+    }
+}