Use AtomicUsize for read limiter + feature gate it

Author: appaquet

capnp/Cargo.toml

@@ -25,7 +25,7 @@ quickcheck = { version = "0.9", optional = true }
 quickcheck = "0.9"
 
 [features]
-default = ["std"]
+default = ["std", "sync_reader"]
 
 rpc_try = []
 
@@ -37,3 +37,8 @@ unaligned = []
 # with the Rust standard library.
 std = []
 
+# If disabled, message reader will not be `Sync`. This can be used
+# to support platforms where `AtomicUsize` is not available.
+# 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

@@ -20,42 +20,17 @@
 
 use alloc::vec::Vec;
 use core::cell::RefCell;
-use core::sync::atomic::{AtomicU64, Ordering};
 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: u64,
-    pub read: AtomicU64,
-}
-
-impl ReadLimiter {
-    pub fn new(limit: u64) -> ReadLimiter {
-        ReadLimiter { 
-            limit,
-            read: AtomicU64::new(0),
-        }
-    }
-
-    #[inline]
-    pub fn can_read(&self, amount: u64) -> Result<()> {
-        let read = self.read.fetch_add(amount, Ordering::Relaxed) + amount;
-
-        if read > self.limit {
-            Err(Error::failed(format!("read limit exceeded")))
-        } else {
-            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)>;
@@ -134,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,92 @@
+// 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: usize,
+        pub read: 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: limit as usize,
+                read: AtomicUsize::new(0),
+            }
+        }
+
+        #[inline]
+        pub fn can_read(&self, amount: usize) -> Result<()> {
+            let read = self.read.load(Ordering::Relaxed) + amount;
+
+            if read > self.limit {
+                Err(Error::failed(format!("read limit exceeded")))
+            } else {
+                self.read.fetch_add(amount, Ordering::Relaxed);
+                Ok(())
+            }
+        }
+    }
+}
+
+#[cfg(not(feature = "sync_reader"))]
+pub use unsync::ReadLimiter;
+
+#[cfg(not(feature = "sync_reader"))]
+mod unsync {
+    use core::cell::Cell;
+    use crate::{Error, Result};
+
+    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(())
+            }
+        }
+    }
+}