rust hostcalls: support hard limits on data sizes

The API doesn't prevent applications from perfoming symmetric operations on data as large as memory allows. Which means that hostcalls may not return quickly enough for some runtimes, which in turn can be an issue if they don't support preemption on hostcalls. A WasiCryptoCtx can now be initialized with limits for individual algorithms. With all symmetric cryptographic operations, there's no significant difference in CPU usage when absorbing and squeezing data, so we can use a single limit per algorithm to keep things simple. No behavior change by default.
WebAssembly · Sep 19, 2022 · 3a374e1 · 3a374e1
1 parent cb0788e
commit 3a374e1
Show file tree

Hide file tree

Showing 8 changed files with 153 additions and 31 deletions.
diff --git a/implementations/hostcalls/rust/src/lib.rs b/implementations/hostcalls/rust/src/lib.rs
@@ -61,8 +61,26 @@ pub enum AlgorithmType {
     KeyExchange,
 }
 
+#[derive(Clone, Debug, PartialEq, Eq, Default)]
+pub struct Limits(HashMap<&'static str, usize>);
+
+impl Limits {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn set(&mut self, algorithm: &'static str, limit: usize) {
+        self.0.insert(algorithm, limit);
+    }
+
+    pub fn get(&self, algorithm: &'static str) -> Option<usize> {
+        self.0.get(algorithm).copied()
+    }
+}
+
 pub struct CryptoCtx {
     pub(crate) handles: HandleManagers,
+    pub(crate) limits: Limits,
 }
 
 impl CryptoCtx {
@@ -82,6 +100,17 @@ impl CryptoCtx {
                 symmetric_tag: HandlesManager::new(0x0a),
                 secrets_manager: HandlesManager::new(0x0b),
             },
+            limits: Limits::default(),
         }
     }
+
+    pub fn set_limits(&mut self, limits: Limits) {
+        self.limits = limits;
+    }
+
+    pub fn new_with_limits(limits: Limits) -> Self {
+        let mut ctx = Self::new();
+        ctx.set_limits(limits);
+        ctx
+    }
 }
diff --git a/implementations/hostcalls/rust/src/symmetric/aes_gcm.rs b/implementations/hostcalls/rust/src/symmetric/aes_gcm.rs
@@ -23,6 +23,7 @@ enum AesGcmVariant {
 pub struct AesGcmSymmetricState {
     pub alg: SymmetricAlgorithm,
     options: SymmetricOptions,
+    size_limit: Option<usize>,
     #[derivative(Debug = "ignore")]
     ctx: AesGcmVariant,
     ad: Vec<u8>,
@@ -107,6 +108,7 @@ impl AesGcmSymmetricState {
         alg: SymmetricAlgorithm,
         key: Option<SymmetricKey>,
         options: Option<SymmetricOptions>,
+        size_limit: Option<usize>,
     ) -> Result<Self, CryptoError> {
         let key = key.ok_or(CryptoError::KeyRequired)?;
         let key = key.inner();
@@ -132,6 +134,7 @@ impl AesGcmSymmetricState {
         let state = AesGcmSymmetricState {
             alg,
             options: options.clone(),
+            size_limit,
             ctx: aes_gcm_impl,
             ad: vec![],
             nonce: Some(nonce),
@@ -153,7 +156,11 @@ impl SymmetricStateLike for AesGcmSymmetricState {
         self.options.get_u64(name)
     }
 
-    fn absorb(&mut self, data: &[u8]) -> Result<(), CryptoError> {
+    fn size_limit(&self) -> Option<usize> {
+        self.size_limit
+    }
+
+    fn absorb_unchecked(&mut self, data: &[u8]) -> Result<(), CryptoError> {
         self.ad.extend_from_slice(data);
         Ok(())
     }

diff --git a/implementations/hostcalls/rust/src/symmetric/chacha_poly.rs b/implementations/hostcalls/rust/src/symmetric/chacha_poly.rs
@@ -22,6 +22,7 @@ enum ChaChaPolyVariant {
 pub struct ChaChaPolySymmetricState {
     pub alg: SymmetricAlgorithm,
     options: SymmetricOptions,
+    size_limit: Option<usize>,
     #[derivative(Debug = "ignore")]
     ctx: ChaChaPolyVariant,
     ad: Vec<u8>,
@@ -105,6 +106,7 @@ impl ChaChaPolySymmetricState {
         alg: SymmetricAlgorithm,
         key: Option<SymmetricKey>,
         options: Option<SymmetricOptions>,
+        size_limits: Option<usize>,
     ) -> Result<Self, CryptoError> {
         let key = key.ok_or(CryptoError::KeyRequired)?;
         let key = key.inner();
@@ -143,6 +145,7 @@ impl ChaChaPolySymmetricState {
         let state = ChaChaPolySymmetricState {
             alg,
             options: options.clone(),
+            size_limit: size_limits,
             ctx: chapoly_impl,
             ad: vec![],
             nonce: Some(nonce),
@@ -164,7 +167,11 @@ impl SymmetricStateLike for ChaChaPolySymmetricState {
         self.options.get_u64(name)
     }
 
-    fn absorb(&mut self, data: &[u8]) -> Result<(), CryptoError> {
+    fn size_limit(&self) -> Option<usize> {
+        self.size_limit
+    }
+
+    fn absorb_unchecked(&mut self, data: &[u8]) -> Result<(), CryptoError> {
         self.ad.extend_from_slice(data);
         Ok(())
     }

diff --git a/implementations/hostcalls/rust/src/symmetric/hkdf.rs b/implementations/hostcalls/rust/src/symmetric/hkdf.rs
@@ -11,6 +11,7 @@ use crate::rand::SecureRandom;
 pub struct HkdfSymmetricState {
     pub alg: SymmetricAlgorithm,
     options: Option<SymmetricOptions>,
+    size_limit: Option<usize>,
     key: Vec<u8>,
     data: Vec<u8>,
 }
@@ -99,6 +100,7 @@ impl HkdfSymmetricState {
         alg: SymmetricAlgorithm,
         key: Option<SymmetricKey>,
         options: Option<SymmetricOptions>,
+        size_limit: Option<usize>,
     ) -> Result<Self, CryptoError> {
         let key = key.ok_or(CryptoError::KeyRequired)?;
         let key = key.inner();
@@ -110,6 +112,7 @@ impl HkdfSymmetricState {
         Ok(HkdfSymmetricState {
             alg,
             options,
+            size_limit,
             key,
             data: vec![],
         })
@@ -135,7 +138,11 @@ impl SymmetricStateLike for HkdfSymmetricState {
             .get_u64(name)
     }
 
-    fn absorb(&mut self, data: &[u8]) -> Result<(), CryptoError> {
+    fn size_limit(&self) -> Option<usize> {
+        self.size_limit
+    }
+
+    fn absorb_unchecked(&mut self, data: &[u8]) -> Result<(), CryptoError> {
         self.data.extend_from_slice(data);
         Ok(())
     }
@@ -158,7 +165,7 @@ impl SymmetricStateLike for HkdfSymmetricState {
         builder.import(&raw_prk)
     }
 
-    fn squeeze(&mut self, out: &mut [u8]) -> Result<(), CryptoError> {
+    fn squeeze_unchecked(&mut self, out: &mut [u8]) -> Result<(), CryptoError> {
         match self.alg {
             SymmetricAlgorithm::HkdfSha256Expand => Hkdf::<Sha256>::from_prk(&self.key)
                 .map_err(|_| CryptoError::InvalidKey)?

diff --git a/implementations/hostcalls/rust/src/symmetric/hmac_sha2.rs b/implementations/hostcalls/rust/src/symmetric/hmac_sha2.rs
@@ -19,6 +19,7 @@ enum HmacVariant {
 pub struct HmacSha2SymmetricState {
     pub alg: SymmetricAlgorithm,
     options: Option<SymmetricOptions>,
+    size_limit: Option<usize>,
     ctx: HmacVariant,
 }
 
@@ -100,6 +101,7 @@ impl HmacSha2SymmetricState {
         alg: SymmetricAlgorithm,
         key: Option<SymmetricKey>,
         options: Option<SymmetricOptions>,
+        size_limit: Option<usize>,
     ) -> Result<Self, CryptoError> {
         let key = key.ok_or(CryptoError::KeyRequired)?;
         let key = key.inner();
@@ -118,7 +120,12 @@ impl HmacSha2SymmetricState {
             ),
             _ => bail!(CryptoError::UnsupportedAlgorithm),
         };
-        Ok(HmacSha2SymmetricState { alg, options, ctx })
+        Ok(HmacSha2SymmetricState {
+            alg,
+            options,
+            size_limit,
+            ctx,
+        })
     }
 }
 
@@ -141,7 +148,11 @@ impl SymmetricStateLike for HmacSha2SymmetricState {
             .get_u64(name)
     }
 
-    fn absorb(&mut self, data: &[u8]) -> Result<(), CryptoError> {
+    fn size_limit(&self) -> Option<usize> {
+        self.size_limit
+    }
+
+    fn absorb_unchecked(&mut self, data: &[u8]) -> Result<(), CryptoError> {
         match &mut self.ctx {
             HmacVariant::Sha256(x) => x.update(data),
             HmacVariant::Sha512(x) => x.update(data),

diff --git a/implementations/hostcalls/rust/src/symmetric/sha2.rs b/implementations/hostcalls/rust/src/symmetric/sha2.rs
@@ -16,6 +16,7 @@ enum HashVariant {
 pub struct Sha2SymmetricState {
     pub alg: SymmetricAlgorithm,
     options: Option<SymmetricOptions>,
+    size_limit: Option<usize>,
     ctx: HashVariant,
 }
 
@@ -24,6 +25,7 @@ impl Sha2SymmetricState {
         alg: SymmetricAlgorithm,
         key: Option<&SymmetricKey>,
         options: Option<SymmetricOptions>,
+        size_limit: Option<usize>,
     ) -> Result<Self, CryptoError> {
         if key.is_some() {
             return Err(CryptoError::KeyNotSupported);
@@ -34,7 +36,12 @@ impl Sha2SymmetricState {
             SymmetricAlgorithm::Sha512_256 => HashVariant::Sha512_256(Sha512_256::new()),
             _ => bail!(CryptoError::UnsupportedAlgorithm),
         };
-        Ok(Sha2SymmetricState { alg, options, ctx })
+        Ok(Sha2SymmetricState {
+            alg,
+            options,
+            size_limit,
+            ctx,
+        })
     }
 }
 
@@ -57,7 +64,11 @@ impl SymmetricStateLike for Sha2SymmetricState {
             .get_u64(name)
     }
 
-    fn absorb(&mut self, data: &[u8]) -> Result<(), CryptoError> {
+    fn size_limit(&self) -> Option<usize> {
+        self.size_limit
+    }
+
+    fn absorb_unchecked(&mut self, data: &[u8]) -> Result<(), CryptoError> {
         match &mut self.ctx {
             HashVariant::Sha256(x) => x.update(data),
             HashVariant::Sha512(x) => x.update(data),
@@ -66,7 +77,7 @@ impl SymmetricStateLike for Sha2SymmetricState {
         Ok(())
     }
 
-    fn squeeze(&mut self, out: &mut [u8]) -> Result<(), CryptoError> {
+    fn squeeze_unchecked(&mut self, out: &mut [u8]) -> Result<(), CryptoError> {
         let raw = match &self.ctx {
             HashVariant::Sha256(x) => x.clone().finalize().to_vec(),
             HashVariant::Sha512(x) => x.clone().finalize().to_vec(),