Crypto: Add overflow checking addition functions.

Added functions to cryptoutil.rs that perform an addition after shifting
the 2nd parameter by a specified constant. These function fail!() if integer
overflow will result. Updated the Sha2 implementation to use these functions.

2 files changed, 71 insertions, 55 deletions

diff --git a/src/libextra/crypto/cryptoutil.rs b/src/libextra/crypto/cryptoutil.rs
index 33ab6a9bcb9..43e3b5c89af 100644
--- a/src/libextra/crypto/cryptoutil.rs
+++ b/src/libextra/crypto/cryptoutil.rs
@@ -8,6 +8,7 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
+use std::num::One;
 use std::vec::bytes::{MutableByteVector, copy_memory};
 
 
@@ -68,6 +69,56 @@ pub fn read_u32v_be(dst: &mut[u32], input: &[u8]) {
 }
 
 
+/// Returns true if adding the two parameters will result in integer overflow
+pub fn will_add_overflow<T: Int + Unsigned>(x: T, y: T) -> bool {
+    // This doesn't handle negative values! Don't copy this code elsewhere without considering if
+    // negative values are important to you!
+    let max: T = Bounded::max_value();
+    return x > max - y;
+}
+
+/// Shifts the second parameter and then adds it to the first. fails!() if there would be unsigned
+/// integer overflow.
+pub fn shift_add_check_overflow<T: Int + Unsigned + Clone>(x: T, mut y: T, shift: T) -> T {
+    if y.leading_zeros() < shift {
+        fail!("Could not add values - integer overflow.");
+    }
+    y = y << shift;
+
+    if will_add_overflow(x.clone(), y.clone()) {
+        fail!("Could not add values - integer overflow.");
+    }
+
+    return x + y;
+}
+
+/// Shifts the second parameter and then adds it to the first, which is a tuple where the first
+/// element is the high order value. fails!() if there would be unsigned integer overflow.
+pub fn shift_add_check_overflow_tuple
+        <T: Int + Unsigned + Clone>
+        (x: (T, T), mut y: T, shift: T) -> (T, T) {
+    if y.leading_zeros() < shift {
+        fail!("Could not add values - integer overflow.");
+    }
+    y = y << shift;
+
+    match x {
+        (hi, low) => {
+            let one: T = One::one();
+            if will_add_overflow(low.clone(), y.clone()) {
+                if will_add_overflow(hi.clone(), one.clone()) {
+                    fail!("Could not add values - integer overflow.");
+                } else {
+                    return (hi + one, low + y);
+                }
+            } else {
+                return (hi, low + y);
+            }
+        }
+    }
+}
+
+
 /// A FixedBuffer, likes its name implies, is a fixed size buffer. When the buffer becomes full, it
 /// must be processed. The input() method takes care of processing and then clearing the buffer
 /// automatically. However, other methods do not and require the caller to process the buffer. Any
diff --git a/src/libextra/crypto/sha2.rs b/src/libextra/crypto/sha2.rs
index b91a54efc92..d92a4be43c3 100644
--- a/src/libextra/crypto/sha2.rs
+++ b/src/libextra/crypto/sha2.rs
@@ -10,8 +10,8 @@
 
 use std::uint;
 
-use cryptoutil::{write_u64_be, write_u32_be, read_u64v_be, read_u32v_be, FixedBuffer,
-    FixedBuffer128, FixedBuffer64, StandardPadding};
+use cryptoutil::{write_u64_be, write_u32_be, read_u64v_be, read_u32v_be, shift_add_check_overflow,
+    shift_add_check_overflow_tuple, FixedBuffer, FixedBuffer128, FixedBuffer64, StandardPadding};
 use digest::Digest;
 
 
@@ -34,47 +34,6 @@ macro_rules! sha2_round(
 )
 
 
-// BitCounter is a specialized structure intended simply for counting the
-// number of bits that have been processed by the SHA-2 512 family of functions.
-// It does very little overflow checking since such checking is not necessary
-// for how it is used. A more generic structure would have to do this checking.
-// So, don't copy this structure and use it elsewhere!
-struct BitCounter {
-    high_bit_count: u64,
-    low_byte_count: u64
-}
-
-impl BitCounter {
-    fn new() -> BitCounter {
-        return BitCounter {
-            high_bit_count: 0,
-            low_byte_count: 0
-        };
-    }
-
-    fn add_bytes(&mut self, bytes: uint) {
-        self.low_byte_count += bytes as u64;
-        if(self.low_byte_count > 0x1fffffffffffffffu64) {
-            self.high_bit_count += (self.low_byte_count >> 61);
-            self.low_byte_count &= 0x1fffffffffffffffu64;
-        }
-    }
-
-    fn reset(&mut self) {
-        self.low_byte_count = 0;
-        self.high_bit_count = 0;
-    }
-
-    fn get_low_bit_count(&self) -> u64 {
-        self.low_byte_count << 3
-    }
-
-    fn get_high_bit_count(&self) -> u64 {
-        self.high_bit_count
-    }
-}
-
-
 // A structure that represents that state of a digest computation for the SHA-2 512 family of digest
 // functions
 struct Engine512State {
@@ -223,7 +182,7 @@ static K64: [u64, ..80] = [
 // A structure that keeps track of the state of the Sha-512 operation and contains the logic
 // necessary to perform the final calculations.
 struct Engine512 {
-    bit_counter: BitCounter,
+    length_bits: (u64, u64),
     buffer: FixedBuffer128,
     state: Engine512State,
     finished: bool,
@@ -232,7 +191,7 @@ struct Engine512 {
 impl Engine512 {
     fn new(h: &[u64, ..8]) -> Engine512 {
         return Engine512 {
-            bit_counter: BitCounter::new(),
+            length_bits: (0, 0),
             buffer: FixedBuffer128::new(),
             state: Engine512State::new(h),
             finished: false
@@ -240,7 +199,7 @@ impl Engine512 {
     }
 
     fn reset(&mut self, h: &[u64, ..8]) {
-        self.bit_counter.reset();
+        self.length_bits = (0, 0);
         self.buffer.reset();
         self.state.reset(h);
         self.finished = false;
@@ -248,7 +207,8 @@ impl Engine512 {
 
     fn input(&mut self, input: &[u8]) {
         assert!(!self.finished)
-        self.bit_counter.add_bytes(input.len());
+        // Assumes that input.len() can be converted to u64 without overflow
+        self.length_bits = shift_add_check_overflow_tuple(self.length_bits, input.len() as u64, 3);
         self.buffer.input(input, |input: &[u8]| { self.state.process_block(input) });
     }
 
@@ -258,8 +218,12 @@ impl Engine512 {
         }
 
         self.buffer.standard_padding(16, |input: &[u8]| { self.state.process_block(input) });
-        write_u64_be(self.buffer.next(8), self.bit_counter.get_high_bit_count());
-        write_u64_be(self.buffer.next(8), self.bit_counter.get_low_bit_count());
+        match self.length_bits {
+            (hi, low) => {
+                write_u64_be(self.buffer.next(8), hi);
+                write_u64_be(self.buffer.next(8), low);
+            }
+        }
         self.state.process_block(self.buffer.full_buffer());
 
         self.finished = true;
@@ -608,7 +572,7 @@ static K32: [u32, ..64] = [
 // A structure that keeps track of the state of the Sha-256 operation and contains the logic
 // necessary to perform the final calculations.
 struct Engine256 {
-    length: u64,
+    length_bits: u64,
     buffer: FixedBuffer64,
     state: Engine256State,
     finished: bool,
@@ -617,7 +581,7 @@ struct Engine256 {
 impl Engine256 {
     fn new(h: &[u32, ..8]) -> Engine256 {
         return Engine256 {
-            length: 0,
+            length_bits: 0,
             buffer: FixedBuffer64::new(),
             state: Engine256State::new(h),
             finished: false
@@ -625,7 +589,7 @@ impl Engine256 {
     }
 
     fn reset(&mut self, h: &[u32, ..8]) {
-        self.length = 0;
+        self.length_bits = 0;
         self.buffer.reset();
         self.state.reset(h);
         self.finished = false;
@@ -633,7 +597,8 @@ impl Engine256 {
 
     fn input(&mut self, input: &[u8]) {
         assert!(!self.finished)
-        self.length += input.len() as u64;
+        // Assumes that input.len() can be converted to u64 without overflow
+        self.length_bits = shift_add_check_overflow(self.length_bits, input.len() as u64, 3);
         self.buffer.input(input, |input: &[u8]| { self.state.process_block(input) });
     }
 
@@ -643,8 +608,8 @@ impl Engine256 {
         }
 
         self.buffer.standard_padding(8, |input: &[u8]| { self.state.process_block(input) });
-        write_u32_be(self.buffer.next(4), (self.length >> 29) as u32 );
-        write_u32_be(self.buffer.next(4), (self.length << 3) as u32);
+        write_u32_be(self.buffer.next(4), (self.length_bits >> 32) as u32 );
+        write_u32_be(self.buffer.next(4), self.length_bits as u32);
         self.state.process_block(self.buffer.full_buffer());
 
         self.finished = true;