From 6c9ee8896a594a50c364a114f994eaac53f1036d Mon Sep 17 00:00:00 2001
From: agureev <artguril@proton.me>
Date: Wed, 10 Jun 2026 21:59:55 +0400
Subject: [PATCH] feat(lee): make EncryptedAccountData guest-availiable

Likewise moves the epk definition outside of host-gated module
---
 lee/state_machine/core/src/encryption/mod.rs  | 46 ++++++++++++++++++-
 .../src/encryption/shared_key_derivation.rs   |  7 +--
 lee/state_machine/core/src/lib.rs             |  4 +-
 3 files changed, 49 insertions(+), 8 deletions(-)

diff --git a/lee/state_machine/core/src/encryption/mod.rs b/lee/state_machine/core/src/encryption/mod.rs
index 37745d4f..983dd330 100644
--- a/lee/state_machine/core/src/encryption/mod.rs
+++ b/lee/state_machine/core/src/encryption/mod.rs
@@ -6,7 +6,7 @@ use chacha20::{
 use risc0_zkvm::sha::{Impl, Sha256 as _};
 use serde::{Deserialize, Serialize};
 #[cfg(feature = "host")]
-pub use shared_key_derivation::{EphemeralPublicKey, MlKem768EncapsulationKey, ViewingPublicKey};
+pub use shared_key_derivation::{MlKem768EncapsulationKey, ViewingPublicKey};
 
 use crate::{Commitment, account::Account, program::PrivateAccountKind};
 #[cfg(feature = "host")]
@@ -17,6 +17,11 @@ pub type Scalar = [u8; 32];
 #[derive(Serialize, Deserialize, Clone, Copy)]
 pub struct SharedSecretKey(pub [u8; 32]);
 
+/// The ML-KEM-768 ciphertext produced during encapsulation; transmitted on-wire in place of the
+/// former ECDH ephemeral public key. Always 1088 bytes for ML-KEM-768.
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
+pub struct EphemeralPublicKey(pub Vec<u8>);
+
 pub struct EncryptionScheme;
 
 #[derive(Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
@@ -36,6 +41,45 @@ impl std::fmt::Debug for Ciphertext {
     }
 }
 
+pub type ViewTag = u8;
+
+/// Encrypted private-account note for one output
+#[derive(Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
+#[cfg_attr(any(feature = "host", test), derive(Debug, Clone, PartialEq, Eq))]
+pub struct EncryptedAccountData {
+    pub ciphertext: Ciphertext,
+    pub epk: EphemeralPublicKey,
+    pub view_tag: ViewTag,
+}
+
+#[cfg(feature = "host")]
+impl EncryptedAccountData {
+    #[must_use]
+    pub fn new(
+        ciphertext: Ciphertext,
+        npk: &crate::NullifierPublicKey,
+        vpk: &ViewingPublicKey,
+        epk: EphemeralPublicKey,
+    ) -> Self {
+        let view_tag = Self::compute_view_tag(npk, vpk);
+        Self {
+            ciphertext,
+            epk,
+            view_tag,
+        }
+    }
+
+    /// Computes the tag as the first byte of SHA256("/LEE/v0.3/ViewTag/" || npk || vpk).
+    #[must_use]
+    pub fn compute_view_tag(npk: &crate::NullifierPublicKey, vpk: &ViewingPublicKey) -> ViewTag {
+        let mut bytes = Vec::new();
+        bytes.extend_from_slice(b"/LEE/v0.3/ViewTag/");
+        bytes.extend_from_slice(&npk.to_byte_array());
+        bytes.extend_from_slice(vpk.to_bytes());
+        Impl::hash_bytes(&bytes).as_bytes()[0]
+    }
+}
+
 impl EncryptionScheme {
     #[must_use]
     pub fn encrypt(
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 a3b2fd32..5c982c6f 100644
--- a/lee/state_machine/core/src/encryption/shared_key_derivation.rs
+++ b/lee/state_machine/core/src/encryption/shared_key_derivation.rs
@@ -2,12 +2,7 @@ use borsh::{BorshDeserialize, BorshSerialize};
 use ml_kem::{Decapsulate as _, Encapsulate as _, KeyExport as _, Seed};
 use serde::{Deserialize, Serialize};
 
-use crate::SharedSecretKey;
-
-/// The ML-KEM-768 ciphertext produced during encapsulation; transmitted on-wire in place of the
-/// former ECDH ephemeral public key.  Always 1088 bytes for ML-KEM-768.
-#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
-pub struct EphemeralPublicKey(pub Vec<u8>);
+use crate::{EphemeralPublicKey, SharedSecretKey};
 
 /// ML-KEM-768 encapsulation key bytes (1184 bytes, opaque to this crate).
 #[derive(
diff --git a/lee/state_machine/core/src/lib.rs b/lee/state_machine/core/src/lib.rs
index 466e1f5d..9ad2858e 100644
--- a/lee/state_machine/core/src/lib.rs
+++ b/lee/state_machine/core/src/lib.rs
@@ -10,7 +10,9 @@ pub use commitment::{
     Commitment, CommitmentSetDigest, DUMMY_COMMITMENT, DUMMY_COMMITMENT_HASH, MembershipProof,
     compute_digest_for_path,
 };
-pub use encryption::{EncryptionScheme, SharedSecretKey};
+pub use encryption::{
+    EncryptedAccountData, EncryptionScheme, EphemeralPublicKey, SharedSecretKey, ViewTag,
+};
 pub use nullifier::{Identifier, Nullifier, NullifierPublicKey, NullifierSecretKey};
 pub use program::PrivateAccountKind;