roshi_client/
access.rs

1//! Client helpers for vault access allowlists.
2
3use roshi_interface::access::{
4    access_merkle_leaf, access_merkle_node, EMPTY_ACCESS_MERKLE_ROOT, MAX_ACCESS_PROOF_LEN,
5};
6use solana_pubkey::Pubkey;
7
8#[derive(Clone, Debug, Eq, PartialEq)]
9pub struct AccessMerkleTree {
10    entries: Vec<AccessMerkleEntry>,
11    layers: Vec<Vec<[u8; 32]>>,
12}
13
14#[derive(Clone, Copy, Debug, Eq, PartialEq)]
15struct AccessMerkleEntry {
16    owner: Pubkey,
17    leaf: [u8; 32],
18}
19
20impl AccessMerkleTree {
21    pub fn new(owners: impl IntoIterator<Item = Pubkey>) -> Self {
22        let mut entries = owners
23            .into_iter()
24            .map(|owner| AccessMerkleEntry {
25                owner,
26                leaf: access_merkle_leaf(&owner),
27            })
28            .collect::<Vec<_>>();
29
30        entries.sort_by(|left, right| {
31            left.leaf
32                .cmp(&right.leaf)
33                .then_with(|| left.owner.to_bytes().cmp(&right.owner.to_bytes()))
34        });
35        entries.dedup_by(|left, right| left.owner == right.owner);
36
37        let mut layers = Vec::new();
38        if !entries.is_empty() {
39            layers.push(entries.iter().map(|entry| entry.leaf).collect::<Vec<_>>());
40
41            while layers.last().is_some_and(|layer| layer.len() > 1) {
42                let current = layers.last().expect("layer exists");
43                let mut next = Vec::with_capacity(current.len().div_ceil(2));
44
45                for pair in current.chunks(2) {
46                    if let [left, right] = pair {
47                        next.push(access_merkle_node(left, right));
48                    } else {
49                        next.push(pair[0]);
50                    }
51                }
52
53                layers.push(next);
54            }
55        }
56
57        Self { entries, layers }
58    }
59
60    pub fn is_empty(&self) -> bool {
61        self.entries.is_empty()
62    }
63
64    pub fn root(&self) -> [u8; 32] {
65        self.layers
66            .last()
67            .and_then(|layer| layer.first())
68            .copied()
69            .unwrap_or(EMPTY_ACCESS_MERKLE_ROOT)
70    }
71
72    pub fn proof(&self, owner: &Pubkey) -> Option<Vec<[u8; 32]>> {
73        let mut index = self
74            .entries
75            .iter()
76            .position(|entry| entry.owner == *owner)?;
77        let mut proof = Vec::new();
78
79        for layer in &self.layers {
80            if layer.len() <= 1 {
81                break;
82            }
83
84            let sibling_index = if index % 2 == 0 {
85                index + 1
86            } else {
87                index.saturating_sub(1)
88            };
89
90            if let Some(sibling) = layer.get(sibling_index) {
91                proof.push(*sibling);
92            }
93
94            index /= 2;
95        }
96
97        if proof.len() <= MAX_ACCESS_PROOF_LEN {
98            Some(proof)
99        } else {
100            None
101        }
102    }
103
104    pub fn contains(&self, owner: &Pubkey) -> bool {
105        self.entries.iter().any(|entry| entry.owner == *owner)
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112    use roshi_interface::access::verify_access_merkle_proof;
113
114    #[test]
115    fn empty_tree_uses_closed_root_and_has_no_proofs() {
116        let tree = AccessMerkleTree::new(Vec::<Pubkey>::new());
117
118        assert!(tree.is_empty());
119        assert_eq!(tree.root(), EMPTY_ACCESS_MERKLE_ROOT);
120        assert_eq!(tree.proof(&Pubkey::new_unique()), None);
121    }
122
123    #[test]
124    fn builds_verifiable_proofs() {
125        let owners = [
126            Pubkey::new_unique(),
127            Pubkey::new_unique(),
128            Pubkey::new_unique(),
129            Pubkey::new_unique(),
130            Pubkey::new_unique(),
131        ];
132        let tree = AccessMerkleTree::new(owners);
133        let root = tree.root();
134
135        for owner in owners {
136            let proof = tree.proof(&owner).unwrap();
137            assert!(verify_access_merkle_proof(&owner, &root, &proof));
138        }
139    }
140
141    #[test]
142    fn root_is_stable_across_input_order() {
143        let owners = [
144            Pubkey::new_unique(),
145            Pubkey::new_unique(),
146            Pubkey::new_unique(),
147            Pubkey::new_unique(),
148            Pubkey::new_unique(),
149        ];
150        let mut reversed = owners;
151        reversed.reverse();
152
153        assert_eq!(
154            AccessMerkleTree::new(owners).root(),
155            AccessMerkleTree::new(reversed).root()
156        );
157    }
158
159    #[test]
160    fn duplicate_owners_do_not_change_root() {
161        let owner = Pubkey::new_unique();
162        let other = Pubkey::new_unique();
163
164        assert_eq!(
165            AccessMerkleTree::new([owner, other]).root(),
166            AccessMerkleTree::new([owner, other, owner]).root()
167        );
168    }
169
170    #[test]
171    fn missing_owner_has_no_proof() {
172        let tree = AccessMerkleTree::new([Pubkey::new_unique()]);
173
174        assert_eq!(tree.proof(&Pubkey::new_unique()), None);
175    }
176}
roshi_client/access.rs

roshi_client/
access.rs
