feat(EXC-1847): Support for ic0.subnet_self() (dfinity#3637)

This PR implements a new system API that enables canisters to learn
which subnet they run on. The new functionality is composed of two
functions: `ic0.subnet_self_size()` and `ic0.subnet_self_copy()`,
described in this [Spec
PR](dfinity/portal#3790).

Author: alexandru-uta

rs/embedders/src/wasm_utils/validation.rs

@@ -637,6 +637,26 @@ fn get_valid_system_apis_common(I: ValType) -> HashMap<String, HashMap<String, F
                 },
             )],
         ),
+        (
+            "subnet_self_size",
+            vec![(
+                API_VERSION_IC0,
+                FunctionSignature {
+                    param_types: vec![],
+                    return_type: vec![I],
+                },
+            )],
+        ),
+        (
+            "subnet_self_copy",
+            vec![(
+                API_VERSION_IC0,
+                FunctionSignature {
+                    param_types: vec![I, I, I],
+                    return_type: vec![],
+                },
+            )],
+        ),
     ];
 
     valid_system_apis

rs/embedders/src/wasmtime_embedder/system_api.rs

@@ -459,7 +459,7 @@ pub fn syscalls<
                 charge_for_cpu(&mut caller, overhead::MSG_METHOD_NAME_SIZE)?;
                 with_system_api(&mut caller, |s| s.ic0_msg_method_name_size()).and_then(|s| {
                     I::try_from(s).map_err(|e| {
-                        anyhow::Error::msg(format!("ic0::msg_metohd_name_size failed: {}", e))
+                        anyhow::Error::msg(format!("ic0::msg_method_name_size failed: {}", e))
                     })
                 })
             }
@@ -1037,6 +1037,38 @@ pub fn syscalls<
         })
         .unwrap();
 
+    linker
+        .func_wrap("ic0", "subnet_self_size", {
+            move |mut caller: Caller<'_, StoreData>| {
+                charge_for_cpu(&mut caller, overhead::SUBNET_SELF_SIZE)?;
+                with_system_api(&mut caller, |s| s.ic0_subnet_self_size()).and_then(|s| {
+                    I::try_from(s).map_err(|e| {
+                        anyhow::Error::msg(format!("ic0::subnet_self_size failed: {}", e))
+                    })
+                })
+            }
+        })
+        .unwrap();
+
+    linker
+        .func_wrap("ic0", "subnet_self_copy", {
+            move |mut caller: Caller<'_, StoreData>, dst: I, offset: I, size: I| {
+                let dst: usize = dst.try_into().expect("Failed to convert I to usize");
+                let offset: usize = offset.try_into().expect("Failed to convert I to usize");
+                let size: usize = size.try_into().expect("Failed to convert I to usize");
+                charge_for_cpu_and_mem(&mut caller, overhead::SUBNET_SELF_COPY, size)?;
+                with_memory_and_system_api(&mut caller, |system_api, memory| {
+                    system_api.ic0_subnet_self_copy(dst, offset, size, memory)
+                })?;
+                if feature_flags.write_barrier == FlagStatus::Enabled {
+                    mark_writes_on_bytemap(&mut caller, dst, size)
+                } else {
+                    Ok(())
+                }
+            }
+        })
+        .unwrap();
+
     match main_memory_type {
         WasmMemoryType::Wasm32 => {
             linker

rs/embedders/src/wasmtime_embedder/system_api_complexity.rs

@@ -62,6 +62,8 @@ pub mod overhead {
     pub const MSG_REPLY_DATA_APPEND: NumInstructions = NumInstructions::new(500);
     pub const MSG_REPLY: NumInstructions = NumInstructions::new(500);
     pub const PERFORMANCE_COUNTER: NumInstructions = NumInstructions::new(200);
+    pub const SUBNET_SELF_SIZE: NumInstructions = NumInstructions::new(500);
+    pub const SUBNET_SELF_COPY: NumInstructions = NumInstructions::new(500);
     pub const STABLE_GROW: NumInstructions = NumInstructions::new(500);
     pub const STABLE_READ: NumInstructions = NumInstructions::new(20);
     pub const STABLE_SIZE: NumInstructions = NumInstructions::new(20);

rs/embedders/tests/wasmtime_embedder.rs

@@ -2185,6 +2185,11 @@ fn wasm64_import_system_api_functions() {
       (import "ic0" "is_controller"
         (func $ic0_is_controller (param i64) (param i64) (result i32)))
 
+      (import "ic0" "subnet_self_size"
+        (func $ic0_subnet_self_size (result i64)))
+      (import "ic0" "subnet_self_copy"
+        (func $ic0_subnet_self_copy (param i64) (param i64) (param i64)))
+
         (global $g1 (export "g1") (mut i64) (i64.const 0))
         (func $test (export "canister_update test")
             (i64.store (i64.const 0) (memory.grow (i64.const 1)))
@@ -2692,6 +2697,130 @@ fn wasm64_canister_self_copy() {
     assert_eq!(wasm_heap, expected_heap);
 }
 
+#[test]
+fn wasm64_subnet_self_size() {
+    let wat = r#"
+    (module
+      (import "ic0" "subnet_self_size"
+        (func $ic0_subnet_self_size (result i64)))
+
+      (global $g1 (export "g1") (mut i64) (i64.const 0))
+      (func $test (export "canister_update test")
+        (call $ic0_subnet_self_size)
+        global.set $g1
+      )
+
+      (memory (export "memory") i64 1)
+    )"#;
+
+    let caller = user_test_id(24).get();
+    let payload: Vec<u8> = vec![1, 3, 5, 7];
+    let api = ic_system_api::ApiType::update(
+        UNIX_EPOCH,
+        payload.clone(),
+        Cycles::zero(),
+        caller,
+        call_context_test_id(13),
+    );
+
+    let mut config = ic_config::embedders::Config::default();
+    config.feature_flags.wasm64 = FlagStatus::Enabled;
+    let mut instance = WasmtimeInstanceBuilder::new()
+        .with_config(config)
+        .with_wat(wat)
+        .with_api_type(api)
+        .build();
+
+    let res = instance
+        .run(FuncRef::Method(WasmMethod::Update("test".to_string())))
+        .unwrap();
+
+    assert_eq!(
+        res.exported_globals[0],
+        Global::I64(
+            instance
+                .store_data()
+                .system_api()
+                .unwrap()
+                .ic0_subnet_self_size()
+                .unwrap() as i64
+        )
+    );
+}
+
+#[test]
+fn wasm64_subnet_self_copy() {
+    let wat = r#"
+    (module
+      (import "ic0" "subnet_self_copy"
+        (func $ic0_subnet_self_copy (param i64) (param i64) (param i64)))
+      (import "ic0" "subnet_self_size"
+        (func $ic0_subnet_self_size (result i64)))
+
+      (global $g1 (export "g1") (mut i64) (i64.const 0))
+      (func $test (export "canister_update test")
+        (call $ic0_subnet_self_size)
+        global.set $g1
+        (call $ic0_subnet_self_copy (i64.const 0) (i64.const 0) (call $ic0_subnet_self_size))
+      )
+
+      (memory (export "memory") i64 1)
+    )"#;
+
+    let caller = user_test_id(24).get();
+    let payload: Vec<u8> = vec![1, 3, 5, 7];
+    let api = ic_system_api::ApiType::update(
+        UNIX_EPOCH,
+        payload.clone(),
+        Cycles::zero(),
+        caller,
+        call_context_test_id(13),
+    );
+
+    let mut config = ic_config::embedders::Config::default();
+    config.feature_flags.wasm64 = FlagStatus::Enabled;
+    let mut instance = WasmtimeInstanceBuilder::new()
+        .with_config(config)
+        .with_wat(wat)
+        .with_api_type(api)
+        .build();
+
+    let res = instance
+        .run(FuncRef::Method(WasmMethod::Update("test".to_string())))
+        .unwrap();
+
+    // Only first os page should have been touched.
+    assert_eq!(res.wasm_dirty_pages, vec![ic_sys::PageIndex::new(0)]);
+
+    // Actual heap is larger, but we can only access first os page, the rest is protected.
+    let dirty_heap_size = ic_sys::PAGE_SIZE;
+
+    let wasm_heap: &[u8] = unsafe {
+        let addr = instance.heap_addr(CanisterMemoryType::Heap);
+        let size_in_bytes =
+            instance.heap_size(CanisterMemoryType::Heap).get() * WASM_PAGE_SIZE_IN_BYTES;
+        assert!(size_in_bytes >= dirty_heap_size);
+        std::slice::from_raw_parts_mut(addr as *mut _, dirty_heap_size)
+    };
+
+    let mut expected_heap = vec![0; dirty_heap_size];
+    let subnet_id_size = instance
+        .store_data()
+        .system_api()
+        .unwrap()
+        .ic0_subnet_self_size()
+        .unwrap();
+
+    instance
+        .store_data()
+        .system_api()
+        .unwrap()
+        .ic0_subnet_self_copy(0, 0, subnet_id_size, &mut expected_heap)
+        .unwrap();
+
+    assert_eq!(wasm_heap, expected_heap);
+}
+
 #[test]
 fn wasm64_trap() {
     let wat = r#"

rs/execution_environment/src/query_handler/query_cache/tests.rs

@@ -1510,6 +1510,8 @@ fn query_cache_future_proof_test() {
         | SystemApiCallId::MsgReplyDataAppend
         | SystemApiCallId::OutOfInstructions
         | SystemApiCallId::PerformanceCounter
+        | SystemApiCallId::SubnetSelfSize
+        | SystemApiCallId::SubnetSelfCopy
         | SystemApiCallId::Stable64Grow
         | SystemApiCallId::Stable64Read
         | SystemApiCallId::Stable64Size

rs/execution_environment/tests/hypervisor.rs

@@ -2006,6 +2006,69 @@ fn ic0_canister_self_copy_works() {
     assert_eq!(WasmResult::Reply(canister_id.get().into_vec()), result);
 }
 
+#[test]
+fn ic0_subnet_self_size_works() {
+    let mut test = ExecutionTestBuilder::new().build();
+    let wat = r#"
+        (module
+            (import "ic0" "subnet_self_size"
+                (func $subnet_self_size (result i32))
+            )
+            (import "ic0" "msg_reply" (func $msg_reply))
+            (import "ic0" "msg_reply_data_append"
+            (func $msg_reply_data_append (param i32 i32)))
+            (func (export "canister_update test")
+                ;; heap[0] = $subnet_self_size()
+                (i32.store (i32.const 0) (call $subnet_self_size))
+                ;; return heap[0-4]
+                (call $msg_reply_data_append (i32.const 0) (i32.const 4))
+                (call $msg_reply)
+            )
+            (memory 1 1)
+        )"#;
+    let canister_id = test.canister_from_wat(wat).unwrap();
+    let result = test.ingress(canister_id, "test", vec![]).unwrap();
+    assert_eq!(
+        WasmResult::Reply(vec![
+            test.get_own_subnet_id().get().as_slice().len() as u8,
+            0,
+            0,
+            0
+        ]),
+        result
+    );
+}
+
+#[test]
+fn ic0_subnet_self_copy_works() {
+    let mut test = ExecutionTestBuilder::new().build();
+    let wat = r#"
+        (module
+            (import "ic0" "subnet_self_copy"
+                (func $subnet_self_copy (param i32 i32 i32))
+            )
+            (import "ic0" "msg_reply" (func $msg_reply))
+            (import "ic0" "msg_reply_data_append"
+            (func $msg_reply_data_append (param i32 i32)))
+            (func (export "canister_update test")
+                ;; heap[0..4] = subnet_id_bytes[0..4]
+                (call $subnet_self_copy (i32.const 0) (i32.const 0) (i32.const 4))
+                ;; heap[4..10] = subnet_id_bytes[4..8]
+                (call $subnet_self_copy (i32.const 4) (i32.const 4) (i32.const 6))
+                ;; return heap[0..10]
+                (call $msg_reply_data_append (i32.const 0) (i32.const 10))
+                (call $msg_reply)
+            )
+            (memory 1 1)
+        )"#;
+    let canister_id = test.canister_from_wat(wat).unwrap();
+    let result = test.ingress(canister_id, "test", vec![]).unwrap();
+    assert_eq!(
+        WasmResult::Reply(test.get_own_subnet_id().get().into_vec()),
+        result
+    );
+}
+
 #[test]
 fn ic0_call_has_no_effect_on_trap() {
     let mut test = ExecutionTestBuilder::new().build();

rs/interfaces/src/execution_environment.rs

@@ -219,6 +219,10 @@ pub enum SystemApiCallId {
     OutOfInstructions,
     /// Tracker for `ic0.performance_counter()`
     PerformanceCounter,
+    /// Tracker for `ic0.subnet_self_size()`
+    SubnetSelfSize,
+    /// Tracker for `ic0.subnet_self_copy()`
+    SubnetSelfCopy,
     /// Tracker for `ic0.stable64_grow()`
     Stable64Grow,
     /// Tracker for `ic0.stable64_read()`
@@ -1172,6 +1176,19 @@ pub trait SystemApi {
         dst: usize,
         heap: &mut [u8],
     ) -> HypervisorResult<()>;
+
+    /// Used to look up the size of the subnet Id of the calling canister.
+    fn ic0_subnet_self_size(&self) -> HypervisorResult<usize>;
+
+    /// Used to copy the subnet Id of the calling canister to its heap
+    /// at the location specified by `dst` and `offset`.
+    fn ic0_subnet_self_copy(
+        &self,
+        dst: usize,
+        offset: usize,
+        size: usize,
+        heap: &mut [u8],
+    ) -> HypervisorResult<()>;
 }
 
 #[derive(Copy, Clone, Eq, PartialEq, Debug)]