Skip to main content

world_id_primitives/
key_set.rs

1use std::{
2    ops::{Deref, DerefMut},
3    str::FromStr,
4};
5
6use ark_babyjubjub::{EdwardsAffine, Fq};
7use ark_ff::AdditiveGroup;
8use arrayvec::ArrayVec;
9use eddsa_babyjubjub::EdDSAPublicKey;
10use ruint::aliases::U256;
11use serde::{Deserialize, Serialize};
12
13use crate::PrimitiveError;
14
15/// The maximum number of authenticator public keys that can be registered at any time
16/// per World ID Account.
17///
18/// This constrained is introduced to maintain proof performance reasonable even
19/// in devices with limited resources.
20pub const MAX_AUTHENTICATOR_KEYS: usize = 7;
21
22/// Errors for decoding sparse authenticator pubkey slots.
23#[derive(Debug, thiserror::Error, Clone, PartialEq, Eq)]
24pub enum SparseAuthenticatorPubkeysError {
25    /// The input contains a non-empty slot index outside supported bounds.
26    #[error(
27        "invalid authenticator pubkey slot {slot_index}; max supported slot is {max_supported_slot}"
28    )]
29    SlotOutOfBounds {
30        /// Slot index returned by the source.
31        slot_index: usize,
32        /// Highest supported slot index.
33        max_supported_slot: usize,
34    },
35    /// A slot contained bytes that are not a valid compressed public key.
36    #[error("invalid authenticator public key at slot {slot_index}: {reason}")]
37    InvalidCompressedPubkey {
38        /// Slot index of the invalid entry.
39        slot_index: usize,
40        /// Parse error details.
41        reason: String,
42    },
43}
44
45/// A set of **off-chain** authenticator public keys for a World ID Account.
46///
47/// Each World ID Account has a number of public keys for each authorized authenticator;
48/// a commitment to the entire set of public keys is stored in the `WorldIDRegistry` contract.
49///
50/// Removed authenticator slots are represented as `None` to preserve stable slot indices.
51#[derive(Debug, Clone, Serialize, Deserialize)]
52pub struct AuthenticatorPublicKeySet(ArrayVec<Option<EdDSAPublicKey>, MAX_AUTHENTICATOR_KEYS>);
53
54impl Default for AuthenticatorPublicKeySet {
55    fn default() -> Self {
56        Self(ArrayVec::new())
57    }
58}
59
60impl AuthenticatorPublicKeySet {
61    /// Creates a new authenticator public key set with the provided active public keys.
62    ///
63    /// The provided keys are inserted in order.
64    /// Sparse holes can be represented later via [`Self::try_clear_at_index`].
65    ///
66    /// # Errors
67    /// Returns an error if the number of public keys exceeds [`MAX_AUTHENTICATOR_KEYS`].
68    pub fn new(pubkeys: Vec<EdDSAPublicKey>) -> Result<Self, PrimitiveError> {
69        if pubkeys.len() > MAX_AUTHENTICATOR_KEYS {
70            return Err(PrimitiveError::OutOfBounds);
71        }
72
73        Ok(Self(
74            pubkeys
75                .into_iter()
76                .map(Some)
77                .collect::<ArrayVec<_, MAX_AUTHENTICATOR_KEYS>>(),
78        ))
79    }
80
81    /// Decodes sparse authenticator pubkey slots while preserving slot positions.
82    ///
83    /// Input may contain `None` entries for removed authenticators. Trailing empty slots are
84    /// trimmed, interior holes are preserved as `None`, and used slots beyond
85    /// [`MAX_AUTHENTICATOR_KEYS`] are rejected.
86    ///
87    /// # Errors
88    /// Returns [`SparseAuthenticatorPubkeysError`] if a used slot is out of bounds or any
89    /// compressed key bytes are invalid.
90    pub fn from_sparse_encoded_pubkeys(
91        pubkeys: Vec<Option<U256>>,
92    ) -> Result<Self, SparseAuthenticatorPubkeysError> {
93        let last_present_idx = pubkeys.iter().rposition(Option::is_some);
94        if let Some(idx) = last_present_idx
95            && idx >= MAX_AUTHENTICATOR_KEYS
96        {
97            return Err(SparseAuthenticatorPubkeysError::SlotOutOfBounds {
98                slot_index: idx,
99                max_supported_slot: MAX_AUTHENTICATOR_KEYS - 1,
100            });
101        }
102
103        let normalized_len = last_present_idx.map_or(0, |idx| idx + 1);
104        let decoded_pubkeys = pubkeys
105            .into_iter()
106            .take(normalized_len)
107            .enumerate()
108            .map(|(idx, pubkey)| match pubkey {
109                Some(pubkey) => EdDSAPublicKey::from_compressed_bytes(pubkey.to_le_bytes())
110                    .map(Some)
111                    .map_err(
112                        |e| SparseAuthenticatorPubkeysError::InvalidCompressedPubkey {
113                            slot_index: idx,
114                            reason: e.to_string(),
115                        },
116                    ),
117                None => Ok(None),
118            })
119            .collect::<Result<Vec<_>, _>>()?;
120
121        Ok(Self(
122            decoded_pubkeys
123                .into_iter()
124                .collect::<ArrayVec<_, MAX_AUTHENTICATOR_KEYS>>(),
125        ))
126    }
127
128    /// Converts the set of public keys to a fixed-length array of Affine points.
129    ///
130    /// This is usually used to serialize to circuit inputs, which require defaulted Affine points
131    /// for empty (`None`) slots.
132    pub fn as_affine_array(&self) -> [EdwardsAffine; MAX_AUTHENTICATOR_KEYS] {
133        let mut array = [EdwardsAffine::default(); MAX_AUTHENTICATOR_KEYS];
134        for (i, pubkey) in self.0.iter().enumerate() {
135            array[i] = pubkey
136                .as_ref()
137                .map_or_else(EdwardsAffine::default, |pubkey| pubkey.pk);
138        }
139        array
140    }
141
142    /// Returns the number of public keys in the set.
143    #[must_use]
144    pub const fn len(&self) -> usize {
145        self.0.len()
146    }
147
148    /// Returns `true` if the set is empty.
149    #[must_use]
150    pub const fn is_empty(&self) -> bool {
151        self.0.is_empty()
152    }
153
154    /// Returns the public key at the given index.
155    ///
156    /// It will return `None` if the index is out of bounds, even if it's less than
157    /// `MAX_AUTHENTICATOR_KEYS` but the key is not initialized.
158    #[must_use]
159    pub fn get(&self, index: usize) -> Option<&EdDSAPublicKey> {
160        self.0.get(index).and_then(Option::as_ref)
161    }
162
163    /// Sets a new public key at the given index if it's within bounds of the initialized set.
164    ///
165    /// # Errors
166    /// Returns an error if the index is out of bounds.
167    pub fn try_set_at_index(
168        &mut self,
169        index: usize,
170        pubkey: EdDSAPublicKey,
171    ) -> Result<(), PrimitiveError> {
172        if index >= self.len() || index >= MAX_AUTHENTICATOR_KEYS {
173            return Err(PrimitiveError::OutOfBounds);
174        }
175        self.0[index] = Some(pubkey);
176        Ok(())
177    }
178
179    /// Clears the public key at the given index while preserving slot position.
180    ///
181    /// # Errors
182    /// Returns an error if the index is out of bounds.
183    pub fn try_clear_at_index(&mut self, index: usize) -> Result<(), PrimitiveError> {
184        if index >= self.len() || index >= MAX_AUTHENTICATOR_KEYS {
185            return Err(PrimitiveError::OutOfBounds);
186        }
187        self.0[index] = None;
188        Ok(())
189    }
190
191    /// Pushes a new public key onto the set.
192    ///
193    /// # Errors
194    /// Returns an error if the set is full.
195    pub fn try_push(&mut self, pubkey: EdDSAPublicKey) -> Result<(), PrimitiveError> {
196        self.0
197            .try_push(Some(pubkey))
198            .map_err(|_| PrimitiveError::OutOfBounds)
199    }
200
201    /// Inserts a public key into the first empty slot, or appends it if no empty slot exists.
202    ///
203    /// Returns the slot index that now contains the inserted key.
204    ///
205    /// # Errors
206    /// Returns an error if the set is already at capacity.
207    pub fn insert_or_reuse(
208        &mut self,
209        new_authenticator_pubkey: EdDSAPublicKey,
210    ) -> Result<usize, PrimitiveError> {
211        if let Some(index) = self.iter().position(Option::is_none) {
212            self.try_set_at_index(index, new_authenticator_pubkey)?;
213            Ok(index)
214        } else {
215            self.try_push(new_authenticator_pubkey)?;
216            Ok(self.len() - 1)
217        }
218    }
219
220    /// Computes the Poseidon2 leaf hash commitment for this key set as stored in the `WorldIDRegistry`.
221    ///
222    /// # Panics
223    /// Panics if the domain separator constant cannot be converted into an `Fq`.
224    #[must_use]
225    pub fn leaf_hash(&self) -> Fq {
226        let mut input = [Fq::ZERO; 16];
227
228        input[0] =
229            Fq::from_str("105702839725298824521994315").expect("domain separator fits in field");
230
231        let pk_array = self.as_affine_array();
232        for i in 0..MAX_AUTHENTICATOR_KEYS {
233            input[i * 2 + 1] = pk_array[i].x;
234            input[i * 2 + 2] = pk_array[i].y;
235        }
236
237        poseidon2::bn254::t16::permutation(&input)[1]
238    }
239}
240
241impl Deref for AuthenticatorPublicKeySet {
242    type Target = ArrayVec<Option<EdDSAPublicKey>, MAX_AUTHENTICATOR_KEYS>;
243    fn deref(&self) -> &Self::Target {
244        &self.0
245    }
246}
247
248impl DerefMut for AuthenticatorPublicKeySet {
249    fn deref_mut(&mut self) -> &mut Self::Target {
250        &mut self.0
251    }
252}
253
254#[cfg(test)]
255mod tests {
256    use super::*;
257    use crate::{MAX_AUTHENTICATOR_KEYS, Signer};
258    use ark_babyjubjub::EdwardsAffine;
259    use ark_serialize::CanonicalSerialize as _;
260
261    fn create_test_pubkey() -> EdDSAPublicKey {
262        EdDSAPublicKey {
263            pk: EdwardsAffine::default(),
264        }
265    }
266
267    fn test_pubkey(seed_byte: u8) -> EdDSAPublicKey {
268        Signer::from_seed_bytes(&[seed_byte; 32])
269            .unwrap()
270            .offchain_signer_pubkey()
271    }
272
273    fn encoded_test_pubkey(seed_byte: u8) -> U256 {
274        let mut compressed = Vec::new();
275        test_pubkey(seed_byte)
276            .pk
277            .serialize_compressed(&mut compressed)
278            .unwrap();
279        U256::from_le_slice(&compressed)
280    }
281
282    #[test]
283    fn test_try_set_at_index_within_bounds() {
284        let mut key_set = AuthenticatorPublicKeySet::default();
285        let pubkey = create_test_pubkey();
286
287        let result = key_set.try_set_at_index(0, pubkey.clone());
288        assert!(result.is_err(), "Should not panic, should return error");
289
290        key_set.try_push(pubkey.clone()).unwrap();
291        key_set.try_set_at_index(0, pubkey).unwrap();
292    }
293
294    #[test]
295    fn test_try_set_at_index_at_length() {
296        let mut key_set = AuthenticatorPublicKeySet::default();
297        let pubkey = create_test_pubkey();
298
299        key_set.try_push(pubkey.clone()).unwrap();
300
301        let result = key_set.try_set_at_index(1, pubkey);
302        assert!(result.is_err(), "Should not panic when index equals length");
303    }
304
305    #[test]
306    fn test_try_set_at_index_out_of_bounds() {
307        let mut key_set = AuthenticatorPublicKeySet::default();
308        let pubkey = create_test_pubkey();
309
310        let result = key_set.try_set_at_index(MAX_AUTHENTICATOR_KEYS, pubkey.clone());
311        assert!(
312            result.is_err(),
313            "Should not panic when index >= MAX_AUTHENTICATOR_KEYS"
314        );
315
316        let result = key_set.try_set_at_index(MAX_AUTHENTICATOR_KEYS + 1, pubkey);
317        assert!(
318            result.is_err(),
319            "Should not panic when index > MAX_AUTHENTICATOR_KEYS"
320        );
321    }
322
323    #[test]
324    fn test_try_push_within_capacity() {
325        let mut key_set = AuthenticatorPublicKeySet::default();
326        let pubkey = create_test_pubkey();
327
328        for i in 0..MAX_AUTHENTICATOR_KEYS {
329            let result = key_set.try_push(pubkey.clone());
330            assert!(
331                result.is_ok(),
332                "Should not panic when pushing element {} of {}",
333                i + 1,
334                MAX_AUTHENTICATOR_KEYS
335            );
336        }
337
338        assert_eq!(key_set.len(), MAX_AUTHENTICATOR_KEYS);
339
340        let result = key_set.try_push(pubkey);
341        assert!(result.is_err());
342    }
343
344    #[test]
345    fn test_as_affine_array_empty_set() {
346        let key_set = AuthenticatorPublicKeySet::default();
347        let array = key_set.as_affine_array();
348
349        assert_eq!(array.len(), MAX_AUTHENTICATOR_KEYS);
350        for affine in &array {
351            assert_eq!(*affine, EdwardsAffine::default());
352        }
353    }
354
355    #[test]
356    fn test_as_affine_array_partial_set() {
357        let mut key_set = AuthenticatorPublicKeySet::default();
358        let pubkey = create_test_pubkey();
359
360        for _ in 0..3 {
361            key_set.try_push(pubkey.clone()).unwrap();
362        }
363
364        let array = key_set.as_affine_array();
365
366        assert_eq!(array.len(), MAX_AUTHENTICATOR_KEYS);
367
368        for item in array.iter().take(3) {
369            assert_eq!(item, &pubkey.pk);
370        }
371
372        for item in array.iter().take(MAX_AUTHENTICATOR_KEYS).skip(3) {
373            assert_eq!(item, &EdwardsAffine::default());
374        }
375    }
376
377    #[test]
378    fn test_decode_sparse_pubkeys_trims_trailing_empty_slots() {
379        let mut encoded_pubkeys = vec![Some(encoded_test_pubkey(1)), Some(encoded_test_pubkey(2))];
380        encoded_pubkeys.extend(vec![None; MAX_AUTHENTICATOR_KEYS + 5]);
381
382        let key_set =
383            AuthenticatorPublicKeySet::from_sparse_encoded_pubkeys(encoded_pubkeys).unwrap();
384
385        assert_eq!(key_set.len(), 2);
386        assert_eq!(key_set[0].as_ref().unwrap().pk, test_pubkey(1).pk);
387        assert_eq!(key_set[1].as_ref().unwrap().pk, test_pubkey(2).pk);
388    }
389
390    #[test]
391    fn test_decode_sparse_pubkeys_preserves_interior_holes() {
392        let key_set = AuthenticatorPublicKeySet::from_sparse_encoded_pubkeys(vec![
393            Some(encoded_test_pubkey(1)),
394            None,
395            Some(encoded_test_pubkey(2)),
396        ])
397        .unwrap();
398
399        assert_eq!(key_set.len(), 3);
400        assert_eq!(key_set[0].as_ref().unwrap().pk, test_pubkey(1).pk);
401        assert_eq!(key_set[1], None);
402        assert_eq!(key_set[2].as_ref().unwrap().pk, test_pubkey(2).pk);
403    }
404
405    #[test]
406    fn test_decode_sparse_pubkeys_rejects_used_slot_beyond_max() {
407        let mut encoded_pubkeys = vec![None; MAX_AUTHENTICATOR_KEYS + 1];
408        encoded_pubkeys[MAX_AUTHENTICATOR_KEYS] = Some(encoded_test_pubkey(1));
409
410        let error =
411            AuthenticatorPublicKeySet::from_sparse_encoded_pubkeys(encoded_pubkeys).unwrap_err();
412        assert!(matches!(
413            error,
414            SparseAuthenticatorPubkeysError::SlotOutOfBounds {
415                slot_index,
416                max_supported_slot
417            } if slot_index == MAX_AUTHENTICATOR_KEYS && max_supported_slot == MAX_AUTHENTICATOR_KEYS - 1
418        ));
419    }
420}