From 2fd8c1b15745d818bf5c16c0216863c8a53c838e Mon Sep 17 00:00:00 2001
From: agureev <artguril@proton.me>
Date: Fri, 19 Jun 2026 22:29:31 +0400
Subject: [PATCH] refactor: expose ML KEM and esk in-guest

---
 lee/state_machine/core/Cargo.toml                |  4 ++--
 lee/state_machine/core/src/encryption/mod.rs     | 16 ++++++++++++++--
 .../core/src/encryption/shared_key_derivation.rs |  7 +++++--
 3 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/lee/state_machine/core/Cargo.toml b/lee/state_machine/core/Cargo.toml
index 6e1f0ff0..585637ae 100644
--- a/lee/state_machine/core/Cargo.toml
+++ b/lee/state_machine/core/Cargo.toml
@@ -16,7 +16,7 @@ thiserror.workspace = true
 bytemuck.workspace = true
 bytesize.workspace = true
 base58.workspace = true
-ml-kem = { workspace = true, optional = true, features = ["getrandom"] }
+ml-kem = { workspace = true }
 chacha20 = { version = "0.10" }
 
 [dev-dependencies]
@@ -24,5 +24,5 @@ serde_json.workspace = true
 
 [features]
 default = []
-host = ["dep:ml-kem"]
+host = ["ml-kem/getrandom"]
 test_utils = ["host"]
diff --git a/lee/state_machine/core/src/encryption/mod.rs b/lee/state_machine/core/src/encryption/mod.rs
index 5fa80b60..f3895924 100644
--- a/lee/state_machine/core/src/encryption/mod.rs
+++ b/lee/state_machine/core/src/encryption/mod.rs
@@ -5,15 +5,15 @@ use chacha20::{
 };
 use risc0_zkvm::sha::{Impl, Sha256 as _};
 use serde::{Deserialize, Serialize};
-#[cfg(feature = "host")]
 pub use shared_key_derivation::{MlKem768EncapsulationKey, ViewingPublicKey};
 
 use crate::{Commitment, account::Account, program::PrivateAccountKind};
-#[cfg(feature = "host")]
 pub mod shared_key_derivation;
 
 pub type Scalar = [u8; 32];
 
+pub type EphemeralSecretKey = [u8; 32];
+
 #[derive(Serialize, Deserialize, Clone, Copy)]
 pub struct SharedSecretKey(pub [u8; 32]);
 
@@ -52,6 +52,18 @@ pub struct EncryptedAccountData {
     pub view_tag: ViewTag,
 }
 
+impl EncryptedAccountData {
+    #[must_use]
+    pub fn compute_view_tag(npk: &crate::NullifierPublicKey, vpk: &ViewingPublicKey) -> ViewTag {
+        const PREFIX: &[u8; 18] = b"/LEE/v0.3/ViewTag/";
+        let mut bytes = [0_u8; 18 + 32 + ViewingPublicKey::LEN];
+        bytes[0..18].copy_from_slice(PREFIX);
+        bytes[18..50].copy_from_slice(&npk.to_byte_array());
+        bytes[50..].copy_from_slice(vpk.to_bytes());
+        Impl::hash_bytes(&bytes).as_bytes()[0]
+    }
+}
+
 #[cfg(feature = "host")]
 impl EncryptedAccountData {
     #[must_use]
diff --git a/lee/state_machine/core/src/encryption/shared_key_derivation.rs b/lee/state_machine/core/src/encryption/shared_key_derivation.rs
index 5c982c6f..2c046cbf 100644
--- a/lee/state_machine/core/src/encryption/shared_key_derivation.rs
+++ b/lee/state_machine/core/src/encryption/shared_key_derivation.rs
@@ -1,5 +1,7 @@
 use borsh::{BorshDeserialize, BorshSerialize};
-use ml_kem::{Decapsulate as _, Encapsulate as _, KeyExport as _, Seed};
+#[cfg(feature = "host")]
+use ml_kem::Encapsulate as _;
+use ml_kem::{Decapsulate as _, KeyExport as _, Seed};
 use serde::{Deserialize, Serialize};
 
 use crate::{EphemeralPublicKey, SharedSecretKey};
@@ -26,6 +28,7 @@ impl MlKem768EncapsulationKey {
     pub const LEN: usize = 1184;
 
     /// Construct from raw bytes, returning an error if the length is not [`Self::LEN`].
+    #[cfg(feature = "host")]
     pub fn from_bytes(bytes: Vec<u8>) -> Result<Self, crate::error::LeeCoreError> {
         if bytes.len() != Self::LEN {
             return Err(crate::error::LeeCoreError::DeserializationError(format!(
@@ -59,6 +62,7 @@ impl SharedSecretKey {
     /// Returns `(shared_secret, ciphertext)`.  The ciphertext must be included in the transaction
     /// as the `EphemeralPublicKey`; the receiver recovers the same shared secret via
     /// [`Self::decapsulate`].
+    #[cfg(feature = "host")]
     #[must_use]
     pub fn encapsulate(ek: &MlKem768EncapsulationKey) -> (Self, EphemeralPublicKey) {
         let ek_bytes: ml_kem::kem::Key<ml_kem::EncapsulationKey768> =
@@ -85,7 +89,6 @@ impl SharedSecretKey {
     /// avoid EPK collisions across multiple outputs in the same test.
     ///
     /// For production use [`Self::encapsulate`], which draws randomness from the OS.
-    #[cfg(any(test, feature = "test_utils"))]
     #[must_use]
     pub fn encapsulate_deterministic(
         ek: &MlKem768EncapsulationKey,