touch up atomic read limiter

dwrensha · dwrensha · commit 37e711035df4 · 2020-12-19T00:23:29.000-05:00
diff --git a/capnp-futures/src/serialize.rs b/capnp-futures/src/serialize.rs
@@ -95,10 +95,12 @@ async fn read_segment_table<R>(mut reader: R,
     // Don't accept a message which the receiver couldn't possibly traverse without hitting the
     // traversal limit. Without this check, a malicious client could transmit a very large segment
     // size to make the receiver allocate excessive space and possibly crash.
-    if segment_lengths_builder.total_words() as u64 > options.traversal_limit_in_words  {
-        return Err(Error::failed(
-            format!("Message has {} words, which is too large. To increase the limit on the \
-             receiving end, see capnp::message::ReaderOptions.", segment_lengths_builder.total_words())))
+    if let Some(traversal_limit_in_words) = options.traversal_limit_in_words {
+        if segment_lengths_builder.total_words() > traversal_limit_in_words  {
+            return Err(Error::failed(
+                format!("Message has {} words, which is too large. To increase the limit on the \
+                         receiving end, see capnp::message::ReaderOptions.", segment_lengths_builder.total_words())))
+        }
     }
 
     Ok(Some(segment_lengths_builder))
diff --git a/capnp/src/message.rs b/capnp/src/message.rs
@@ -50,7 +50,7 @@ pub struct ReaderOptions {
     /// Together with sensible coding practices (e.g. trying to avoid calling sub-object getters
     /// multiple times, which is expensive anyway), this should provide adequate protection without
     /// inconvenience.
-    pub traversal_limit_in_words: u64,
+    pub traversal_limit_in_words: Option<usize>,
 
     /// Limits how deeply nested a message structure can be, e.g. structs containing other structs or
     /// lists of structs.
@@ -64,7 +64,7 @@ pub struct ReaderOptions {
 }
 
 pub const DEFAULT_READER_OPTIONS: ReaderOptions =
-    ReaderOptions { traversal_limit_in_words: 8 * 1024 * 1024, nesting_limit: 64 };
+    ReaderOptions { traversal_limit_in_words: Some(8 * 1024 * 1024), nesting_limit: 64 };
 
 
 impl Default for ReaderOptions {
@@ -81,7 +81,7 @@ impl ReaderOptions {
         self
     }
 
-    pub fn traversal_limit_in_words<'a>(&'a mut self, value: u64) -> &'a mut ReaderOptions {
+    pub fn traversal_limit_in_words<'a>(&'a mut self, value: Option<usize>) -> &'a mut ReaderOptions {
         self.traversal_limit_in_words = value;
         self
     }
@@ -382,7 +382,7 @@ impl <A> Builder<A> where A: Allocator {
 
     pub fn into_reader(self) -> Reader<Builder<A>> {
         Reader::new(self, ReaderOptions {
-            traversal_limit_in_words: u64::max_value(),
+            traversal_limit_in_words: None,
             nesting_limit: i32::max_value()
         })
     }
diff --git a/capnp/src/private/read_limiter.rs b/capnp/src/private/read_limiter.rs
@@ -28,35 +28,45 @@ mod sync {
     use core::sync::atomic::{AtomicUsize, Ordering};
 
     pub struct ReadLimiter {
-        pub limit: AtomicUsize,
+        limit: AtomicUsize,
+        error_on_limit_exceeded: bool,
     }
 
     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),
+        pub fn new(limit: Option<usize>) -> ReadLimiter {
+            match limit {
+                Some(value) => {
+                    ReadLimiter {
+                        limit: AtomicUsize::new(value),
+                        error_on_limit_exceeded: true,
+                    }
+                }
+                None => {
+                    ReadLimiter {
+                        limit: AtomicUsize::new(usize::MAX),
+                        error_on_limit_exceeded: false,
+                    }
+                }
             }
         }
 
         #[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")));
-            }
+            // We use separate AtomicUsize::load() and AtomicUsize::store() steps, which may
+            // result in undercounting reads if multiple threads are reading at the same.
+            // That's okay -- a denial of service attack will eventually hit the limit anyway.
+            //
+            // We could instead do a single fetch_sub() step, but that seems to be slower.
 
-            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);
+            let current = self.limit.load(Ordering::Relaxed);
+            if amount > current && self.error_on_limit_exceeded {
                 return Err(Error::failed(format!("read limit exceeded")));
+            } else {
+                // The common case is current >= amount. Note that we only branch once in that case.
+                // If we combined the fields into an Option<AtomicUsize>, we would
+                // need to branch twice in the common case.
+                self.limit.store(current.wrapping_sub(amount), Ordering::Relaxed);
             }
-
             Ok(())
         }
     }
@@ -71,24 +81,38 @@ mod unsync {
     use core::cell::Cell;
 
     pub struct ReadLimiter {
-        pub limit: Cell<u64>,
+        limit: Cell<usize>,
+        error_on_limit_exceeded: bool,
     }
 
     impl ReadLimiter {
-        pub fn new(limit: u64) -> ReadLimiter {
-            ReadLimiter {
-                limit: Cell::new(limit),
+        pub fn new(limit: Option<usize>) -> ReadLimiter {
+            match limit {
+                Some(value) => {
+                    ReadLimiter {
+                        limit: Cell::new(value),
+                        error_on_limit_exceeded: true,
+                    }
+                }
+                None => {
+                    ReadLimiter {
+                        limit: Cell::new(usize::MAX),
+                        error_on_limit_exceeded: false,
+                    }
+                }
             }
         }
 
         #[inline]
         pub fn can_read(&self, amount: usize) -> Result<()> {
-            let amount = amount as u64;
             let current = self.limit.get();
-            if amount > current {
+            if amount > current && self.error_on_limit_exceeded {
                 Err(Error::failed(format!("read limit exceeded")))
             } else {
-                self.limit.set(current - amount);
+                // The common case is current >= amount. Note that we only branch once in that case.
+                // If we combined the fields into an Option<Cell<usize>>, we would
+                // need to branch twice in the common case.
+                self.limit.set(current.wrapping_sub(amount));
                 Ok(())
             }
         }
diff --git a/capnp/src/serialize.rs b/capnp/src/serialize.rs
@@ -257,10 +257,12 @@ fn read_segment_table<R>(read: &mut R,
     // Don't accept a message which the receiver couldn't possibly traverse without hitting the
     // traversal limit. Without this check, a malicious client could transmit a very large segment
     // size to make the receiver allocate excessive space and possibly crash.
-    if segment_lengths_builder.total_words() as u64 > options.traversal_limit_in_words  {
-        return Err(Error::failed(
-            format!("Message has {} words, which is too large. To increase the limit on the \
-             receiving end, see capnp::message::ReaderOptions.", segment_lengths_builder.total_words())))
+    if let Some(limit) = options.traversal_limit_in_words {
+        if segment_lengths_builder.total_words() > limit {
+            return Err(Error::failed(
+                format!("Message has {} words, which is too large. To increase the limit on the \
+                         receiving end, see capnp::message::ReaderOptions.", segment_lengths_builder.total_words())))
+        }
     }
 
     Ok(Some(segment_lengths_builder))
diff --git a/capnpc/test/test.rs b/capnpc/test/test.rs
@@ -1250,7 +1250,7 @@ mod tests {
 
         let segments = message.get_segments_for_output();
         let reader = message::Reader::new(message::SegmentArray::new(&segments),
-                                          *ReaderOptions::new().traversal_limit_in_words(2));
+                                          *ReaderOptions::new().traversal_limit_in_words(Some(2)));
         match reader.get_root::<test_all_types::Reader>() {
             Err(e) => assert_eq!(e.description, "read limit exceeded"),
             Ok(_) => panic!("expected error"),
