[PVM-host] Fix accumulation pipeline and add arithmetic overflow guards

common/src/types/workloads/work_report.rs

@@ -491,7 +491,10 @@ impl WorkReport {
     }
 
     pub fn total_accumulation_gas_allotted(&self) -> UnsignedGas {
-        self.digests.iter().map(|wd| wd.accumulate_gas_limit).sum()
+        self.digests
+            .iter()
+            .try_fold(0u64, |acc, wd| acc.checked_add(wd.accumulate_gas_limit))
+            .unwrap_or(UnsignedGas::MAX)
     }
 
     pub fn extract_exports_manifest(&self) -> ReportedWorkPackage {

pvm/pvm-core/src/state/memory.rs

@@ -159,7 +159,17 @@ impl Memory {
         if length == 0 {
             return true;
         }
-        let end = start as usize + length - 1;
+        let start_usize = start as usize;
+        let end = match start_usize
+            .checked_add(length)
+            .and_then(|end| end.checked_sub(1))
+        {
+            Some(end) => end,
+            None => return false,
+        };
+        if end >= self.data.len() {
+            return false;
+        }
         let (start_page, _) = self.get_page_and_offset(start);
         let (end_page, _) = self.get_page_and_offset(end as MemAddress);
         self.is_page_range_readable(start_page..end_page + 1)
@@ -205,7 +215,17 @@ impl Memory {
         if length == 0 {
             return true;
         }
-        let end = start as usize + length - 1;
+        let start_usize = start as usize;
+        let end = match start_usize
+            .checked_add(length)
+            .and_then(|end| end.checked_sub(1))
+        {
+            Some(end) => end,
+            None => return false,
+        };
+        if end >= self.data.len() {
+            return false;
+        }
         let (start_page, _) = self.get_page_and_offset(start);
         let (end_page, _) = self.get_page_and_offset(end as MemAddress);
         self.is_page_range_writable(start_page..end_page + 1)

pvm/pvm-host/src/host_functions/accumulate/mod.rs

@@ -78,15 +78,18 @@ impl<S: HostStateProvider> AccumulateHostFunction<S> {
         }
 
         // Read always-accumulate services from the memory
+        let Some(always_accumulate_len) = always_accumulates_count.checked_mul(12) else {
+            host_call_panic!()
+        };
         if !vm
             .memory
-            .is_address_range_readable(always_accumulate_offset, 12 * always_accumulates_count)
+            .is_address_range_readable(always_accumulate_offset, always_accumulate_len)
         {
             host_call_panic!()
         }
         let Ok(always_accumulate_services_data) = vm
             .memory
-            .read_bytes(always_accumulate_offset, 12 * always_accumulates_count)
+            .read_bytes(always_accumulate_offset, always_accumulate_len)
         else {
             host_call_panic!()
         };
@@ -552,7 +555,10 @@ impl<S: HostStateProvider> AccumulateHostFunction<S> {
         // --- Check Sender Balance (Err: CASH)
 
         let amount = vm.read_reg(8);
-        if accumulator_balance.saturating_sub(amount) < accumulator_threshold_balance {
+        let Some(remaining_balance) = accumulator_balance.checked_sub(amount) else {
+            continue_cash!()
+        };
+        if remaining_balance < accumulator_threshold_balance {
             continue_cash!()
         }
 
@@ -1198,7 +1204,10 @@ impl<S: HostStateProvider> AccumulateHostFunction<S> {
         let service_id = if service_id_reg == u64::MAX {
             service_id
         } else {
-            service_id_reg as ServiceId
+            let Ok(service_id) = vm.read_reg_as_service_id(7) else {
+                continue_who!()
+            };
+            service_id
         };
 
         // Service account not found

pvm/pvm-host/src/host_functions/accumulate/tests.rs

@@ -4233,6 +4233,84 @@ mod provide_tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_provide_service_id_out_of_range_returns_who() -> Result<(), Box<dyn Error>> {
+        let fixture = ProvideTestFixture::default();
+        let out_of_range_service = (1u64 << 32) + fixture.provide_service as u64;
+        let vm = fixture
+            .prepare_vm_builder()?
+            .with_reg(7, out_of_range_service)
+            .with_mem_readable_range(fixture.mem_readable_range.clone())?
+            .build();
+        let state_provider = Arc::new(fixture.prepare_state_provider());
+        let mut context = fixture
+            .prepare_invocation_context(state_provider.clone())
+            .await?;
+
+        // Check host-call result
+        let res = AccumulateHostFunction::<MockStateManager>::host_provide(
+            fixture.provide_service,
+            &vm,
+            state_provider.clone(),
+            &mut context,
+        )
+        .await?;
+        assert_eq!(res, ProvideTestFixture::host_call_result_who());
+
+        // Check partial state after host-call (should remain unchanged)
+        let x = context.get_mut_accumulate_x().unwrap();
+        assert!(!x.provided_preimages.contains(&(
+            fixture.provide_service,
+            Octets::from_vec(fixture.preimage_data.clone())
+        )));
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_provide_invalid_service_id_returns_who() -> Result<(), Box<dyn Error>> {
+        let fixture = ProvideTestFixture::default();
+        let invalid_service_id = u64::from(u32::MAX) + 1;
+        let spoofed_service_id = 0;
+        let vm = fixture
+            .prepare_vm_builder()?
+            .with_reg(7, invalid_service_id)
+            .with_mem_readable_range(fixture.mem_readable_range.clone())?
+            .build();
+        let state_provider = Arc::new(
+            MockStateManager::builder()
+                .with_empty_account(fixture.accumulate_host)
+                .with_empty_account(spoofed_service_id)
+                .with_lookups_entry(
+                    spoofed_service_id,
+                    (fixture.preimage_hash.clone(), fixture.preimage_size),
+                    AccountLookupsEntry {
+                        value: AccountLookupsEntryTimeslots::try_from(vec![]).unwrap(),
+                    },
+                ),
+        );
+        let mut context = fixture
+            .prepare_invocation_context(state_provider.clone())
+            .await?;
+
+        // Check host-call result
+        let res = AccumulateHostFunction::<MockStateManager>::host_provide(
+            fixture.accumulate_host,
+            &vm,
+            state_provider.clone(),
+            &mut context,
+        )
+        .await?;
+        assert_eq!(res, ProvideTestFixture::host_call_result_who());
+
+        // Check partial state after host-call (should remain unchanged)
+        let x = context.get_mut_accumulate_x().unwrap();
+        assert!(!x.provided_preimages.contains(&(
+            spoofed_service_id,
+            Octets::from_vec(fixture.preimage_data.clone())
+        )));
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_provide_preimage_not_solicited() -> Result<(), Box<dyn Error>> {
         let fixture = ProvideTestFixture::default();

pvm/pvm-host/src/host_functions/general/mod.rs

@@ -263,10 +263,13 @@ impl<S: HostStateProvider> GeneralHostFunction<S> {
         // --- Check Preimage (Err: NONE)
 
         let service_id_reg = vm.read_reg(7);
-        let service_id = if service_id_reg == u64::MAX || service_id_reg == service_id as u64 {
+        let service_id = if service_id_reg == u64::MAX {
             service_id
         } else {
-            service_id_reg as ServiceId
+            let Ok(service_id_reg) = vm.read_reg_as_service_id(7) else {
+                continue_none!()
+            };
+            service_id_reg
         };
         let Ok(Some(entry)) = accounts_sandbox
             .get_account_preimages_entry(state_provider, service_id, &hash)
@@ -340,7 +343,10 @@ impl<S: HostStateProvider> GeneralHostFunction<S> {
         let service_id = if service_id_reg == u64::MAX {
             service_id
         } else {
-            service_id_reg as ServiceId
+            let Ok(service_id_reg) = vm.read_reg_as_service_id(7) else {
+                continue_none!()
+            };
+            service_id_reg
         };
         let Ok(Some(entry)) = accounts_sandbox
             .get_account_storage_entry(state_provider, service_id, &storage_key)
@@ -533,7 +539,10 @@ impl<S: HostStateProvider> GeneralHostFunction<S> {
         let service_id = if service_id_reg == u64::MAX {
             service_id
         } else {
-            service_id_reg as ServiceId
+            let Ok(service_id_reg) = vm.read_reg_as_service_id(7) else {
+                continue_none!()
+            };
+            service_id_reg
         };
         let Ok(Some(metadata)) = accounts_sandbox
             .get_account_metadata(state_provider, service_id)

pvm/pvm-host/src/host_functions/general/tests.rs

@@ -1103,6 +1103,34 @@ mod lookup_tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_lookup_service_id_out_of_range_returns_none() -> Result<(), Box<dyn Error>> {
+        setup_tracing();
+        let fixture = LookupTestFixture::default();
+        let out_of_range_service_id = (1u64 << 32) + fixture.other_service_id as u64;
+        let vm = fixture
+            .prepare_vm_builder()?
+            .with_reg(7, out_of_range_service_id)
+            .with_mem_readable_range(fixture.mem_readable_range.clone())?
+            .with_mem_writable_range(fixture.mem_writable_range.clone())?
+            .build();
+
+        let state_provider =
+            Arc::new(fixture.prepare_state_provider(Some(fixture.other_service_id)));
+        let mut context = fixture
+            .prepare_invocation_context(state_provider.clone())
+            .await?;
+        let res = GeneralHostFunction::<MockStateManager>::host_lookup(
+            fixture.accumulate_host,
+            &vm,
+            state_provider,
+            &mut context,
+        )
+        .await?;
+        assert_eq!(res, LookupTestFixture::host_call_result_none());
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_lookup_account_not_found() -> Result<(), Box<dyn Error>> {
         setup_tracing();
@@ -1402,6 +1430,38 @@ mod read_tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_read_service_id_out_of_range_returns_none() -> Result<(), Box<dyn Error>> {
+        setup_tracing();
+        let fixture = ReadTestFixture::default();
+        let out_of_range_service_id = (1u64 << 32) + fixture.other_service_id as u64;
+        let vm = fixture
+            .prepare_vm_builder()?
+            .with_reg(7, out_of_range_service_id)
+            .with_mem_data(
+                fixture.storage_key_mem_offset,
+                fixture.storage_key.as_slice(),
+            )?
+            .with_mem_readable_range(fixture.mem_readable_range.clone())?
+            .with_mem_writable_range(fixture.mem_writable_range.clone())?
+            .build();
+
+        let state_provider =
+            Arc::new(fixture.prepare_state_provider(Some(fixture.other_service_id)));
+        let mut context = fixture
+            .prepare_invocation_context(state_provider.clone())
+            .await?;
+        let res = GeneralHostFunction::<MockStateManager>::host_read(
+            fixture.accumulate_host,
+            &vm,
+            state_provider,
+            &mut context,
+        )
+        .await?;
+        assert_eq!(res, ReadTestFixture::host_call_result_none());
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_read_account_not_found() -> Result<(), Box<dyn Error>> {
         setup_tracing();
@@ -2061,6 +2121,30 @@ mod info_tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_info_service_id_out_of_range_returns_none() -> Result<(), Box<dyn Error>> {
+        let fixture = InfoTestFixture::default();
+        let out_of_range_service_id = (1u64 << 32) + fixture.info_service_id as u64;
+        let vm = fixture
+            .prepare_vm_builder()?
+            .with_reg(7, out_of_range_service_id)
+            .build();
+        let state_provider = Arc::new(fixture.prepare_state_provider());
+        let mut context = fixture
+            .prepare_invocation_context(state_provider.clone())
+            .await?;
+
+        let res = GeneralHostFunction::<MockStateManager>::host_info(
+            fixture.accumulate_host,
+            &vm,
+            state_provider,
+            &mut context,
+        )
+        .await?;
+        assert_eq!(res, InfoTestFixture::host_call_result_none());
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_info_successful_accumulate_host() -> Result<(), Box<dyn Error>> {
         let fixture = InfoTestFixture::default();

pvm/pvm-host/src/host_functions/refine/mod.rs

@@ -64,7 +64,7 @@ impl<S: HostStateProvider> RefineHostFunction<S> {
 
         let service_id_reg = vm.read_reg(7);
         let service_id = if service_id_reg == u64::MAX
-            || state_provider.account_exists(refine_service_id).await?
+            && state_provider.account_exists(refine_service_id).await?
         {
             refine_service_id
         } else {