Skip to main content

cougr_core/zk/
systems.rs

1use crate::simple_world::SimpleWorld;
2use soroban_sdk::{Bytes, BytesN, Env, Symbol};
3
4use super::components::{COMMIT_REVEAL_TYPE, VERIFIED_MARKER_TYPE};
5use super::interfaces::{Groth16ProofVerifier, ProofVerifier};
6use super::types::{Groth16Proof, Scalar, VerificationKey};
7
8/// Read a big-endian `u64` from `data` at byte offset `offset`.
9fn read_u64(data: &Bytes, offset: u32) -> u64 {
10    let mut arr = [0u8; 8];
11    for i in 0..8u32 {
12        arr[i as usize] = data.get(offset + i).unwrap();
13    }
14    u64::from_be_bytes(arr)
15}
16
17/// Read a boolean from `data` at byte offset `offset`.
18fn read_bool(data: &Bytes, offset: u32) -> bool {
19    data.get(offset).unwrap() != 0
20}
21
22/// Encode a verified-marker component as `Bytes`.
23///
24/// Layout: `[verified_at: u64 (8 bytes, big-endian)]`.
25pub fn encode_verified_marker(env: &Env, verified_at: u64) -> Bytes {
26    Bytes::from_slice(env, &verified_at.to_be_bytes())
27}
28
29/// Decode a verified-marker component's `verified_at` timestamp.
30pub fn decode_verified_at(data: &Bytes) -> u64 {
31    read_u64(data, 0)
32}
33
34/// Encode a commit-reveal component as `Bytes`.
35///
36/// Layout: `[commitment: 32 bytes | reveal_deadline: u64 (8 bytes) | revealed: u8 (1 byte)]`.
37/// Total: 41 bytes.
38pub fn encode_commit_reveal(
39    env: &Env,
40    commitment: &BytesN<32>,
41    reveal_deadline: u64,
42    revealed: bool,
43) -> Bytes {
44    let mut b: Bytes = commitment.clone().into();
45    b.append(&Bytes::from_slice(env, &reveal_deadline.to_be_bytes()));
46    b.push_back(if revealed { 1 } else { 0 });
47    b
48}
49
50/// Verify a proof for a specific entity and mark it as verified on success.
51///
52/// Unlike a world-scanning system, this takes the proof as a parameter
53/// (proofs are large and should not be stored in the ECS). On successful
54/// verification, a `VerifiedMarker` component is added to the entity.
55///
56/// Returns `true` if the proof was valid.
57pub fn verify_proofs_with<
58    V: ProofVerifier<VerificationKey = VerificationKey, Proof = Groth16Proof, PublicInput = Scalar>,
59>(
60    world: &mut SimpleWorld,
61    env: &Env,
62    entity_id: u32,
63    verifier: &V,
64    vk: &VerificationKey,
65    proof: &Groth16Proof,
66    public_inputs: &[Scalar],
67) -> Result<bool, super::error::ZKError> {
68    let verified_sym = Symbol::new(env, VERIFIED_MARKER_TYPE);
69    let is_valid = verifier.verify(env, vk, proof, public_inputs)?;
70
71    if is_valid {
72        let now = env.ledger().timestamp();
73        let marker_data = encode_verified_marker(env, now);
74        world.add_component(entity_id, verified_sym, marker_data);
75    }
76
77    Ok(is_valid)
78}
79
80pub fn verify_proofs_system(
81    world: &mut SimpleWorld,
82    env: &Env,
83    entity_id: u32,
84    vk: &VerificationKey,
85    proof: &Groth16Proof,
86    public_inputs: &[Scalar],
87) -> bool {
88    verify_proofs_with(
89        world,
90        env,
91        entity_id,
92        &Groth16ProofVerifier,
93        vk,
94        proof,
95        public_inputs,
96    )
97    .unwrap_or(false)
98}
99
100/// Check for expired commit-reveal deadlines.
101///
102/// Removes `CommitReveal` components that have passed their deadline
103/// without being revealed.
104pub fn commit_reveal_deadline_system(world: &mut SimpleWorld, env: &Env) {
105    let cr_sym = Symbol::new(env, COMMIT_REVEAL_TYPE);
106    let entities = world.get_entities_with_component(&cr_sym, env);
107    let now = env.ledger().timestamp();
108
109    for i in 0..entities.len() {
110        let entity_id = entities.get(i).unwrap();
111
112        if let Some(data) = world.get_component(entity_id, &cr_sym) {
113            let deadline = read_u64(&data, 32);
114            let revealed = read_bool(&data, 40);
115
116            if !revealed && now > deadline {
117                world.remove_component(entity_id, &cr_sym);
118            }
119        }
120    }
121}
122
123/// Remove `VerifiedMarker` components older than `max_age` ledgers.
124///
125/// This prevents verified markers from accumulating indefinitely.
126pub fn cleanup_verified_system(world: &mut SimpleWorld, env: &Env, max_age: u64) {
127    let verified_sym = Symbol::new(env, VERIFIED_MARKER_TYPE);
128    let entities = world.get_entities_with_component(&verified_sym, env);
129    let now = env.ledger().timestamp();
130
131    for i in 0..entities.len() {
132        let entity_id = entities.get(i).unwrap();
133
134        if let Some(data) = world.get_component(entity_id, &verified_sym) {
135            let verified_at = read_u64(&data, 0);
136
137            if now.saturating_sub(verified_at) > max_age {
138                world.remove_component(entity_id, &verified_sym);
139            }
140        }
141    }
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147    use crate::zk::error::ZKError;
148    use soroban_sdk::Env;
149
150    struct RejectingVerifier;
151
152    impl ProofVerifier for RejectingVerifier {
153        type VerificationKey = VerificationKey;
154        type Proof = Groth16Proof;
155        type PublicInput = Scalar;
156
157        fn verify(
158            &self,
159            _env: &Env,
160            _verification_key: &Self::VerificationKey,
161            _proof: &Self::Proof,
162            _public_inputs: &[Self::PublicInput],
163        ) -> Result<bool, ZKError> {
164            Ok(false)
165        }
166    }
167
168    #[test]
169    fn test_commit_reveal_deadline_keeps_non_expired() {
170        let env = Env::default();
171        let mut world = SimpleWorld::new(&env);
172        let e1 = world.spawn_entity();
173
174        let commitment = BytesN::from_array(&env, &[0xABu8; 32]);
175        let cr_data = encode_commit_reveal(&env, &commitment, 1000, false);
176        let cr_sym = Symbol::new(&env, COMMIT_REVEAL_TYPE);
177        world.add_component(e1, cr_sym.clone(), cr_data);
178
179        // now = 0, deadline = 1000, not expired
180        commit_reveal_deadline_system(&mut world, &env);
181        assert!(world.has_component(e1, &cr_sym));
182    }
183
184    #[test]
185    fn test_commit_reveal_keeps_revealed() {
186        let env = Env::default();
187        let mut world = SimpleWorld::new(&env);
188        let e1 = world.spawn_entity();
189
190        let commitment = BytesN::from_array(&env, &[0xABu8; 32]);
191        let cr_data = encode_commit_reveal(&env, &commitment, 0, true);
192        let cr_sym = Symbol::new(&env, COMMIT_REVEAL_TYPE);
193        world.add_component(e1, cr_sym.clone(), cr_data);
194
195        commit_reveal_deadline_system(&mut world, &env);
196        // Revealed commitments are not removed even if past deadline
197        assert!(world.has_component(e1, &cr_sym));
198    }
199
200    #[test]
201    fn test_cleanup_verified_no_markers() {
202        let env = Env::default();
203        let mut world = SimpleWorld::new(&env);
204        // Should not panic with empty world
205        cleanup_verified_system(&mut world, &env, 100);
206    }
207
208    #[test]
209    fn test_cleanup_verified_keeps_recent() {
210        let env = Env::default();
211        let mut world = SimpleWorld::new(&env);
212        let e1 = world.spawn_entity();
213
214        let marker_data = encode_verified_marker(&env, 0);
215        let verified_sym = Symbol::new(&env, VERIFIED_MARKER_TYPE);
216        world.add_component(e1, verified_sym.clone(), marker_data);
217
218        // max_age is 1000, marker is at time 0, now is 0, age = 0 <= 1000
219        cleanup_verified_system(&mut world, &env, 1000);
220        assert!(world.has_component(e1, &verified_sym));
221    }
222
223    #[test]
224    fn test_verify_proofs_with_invalid_result_does_not_mark_entity() {
225        let env = Env::default();
226        let mut world = SimpleWorld::new(&env);
227        let entity_id = world.spawn_entity();
228        let verifier = RejectingVerifier;
229        let vk = VerificationKey {
230            alpha: super::super::types::G1Point {
231                bytes: BytesN::from_array(&env, &[0u8; 64]),
232            },
233            beta: super::super::types::G2Point {
234                bytes: BytesN::from_array(&env, &[0u8; 128]),
235            },
236            gamma: super::super::types::G2Point {
237                bytes: BytesN::from_array(&env, &[0u8; 128]),
238            },
239            delta: super::super::types::G2Point {
240                bytes: BytesN::from_array(&env, &[0u8; 128]),
241            },
242            ic: soroban_sdk::Vec::new(&env),
243        };
244        let proof = Groth16Proof {
245            a: super::super::types::G1Point {
246                bytes: BytesN::from_array(&env, &[0u8; 64]),
247            },
248            b: super::super::types::G2Point {
249                bytes: BytesN::from_array(&env, &[0u8; 128]),
250            },
251            c: super::super::types::G1Point {
252                bytes: BytesN::from_array(&env, &[0u8; 64]),
253            },
254        };
255
256        let result =
257            verify_proofs_with(&mut world, &env, entity_id, &verifier, &vk, &proof, &[]).unwrap();
258        assert!(!result);
259        assert!(!world.has_component(entity_id, &Symbol::new(&env, VERIFIED_MARKER_TYPE)));
260    }
261}