1use crate::error::PrivacyError;
10use crate::types::{
11 CapabilityDescriptor, CapabilityKind, PrivacyCapability, PrivacyProof, ProofScheme,
12};
13use sha2::{Digest, Sha256};
14use serde::{Deserialize, Serialize};
15
16use pqc_sig::fips204::MlDsa65Keypair;
17use pqc_sig::types::{SigAlgorithm, SigPublicKey, Signature};
18
19extern crate alloc;
20use alloc::{string::{String, ToString}, vec::Vec};
21
22pub struct UniNode {
26 pub node_id: String,
27 capabilities: Vec<PrivacyCapability>,
28}
29
30impl UniNode {
31 pub fn new(node_id: impl Into<String>) -> Self {
32 Self {
33 node_id: node_id.into(),
34 capabilities: Vec::new(),
35 }
36 }
37
38 pub fn register_capability(&mut self, cap: PrivacyCapability) {
40 self.capabilities.push(cap);
41 }
42
43 pub fn advertise(
48 &self,
49 chaos_seed: &[u8; 32],
50 ) -> Result<CapabilityDescriptor, PrivacyError> {
51 if self.capabilities.is_empty() {
52 return Err(PrivacyError::CapabilityNotFound("No capabilities registered".into()));
53 }
54
55 let mut nonce_hasher = Sha256::new();
57 nonce_hasher.update(chaos_seed);
58 nonce_hasher.update(self.node_id.as_bytes());
59 nonce_hasher.update(b"uni-nonce-v1");
60 let nonce = hex::encode(nonce_hasher.finalize())[..16].to_string();
61
62 let mut payload = Vec::new();
64 payload.extend_from_slice(self.node_id.as_bytes());
65 for cap in &self.capabilities {
66 payload.extend_from_slice(alloc::format!("{:?}", cap.kind).as_bytes());
67 }
68 payload.extend_from_slice(nonce.as_bytes());
69 payload.extend_from_slice(chaos_seed);
70
71 let keypair = MlDsa65Keypair::from_secret_key_bytes(chaos_seed)
73 .map_err(|e| PrivacyError::Internal(alloc::format!("ML-DSA-65 keygen failed: {:?}", e)))?;
74
75 let sig = keypair.sign_deterministic(&payload)
76 .map_err(|e| PrivacyError::Internal(alloc::format!("ML-DSA-65 sign failed: {:?}", e)))?;
77
78 let signature = hex::encode(&sig.bytes);
79
80 Ok(CapabilityDescriptor {
81 node_id: self.node_id.clone(),
82 capabilities: self.capabilities.clone(),
83 nonce,
84 signature,
85 })
86 }
87
88 pub fn verify_advertisement(
93 descriptor: &CapabilityDescriptor,
94 chaos_seed: &[u8; 32],
95 ) -> Result<bool, PrivacyError> {
96 let mut payload = Vec::new();
98 payload.extend_from_slice(descriptor.node_id.as_bytes());
99 for cap in &descriptor.capabilities {
100 payload.extend_from_slice(alloc::format!("{:?}", cap.kind).as_bytes());
101 }
102 payload.extend_from_slice(descriptor.nonce.as_bytes());
103 payload.extend_from_slice(chaos_seed);
104
105 let keypair = MlDsa65Keypair::from_secret_key_bytes(chaos_seed)
107 .map_err(|e| PrivacyError::Internal(alloc::format!("ML-DSA-65 keygen failed: {:?}", e)))?;
108 let pk = keypair.public_key();
109
110 let sig_bytes = hex::decode(&descriptor.signature)
112 .map_err(|e| PrivacyError::Internal(alloc::format!("hex decode signature: {}", e)))?;
113
114 let pk_typed = SigPublicKey::new(SigAlgorithm::MlDsa65, pk.bytes);
115 let sig_typed = Signature::new(SigAlgorithm::MlDsa65, sig_bytes);
116
117 Ok(MlDsa65Keypair::verify(&pk_typed, &payload, &sig_typed).is_ok())
118 }
119
120 pub fn route_ibc_call(
122 &self,
123 capability: CapabilityKind,
124 payload: &[u8],
125 chaos_seed: &[u8; 32],
126 ) -> Result<Vec<u8>, PrivacyError> {
127 let cap = self.capabilities.iter()
128 .find(|c| c.kind == capability)
129 .ok_or_else(|| PrivacyError::CapabilityNotFound(alloc::format!("{capability:?}")))?;
130
131 let mut hasher = Sha256::new();
133 hasher.update(payload);
134 hasher.update(chaos_seed);
135 hasher.update(alloc::format!("{:?}", cap.kind).as_bytes());
136 hasher.update(b"ibc-route-v1");
137 Ok(hasher.finalize().to_vec())
138 }
139}
140
141#[derive(Debug, Clone, Serialize, Deserialize)]
145pub struct WasmCircuit {
146 pub circuit_id: String,
147 pub wasm_hash: String,
149 pub input: Vec<u8>,
151 pub privacy_tags: Vec<String>,
153 pub tenant_did: String,
155}
156
157#[derive(Debug, Clone, Serialize, Deserialize)]
159pub struct WasmResult {
160 pub circuit_id: String,
161 pub output: Vec<u8>,
163 pub attestation: PrivacyProof,
165}
166
167pub struct UviInterface {
169 pub node_id: String,
170}
171
172impl UviInterface {
173 pub fn new(node_id: impl Into<String>) -> Self {
174 Self { node_id: node_id.into() }
175 }
176
177 pub fn execute(
182 &self,
183 circuit: &WasmCircuit,
184 chaos_seed: &[u8; 32],
185 ) -> Result<WasmResult, PrivacyError> {
186 if circuit.wasm_hash.is_empty() {
188 return Err(PrivacyError::WasmExecutionFailed("Empty WASM hash".into()));
189 }
190
191 let mut exec_hasher = Sha256::new();
193 exec_hasher.update(&circuit.input);
194 exec_hasher.update(chaos_seed);
195 exec_hasher.update(circuit.wasm_hash.as_bytes());
196 for tag in &circuit.privacy_tags {
197 exec_hasher.update(tag.as_bytes());
198 }
199 exec_hasher.update(b"aqvm-exec-v1");
200 let output: Vec<u8> = exec_hasher.finalize().to_vec();
201
202 let attestation = self.attest(&output, &circuit.tenant_did, chaos_seed);
204
205 Ok(WasmResult {
206 circuit_id: circuit.circuit_id.clone(),
207 output,
208 attestation,
209 })
210 }
211
212 fn attest(&self, output: &[u8], tenant_did: &str, chaos_seed: &[u8; 32]) -> PrivacyProof {
213 let mut hasher = Sha256::new();
214 hasher.update(output);
215 hasher.update(tenant_did.as_bytes());
216 hasher.update(chaos_seed);
217 hasher.update(b"pqc-attest-v1");
218 let commitment: [u8; 32] = hasher.finalize().into();
219 PrivacyProof {
220 proof_bytes: hex::encode(commitment),
221 public_inputs: hex::encode(Sha256::digest(output)),
222 scheme: ProofScheme::Snark,
223 security_bits: 128,
224 proof_size: 64,
225 chsh_value: 0.0,
226 lyapunov: 4.5,
227 }
228 }
229}
230
231#[cfg(test)]
232mod tests {
233 use super::*;
234
235 #[test]
236 fn test_uni_advertise() {
237 let mut uni = UniNode::new("node-1");
238 uni.register_capability(PrivacyCapability {
239 kind: CapabilityKind::Fhe,
240 scheme: "CKKS".into(),
241 security_level: 128,
242 recursion: false,
243 });
244 let seed = [0u8; 32];
245 let desc = uni.advertise(&seed).unwrap();
246 assert_eq!(desc.node_id, "node-1");
247 assert_eq!(desc.capabilities.len(), 1);
248 assert_eq!(desc.signature.len(), 3309 * 2);
250 }
251
252 #[test]
253 fn test_uni_advertise_verify() {
254 let mut uni = UniNode::new("node-2");
255 uni.register_capability(PrivacyCapability {
256 kind: CapabilityKind::Zk,
257 scheme: "Plonk".into(),
258 security_level: 128,
259 recursion: true,
260 });
261 let seed = [1u8; 32];
262 let desc = uni.advertise(&seed).unwrap();
263 let valid = UniNode::verify_advertisement(&desc, &seed).unwrap();
264 assert!(valid);
265 }
266
267 #[test]
268 fn test_uni_advertise_verify_wrong_seed_fails() {
269 let mut uni = UniNode::new("node-3");
270 uni.register_capability(PrivacyCapability {
271 kind: CapabilityKind::Dp,
272 scheme: "Laplace".into(),
273 security_level: 128,
274 recursion: false,
275 });
276 let seed = [2u8; 32];
277 let wrong_seed = [3u8; 32];
278 let desc = uni.advertise(&seed).unwrap();
279 let valid = UniNode::verify_advertisement(&desc, &wrong_seed).unwrap();
281 assert!(!valid);
282 }
283
284 #[test]
285 fn test_uvi_execute() {
286 let uvi = UviInterface::new("uvi-1");
287 let seed = [0u8; 32];
288 let circuit = WasmCircuit {
289 circuit_id: "c1".into(),
290 wasm_hash: "deadbeef".into(),
291 input: b"input".to_vec(),
292 privacy_tags: vec!["fhe-enabled".into()],
293 tenant_did: "did:wyqcc:tenant1".into(),
294 };
295 let result = uvi.execute(&circuit, &seed).unwrap();
296 assert_eq!(result.circuit_id, "c1");
297 assert!(!result.output.is_empty());
298 }
299}