From 2619a2e7bb194b2512752fd9811b82a10631305f Mon Sep 17 00:00:00 2001
From: Milosz Muszynski <milosz.muszynski@gmail.com>
Date: Tue, 13 Feb 2024 11:39:35 +0100
Subject: [PATCH] piecrust-uplink: support bytecheck integrity check of
 arguments

---
 piecrust-uplink/CHANGELOG.md       | 5 +++++
 piecrust-uplink/Cargo.toml         | 2 +-
 piecrust-uplink/src/abi/helpers.rs | 9 ++++++---
 3 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/piecrust-uplink/CHANGELOG.md b/piecrust-uplink/CHANGELOG.md
index 08b87132..5be7d1c8 100644
--- a/piecrust-uplink/CHANGELOG.md
+++ b/piecrust-uplink/CHANGELOG.md
@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Changed
+
+- Change `wrap_call` function to support `bytecheck`-based integrity check of arguments [#324]
+
 ## [0.10.0] - 2024-01-24
 
 ### Added
@@ -152,6 +156,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - First `piecrust-uplink` release
 
 <!-- ISSUES -->
+[#324]: https://github.com/dusk-network/piecrust/issues/324
 [#301]: https://github.com/dusk-network/piecrust/issues/301
 [#271]: https://github.com/dusk-network/piecrust/issues/271
 [#268]: https://github.com/dusk-network/piecrust/issues/268
diff --git a/piecrust-uplink/Cargo.toml b/piecrust-uplink/Cargo.toml
index f39d99f9..08fa6a38 100644
--- a/piecrust-uplink/Cargo.toml
+++ b/piecrust-uplink/Cargo.toml
@@ -13,7 +13,7 @@ edition = "2021"
 license = "MPL-2.0"
 
 [dependencies]
-rkyv = { version = "0.7", default-features = false, features = ["size_32", "alloc"] }
+rkyv = { version = "0.7", default-features = false, features = ["size_32", "alloc", "validation"] }
 bytecheck = { version = "0.6", default-features = false }
 dlmalloc = { version = "0.2", optional = true, features = ["global"] }
 
diff --git a/piecrust-uplink/src/abi/helpers.rs b/piecrust-uplink/src/abi/helpers.rs
index ef9b9d56..8624cc55 100644
--- a/piecrust-uplink/src/abi/helpers.rs
+++ b/piecrust-uplink/src/abi/helpers.rs
@@ -11,7 +11,8 @@ use rkyv::ser::serializers::{
     BufferScratch, BufferSerializer, CompositeSerializer,
 };
 use rkyv::ser::Serializer;
-use rkyv::{archived_root, Archive, Deserialize, Infallible, Serialize};
+use rkyv::validation::validators::DefaultValidator;
+use rkyv::{check_archived_root, Archive, Deserialize, Infallible, Serialize};
 
 use crate::types::StandardBufSerializer;
 
@@ -21,14 +22,16 @@ use crate::types::StandardBufSerializer;
 pub fn wrap_call<A, R, F>(arg_len: u32, f: F) -> u32
 where
     A: Archive,
-    A::Archived: Deserialize<A, Infallible>,
+    A::Archived: Deserialize<A, Infallible>
+        + for<'b> bytecheck::CheckBytes<DefaultValidator<'b>>,
     R: for<'a> Serialize<StandardBufSerializer<'a>>,
     F: Fn(A) -> R,
 {
     with_arg_buf(|buf| {
         let slice = &buf[..arg_len as usize];
 
-        let aa: &A::Archived = unsafe { archived_root::<A>(slice) };
+        let aa: &A::Archived = check_archived_root::<A>(slice)
+            .expect("Argument should correctly deserialize");
         let a: A = aa.deserialize(&mut Infallible).unwrap();
 
         let ret = f(a);