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
15pub const MAX_AUTHENTICATOR_KEYS: usize = 7;
21
22#[derive(Debug, thiserror::Error, Clone, PartialEq, Eq)]
24pub enum SparseAuthenticatorPubkeysError {
25 #[error(
27 "invalid authenticator pubkey slot {slot_index}; max supported slot is {max_supported_slot}"
28 )]
29 SlotOutOfBounds {
30 slot_index: usize,
32 max_supported_slot: usize,
34 },
35 #[error("invalid authenticator public key at slot {slot_index}: {reason}")]
37 InvalidCompressedPubkey {
38 slot_index: usize,
40 reason: String,
42 },
43}
44
45#[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 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 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 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 #[must_use]
144 pub const fn len(&self) -> usize {
145 self.0.len()
146 }
147
148 #[must_use]
150 pub const fn is_empty(&self) -> bool {
151 self.0.is_empty()
152 }
153
154 #[must_use]
159 pub fn get(&self, index: usize) -> Option<&EdDSAPublicKey> {
160 self.0.get(index).and_then(Option::as_ref)
161 }
162
163 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 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 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 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 #[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}