Merge pull request #517 from enigmampc/query-recursion-limit

Added recursion limit of 5 to queries

Author: reuvenpo

cosmwasm/packages/enclave-ffi-types/src/types.rs

@@ -154,6 +154,8 @@ pub enum EnclaveError {
     Panic,
     #[display(fmt = "enclave ran out of heap memory")]
     OutOfMemory,
+    #[display(fmt = "depth of nested contract calls exceeded")]
+    ExceededRecursionLimit,
     /// Unexpected Error happened, no more details available
     #[display(fmt = "unknown error")]
     Unknown,

cosmwasm/packages/std/src/errors/system_error.rs

@@ -21,6 +21,7 @@ pub enum SystemError {
     NoSuchContract { addr: HumanAddr },
     Unknown {},
     UnsupportedRequest { kind: String },
+    ExceededRecursionLimit {},
 }
 
 impl std::error::Error for SystemError {}
@@ -45,6 +46,7 @@ impl std::fmt::Display for SystemError {
             SystemError::UnsupportedRequest { kind } => {
                 write!(f, "Unsupported query type: {}", kind)
             }
+            SystemError::ExceededRecursionLimit {} => write!(f, "Query recursion limit exceeded"),
         }
     }
 }

cosmwasm/packages/wasmi-runtime/Enclave.config.prod.xml

@@ -3,7 +3,7 @@
   <ProdID>0</ProdID>
   <ISVSVN>0</ISVSVN>
   <StackMaxSize>0x800000</StackMaxSize>
-  <HeapMaxSize>0x8000000</HeapMaxSize>
+  <HeapMaxSize>0x10000000</HeapMaxSize>
   <TCSNum>1</TCSNum>
   <TCSPolicy>1</TCSPolicy>
   <DisableDebug>1</DisableDebug>

cosmwasm/packages/wasmi-runtime/Enclave.config.xml

@@ -3,7 +3,7 @@
   <ProdID>0</ProdID>
   <ISVSVN>0</ISVSVN>
   <StackMaxSize>0x800000</StackMaxSize>
-  <HeapMaxSize>0x4000000</HeapMaxSize>
+  <HeapMaxSize>0x10000000</HeapMaxSize>
   <TCSNum>1</TCSNum>
   <TCSPolicy>1</TCSPolicy>
   <DisableDebug>0</DisableDebug>

cosmwasm/packages/wasmi-runtime/src/cosmwasm/system_error.rs

@@ -23,6 +23,7 @@ pub enum SystemError {
     NoSuchContract { addr: HumanAddr },
     Unknown {},
     UnsupportedRequest { kind: String },
+    ExceededRecursionLimit {},
 }
 
 pub type SystemResult<T> = Result<T, SystemError>;

cosmwasm/packages/wasmi-runtime/src/exports.rs

@@ -13,7 +13,7 @@ use crate::results::{
     result_query_success_to_queryresult,
 };
 use crate::{
-    oom_handler,
+    oom_handler, recursion_depth,
     utils::{validate_const_ptr, validate_mut_ptr},
 };
 
@@ -116,6 +116,19 @@ pub unsafe extern "C" fn ecall_init(
     sig_info: *const u8,
     sig_info_len: usize,
 ) -> InitResult {
+    let _recursion_guard = match recursion_depth::guard() {
+        Ok(rg) => rg,
+        Err(err) => {
+            // https://github.com/enigmampc/SecretNetwork/pull/517#discussion_r481924571
+            // I believe that this error condition is currently unreachable.
+            // I think we can safely remove it completely right now, and have
+            // recursion_depth::increment() simply increment the counter with no further checks,
+            // but i wanted to stay on the safe side here, in case something changes in the
+            // future, and we can easily spot that we forgot to add a limit somewhere.
+            error!("recursion limit exceeded, can not perform init!");
+            return InitResult::Failure { err };
+        }
+    };
     if let Err(err) = oom_handler::register_oom_handler() {
         error!("Could not register OOM handler!");
         return InitResult::Failure { err };
@@ -200,6 +213,19 @@ pub unsafe extern "C" fn ecall_handle(
     sig_info: *const u8,
     sig_info_len: usize,
 ) -> HandleResult {
+    let _recursion_guard = match recursion_depth::guard() {
+        Ok(rg) => rg,
+        Err(err) => {
+            // https://github.com/enigmampc/SecretNetwork/pull/517#discussion_r481924571
+            // I believe that this error condition is currently unreachable.
+            // I think we can safely remove it completely right now, and have
+            // recursion_depth::increment() simply increment the counter with no further checks,
+            // but i wanted to stay on the safe side here, in case something changes in the
+            // future, and we can easily spot that we forgot to add a limit somewhere.
+            error!("recursion limit exceeded, can not perform handle!");
+            return HandleResult::Failure { err };
+        }
+    };
     if let Err(err) = oom_handler::register_oom_handler() {
         error!("Could not register OOM handler!");
         return HandleResult::Failure { err };
@@ -280,6 +306,19 @@ pub unsafe extern "C" fn ecall_query(
     msg: *const u8,
     msg_len: usize,
 ) -> QueryResult {
+    let _recursion_guard = match recursion_depth::guard() {
+        Ok(rg) => rg,
+        Err(err) => {
+            // https://github.com/enigmampc/SecretNetwork/pull/517#discussion_r481924571
+            // I believe that this error condition is currently unreachable.
+            // I think we can safely remove it completely right now, and have
+            // recursion_depth::increment() simply increment the counter with no further checks,
+            // but i wanted to stay on the safe side here, in case something changes in the
+            // future, and we can easily spot that we forgot to add a limit somewhere.
+            error!("recursion limit exceeded, can not perform query!");
+            return QueryResult::Failure { err };
+        }
+    };
     if let Err(err) = oom_handler::register_oom_handler() {
         error!("Could not register OOM handler!");
         return QueryResult::Failure { err };

cosmwasm/packages/wasmi-runtime/src/lib.rs

@@ -21,6 +21,7 @@ pub mod exports;
 pub mod imports;
 pub mod logger;
 mod oom_handler;
+mod recursion_depth;
 pub mod registration;
 use std::env;
 

cosmwasm/packages/wasmi-runtime/src/oom_handler.rs

@@ -72,12 +72,13 @@ impl SafetyBuffer {
 }
 
 lazy_static! {
-    /// SAFETY_BUFFER is a 32 MiB of SafetyBuffer. This is occupying 50% of available memory
-    /// to be extra sure this is enough.
+    /// SAFETY_BUFFER is a 4 MiB of SafetyBuffer. This is twice the bare minimum to unwind after
+    /// a best-case OOM event. thanks to the recursion limit on queries, together with other memory
+    /// limits, we don't expect to hit OOM, and this mechanism remains in place just in case.
     /// 2 MiB is the minimum allowed buffer. If we don't succeed to allocate 2 MiB, we throw a panic,
-    /// if we do succeed to allocate 2 MiB but less than 32 MiB than we move on and will try to allocate
+    /// if we do succeed to allocate 2 MiB but less than 4 MiB than we move on and will try to allocate
     /// the rest on the next entry to the enclave.
-    static ref SAFETY_BUFFER: SgxMutex<SafetyBuffer> = SgxMutex::new(SafetyBuffer::new(16 * 1024, 2 * 1024));
+    static ref SAFETY_BUFFER: SgxMutex<SafetyBuffer> = SgxMutex::new(SafetyBuffer::new(4 * 1024, 2 * 1024));
 }
 
 static OOM_HAPPENED: AtomicBool = AtomicBool::new(false);

cosmwasm/packages/wasmi-runtime/src/recursion_depth.rs

@@ -0,0 +1,55 @@
+use std::sync::SgxMutex;
+
+use lazy_static::lazy_static;
+
+use enclave_ffi_types::EnclaveError;
+
+const RECURSION_LIMIT: u8 = 5;
+
+lazy_static! {
+    /// This counter tracks the recursion depth of queries,
+    /// and effectively the amount of loaded instances of WASMI.
+    ///
+    /// It is incremented before each computation begins and is decremented after each computation ends.
+    static ref RECURSION_DEPTH: SgxMutex<u8> = SgxMutex::new(0);
+}
+
+fn increment() -> Result<(), EnclaveError> {
+    let mut depth = RECURSION_DEPTH.lock().unwrap();
+    if *depth == RECURSION_LIMIT {
+        return Err(EnclaveError::ExceededRecursionLimit);
+    }
+    *depth = depth.saturating_add(1);
+    Ok(())
+}
+
+fn decrement() {
+    let mut depth = RECURSION_DEPTH.lock().unwrap();
+    *depth = depth.saturating_sub(1);
+}
+
+/// Returns whether or not this is the last possible level of recursion
+pub fn limit_reached() -> bool {
+    *RECURSION_DEPTH.lock().unwrap() == RECURSION_LIMIT
+}
+
+pub struct RecursionGuard {
+    _private: (), // prevent direct instantiation outside this module
+}
+
+impl RecursionGuard {
+    pub fn new() -> Result<Self, EnclaveError> {
+        increment()?;
+        Ok(Self { _private: () })
+    }
+}
+
+impl Drop for RecursionGuard {
+    fn drop(&mut self) {
+        decrement();
+    }
+}
+
+pub fn guard() -> Result<RecursionGuard, EnclaveError> {
+    RecursionGuard::new()
+}

cosmwasm/packages/wasmi-runtime/src/wasm/query_chain.rs

@@ -1,11 +1,12 @@
 use super::errors::WasmEngineError;
 use crate::crypto::Ed25519PublicKey;
+use crate::recursion_depth;
 use crate::wasm::types::{IoNonce, SecretMessage};
 use crate::{exports, imports};
 
-use crate::cosmwasm::encoding::Binary;
-use crate::cosmwasm::query::{QueryRequest, WasmQuery};
 use crate::cosmwasm::{
+    encoding::Binary,
+    query::{QueryRequest, WasmQuery},
     std_error::{StdError, StdResult},
     system_error::{SystemError, SystemResult},
 };
@@ -22,6 +23,10 @@ pub fn encrypt_and_query_chain(
     gas_used: &mut u64,
     gas_limit: u64,
 ) -> Result<Vec<u8>, WasmEngineError> {
+    if let Some(answer) = check_recursion_limit() {
+        return serialize_error_response(&answer);
+    }
+
     let mut query_struct: QueryRequest = match serde_json::from_slice(query) {
         Ok(query_struct) => query_struct,
         Err(err) => {
@@ -182,6 +187,21 @@ fn query_chain(
     (Ok(value), gas_used)
 }
 
+/// Check whether the query is allowed to run.
+///
+/// We make sure that a recursion limit is in place in order to
+/// mitigate cases where the enclave runs out of memory.
+fn check_recursion_limit() -> Option<SystemResult<StdResult<Binary>>> {
+    if recursion_depth::limit_reached() {
+        debug!(
+            "Recursion limit reached while performing nested queries. Returning error to contract."
+        );
+        Some(Err(SystemError::ExceededRecursionLimit {}))
+    } else {
+        None
+    }
+}
+
 fn system_error_invalid_request<T>(request: &[u8], err: T) -> Result<Vec<u8>, WasmEngineError>
 where
     T: std::fmt::Debug + ToString,
@@ -196,16 +216,7 @@ where
         error: err.to_string(),
     });
 
-    serde_json::to_vec(&answer).map_err(|err| {
-        // this should never happen
-        debug!(
-            "encrypt_and_query_chain() got an error while trying to serialize the error {:?} returned to WASM: {:?}",
-            answer,
-            err
-        );
-
-        WasmEngineError::SerializationError
-    })
+    serialize_error_response(&answer)
 }
 
 fn system_error_invalid_response<T>(response: Vec<u8>, err: T) -> Result<Vec<u8>, WasmEngineError>
@@ -217,7 +228,13 @@ where
         error: err.to_string(),
     });
 
-    serde_json::to_vec(&answer).map_err(|err| {
+    serialize_error_response(&answer)
+}
+
+fn serialize_error_response(
+    answer: &SystemResult<StdResult<Binary>>,
+) -> Result<Vec<u8>, WasmEngineError> {
+    serde_json::to_vec(answer).map_err(|err| {
         // this should never happen
         debug!(
             "encrypt_and_query_chain() got an error while trying to serialize the error {:?} returned to WASM: {:?}",

go-cosmwasm/types/systemerror.go

@@ -12,6 +12,7 @@ type SystemError struct {
 	NoSuchContract     *NoSuchContract     `json:"no_such_contract,omitempty"`
 	Unknown            *Unknown            `json:"unknown,omitempty"`
 	UnsupportedRequest *UnsupportedRequest `json:"unsupported_request,omitempty"`
+	ExceededRecursionLimit	*ExceededRecursionLimit	`json:"exceeded_recursion_limit,omitempty"`
 }
 
 var (
@@ -21,6 +22,7 @@ var (
 	_ error = NoSuchContract{}
 	_ error = Unknown{}
 	_ error = UnsupportedRequest{}
+	_ error = ExceededRecursionLimit{}
 )
 
 func (a SystemError) Error() string {
@@ -35,6 +37,8 @@ func (a SystemError) Error() string {
 		return a.Unknown.Error()
 	case a.UnsupportedRequest != nil:
 		return a.UnsupportedRequest.Error()
+	case a.ExceededRecursionLimit != nil:
+		return a.ExceededRecursionLimit.Error()
 	default:
 		panic("unknown error variant")
 	}
@@ -80,6 +84,12 @@ func (e UnsupportedRequest) Error() string {
 	return fmt.Sprintf("unsupported request: %s", e.Kind)
 }
 
+type ExceededRecursionLimit struct{}
+
+func (e ExceededRecursionLimit) Error() string {
+	return "unknown system error"
+}
+
 // ToSystemError will try to convert the given error to an SystemError.
 // This is important to returning any Go error back to Rust.
 //
@@ -118,6 +128,10 @@ func ToSystemError(err error) *SystemError {
 		return &SystemError{UnsupportedRequest: &t}
 	case *UnsupportedRequest:
 		return &SystemError{UnsupportedRequest: t}
+	case ExceededRecursionLimit:
+		return &SystemError{ExceededRecursionLimit: &t}
+	case *ExceededRecursionLimit:
+		return &SystemError{ExceededRecursionLimit: t}
 	default:
 		return nil
 	}

x/compute/internal/keeper/recurse_test.go

@@ -281,6 +281,7 @@ func TestLimitRecursiveQueryGas(t *testing.T) {
 		expectedGas               uint64
 		expectOutOfGas            bool
 		expectOOM                 bool
+		expectRecursionLimit      bool
 	}{
 		"no recursion, lots of work": {
 			gasLimit: 4_000_000,
@@ -291,6 +292,18 @@ func TestLimitRecursiveQueryGas(t *testing.T) {
 			expectQueriesFromContract: 0,
 			expectedGas:               GasWork2k,
 		},
+		"recursion 4, lots of work": {
+			gasLimit: 4_000_000,
+			msg: Recurse{
+				Depth: 4,
+				Work:  2000,
+			},
+			expectQueriesFromContract: 4,
+			expectedGas:               GasWork2k + 9*(GasWork2k+GasReturnHashed),
+			expectOutOfGas:            false,
+			expectOOM:                 false,
+			expectRecursionLimit:      false,
+		},
 		"recursion 9, lots of work": {
 			gasLimit: 4_000_000,
 			msg: Recurse{
@@ -300,7 +313,8 @@ func TestLimitRecursiveQueryGas(t *testing.T) {
 			expectQueriesFromContract: 9,
 			expectedGas:               GasWork2k + 9*(GasWork2k+GasReturnHashed),
 			expectOutOfGas:            false,
-			expectOOM:                 true,
+			expectOOM:                 false,
+			expectRecursionLimit:       true,
 		},
 		// this is where we expect an error...
 		// it has enough gas to run 4 times and die on the 5th (4th time dispatching to sub-contract)
@@ -314,7 +328,8 @@ func TestLimitRecursiveQueryGas(t *testing.T) {
 			},
 			expectQueriesFromContract: 4,
 			expectOutOfGas:            false,
-			expectOOM:                 true,
+			expectOOM:                 false,
+			expectRecursionLimit:       true,
 		},
 	}
 
@@ -354,6 +369,14 @@ func TestLimitRecursiveQueryGas(t *testing.T) {
 				return
 			}
 
+			if tc.expectRecursionLimit {
+				_, qErr := queryHelper(t, keeper, ctx, contractAddr, string(msg), true, tc.gasLimit)
+				require.NotEmpty(t, qErr)
+				require.NotNil(t, qErr.GenericErr)
+				require.Contains(t, qErr.GenericErr.Msg, "Querier system error: Query recursion limit exceeded")
+				return
+			}
+
 			// otherwise, we expect a successful call
 			_, qErr := queryHelper(t, keeper, ctx, contractAddr, string(msg), true, tc.gasLimit)
 			require.Empty(t, qErr)