Skip to main content

junobuild_shared/segments/
access_keys.rs

1use crate::constants::internal::REVOKED_CONTROLLERS;
2use crate::constants::shared::{MAX_NUMBER_OF_ACCESS_KEYS, MAX_NUMBER_OF_ADMIN_CONTROLLERS};
3use crate::env::{CONSOLE, OBSERVATORY};
4use crate::errors::{
5    JUNO_ERROR_CONTROLLERS_ADMIN_NO_EXPIRY, JUNO_ERROR_CONTROLLERS_ANONYMOUS_NOT_ALLOWED,
6    JUNO_ERROR_CONTROLLERS_EXPIRY_IN_PAST, JUNO_ERROR_CONTROLLERS_MAX_NUMBER,
7    JUNO_ERROR_CONTROLLERS_REVOKED_NOT_ALLOWED,
8};
9use crate::ic::api::{id, is_canister_controller, time};
10use crate::types::interface::SetAccessKey;
11use crate::types::state::{AccessKey, AccessKeyId, AccessKeyScope, AccessKeys};
12use crate::utils::{principal_anonymous, principal_equal, principal_not_anonymous};
13use candid::Principal;
14use std::collections::HashMap;
15
16/// Initializes a set of access keys with default administrative scope.
17///
18/// # Arguments
19/// - `new_access_keys`: Slice of `AccessKeyId` representing the new access keys to be initialized.
20///
21/// # Returns
22/// A `AccessKeys` collection populated with the specified new access keys.
23pub fn init_admin_access_keys(new_access_keys: &[AccessKeyId]) -> AccessKeys {
24    let mut access_keys: AccessKeys = AccessKeys::new();
25
26    let access_key_data: SetAccessKey = SetAccessKey {
27        metadata: HashMap::new(),
28        expires_at: None,
29        scope: AccessKeyScope::Admin,
30        kind: None,
31    };
32
33    set_access_keys(new_access_keys, &access_key_data, &mut access_keys);
34
35    access_keys
36}
37
38/// Sets or updates access keys with specified data.
39///
40/// # Arguments
41/// - `new_access_key_ids`: Slice of `AccessKeyId` for the access keys to be set or updated.
42/// - `access_key_data`: `SetController` data to apply to the access keys.
43/// - `access_keys`: Mutable reference to the current set of access keys to update.
44pub fn set_access_keys(
45    new_access_key_ids: &[AccessKeyId],
46    access_key_data: &SetAccessKey,
47    access_keys: &mut AccessKeys,
48) {
49    for access_key_id in new_access_key_ids {
50        let existing_access_key = access_keys.get(access_key_id);
51
52        let now = time();
53
54        let created_at: u64 = match existing_access_key {
55            None => now,
56            Some(existing_access_key) => existing_access_key.created_at,
57        };
58
59        let updated_at: u64 = now;
60
61        let access_key: AccessKey = AccessKey {
62            metadata: access_key_data.metadata.clone(),
63            created_at,
64            updated_at,
65            expires_at: access_key_data.expires_at,
66            scope: access_key_data.scope.clone(),
67            kind: access_key_data.kind.clone(),
68        };
69
70        access_keys.insert(*access_key_id, access_key);
71    }
72}
73
74/// Removes specified access keys from the set.
75///
76/// # Arguments
77/// - `remove_access_key_ids`: Slice of `AccessKeyId` for the access keys to be removed.
78/// - `access_keys`: Mutable reference to the current set of access keys to update.
79pub fn delete_access_keys(remove_access_key_ids: &[AccessKeyId], access_keys: &mut AccessKeys) {
80    for id in remove_access_key_ids {
81        access_keys.remove(id);
82    }
83}
84
85/// Checks if an id is a non-expired access key with admin or write scope (permissions).
86///
87/// # Arguments
88/// - `id`: `Principal` of e.g. the caller.
89/// - `access_keys`: Reference to the current set of access keys.
90///
91/// # Returns
92/// `true` if the id is an access key (not anonymous, calling itself or one of the known write or admin access keys), otherwise `false`.
93pub fn is_write_access_key(id: Principal, access_keys: &AccessKeys) -> bool {
94    principal_not_anonymous(id)
95        && (is_self(id)
96            || access_keys
97                .iter()
98                .any(|(&access_key_id, access_key)| match access_key.scope {
99                    AccessKeyScope::Submit => false,
100                    _ => {
101                        principal_equal(access_key_id, id) && is_access_key_not_expired(access_key)
102                    }
103                }))
104}
105
106/// Checks if an id is a non-expired access key regardless of scope (admin, write, or submit).
107///
108/// # Arguments
109/// - `id`: `Principal` e.g. of the caller.
110/// - `access_keys`: Reference to the current set of access keys.
111///
112/// # Returns
113/// `true` if the id is an access key (not anonymous, calling itself or one of the known controllers), otherwise `false`.
114pub fn is_valid_access_key(id: Principal, access_keys: &AccessKeys) -> bool {
115    principal_not_anonymous(id)
116        && (is_self(id)
117            || access_keys.iter().any(|(&access_key_id, access_key)| {
118                principal_equal(access_key_id, id) && is_access_key_not_expired(access_key)
119            }))
120}
121
122/// Checks if an access key has not expired.
123///
124/// Admin access keys never expire. Other access keys are considered not expired if:
125/// - They have no expiration date set, or
126/// - Their expiration date is in the future
127///
128/// # Arguments
129/// - `access_key`: The access key to check
130///
131/// # Returns
132/// `true` if the access key has not expired, `false` otherwise.
133fn is_access_key_not_expired(access_key: &AccessKey) -> bool {
134    !is_access_key_expired(access_key)
135}
136
137/// Checks if an access key has expired.
138///
139/// Admin access keys never expire (because those are controllers on the Internet computer).
140/// Other keys are considered expired if:
141/// - They have an expiration date set, and
142/// - That expiration date is in the past
143///
144/// # Arguments
145/// - `access_key`: The access key to check
146///
147/// # Returns
148/// `true` if the access key has expired, `false` otherwise.
149fn is_access_key_expired(access_key: &AccessKey) -> bool {
150    // Admin controller cannot expire
151    if matches!(access_key.scope, AccessKeyScope::Admin) {
152        return false;
153    }
154
155    access_key
156        .expires_at
157        .is_some_and(|expires_at| expires_at < time())
158}
159
160/// Checks if an id is an admin access key and a controller has known by the Internet Computer.
161///
162/// # Arguments
163/// - `id`: `Principal` e.g. of the caller.
164/// - `access_keys`: Reference to the current set of access_keys.
165///
166/// # Returns
167/// `true` if the id is an admin access key and a controller, otherwise `false`.
168pub fn is_admin_controller(id: Principal, access_keys: &AccessKeys) -> bool {
169    is_canister_controller(&id)
170        && principal_not_anonymous(id)
171        && access_keys
172            .iter()
173            .any(|(&access_key_id, access_key)| match access_key.scope {
174                AccessKeyScope::Admin => principal_equal(access_key_id, id),
175                _ => false,
176            })
177}
178
179/// Converts the access keys set into a vector of access key IDs.
180///
181/// # Arguments
182/// - `access_keys`: Reference to the current set of access keys.
183///
184/// # Returns
185/// A vector of `AccessKeyId`.
186pub fn into_access_key_ids(access_keys: &AccessKeys) -> Vec<AccessKeyId> {
187    access_keys
188        .clone()
189        .into_keys()
190        .collect::<Vec<AccessKeyId>>()
191}
192
193/// Asserts that the number of access keys does not exceed the maximum allowed.
194///
195/// # Arguments
196/// - `current_access_keys`: Reference to the current set of access keys.
197/// - `access_keys_ids`: Slice of `AccessKeyId` representing the access keys to be added.
198/// - `scope`: The scope of the new access key to be added.
199/// - `max_access_keys`: Optional custom maximum number of allowed access keys.
200///
201/// # Returns
202/// `Ok(())` if the operation is successful, or `Err(String)` if the maximum is exceeded.
203pub fn assert_max_number_of_access_keys(
204    current_access_keys: &AccessKeys,
205    access_keys_ids: &[AccessKeyId],
206    scope: &AccessKeyScope,
207    max_access_keys: Option<usize>,
208) -> Result<(), String> {
209    let filtered_access_keys = filter_access_keys(current_access_keys, scope);
210
211    let max_length = match scope {
212        AccessKeyScope::Admin => max_access_keys.unwrap_or(MAX_NUMBER_OF_ADMIN_CONTROLLERS),
213        _ => max_access_keys.unwrap_or(MAX_NUMBER_OF_ACCESS_KEYS),
214    };
215
216    assert_access_keys_length(&filtered_access_keys, access_keys_ids, max_length)
217}
218
219/// Asserts that the number of access keys does not exceed the maximum allowed.
220///
221/// # Arguments
222/// - `current_access_keys`: Reference to the current set of access keys.
223/// - `access_keys_ids`: Slice of `AccessKeyId` representing the access key to be added.
224/// - `max_access_keys`: Maximum number of allowed access keys.
225///
226/// # Returns
227/// `Ok(())` if the operation is successful, or `Err(String)` if the maximum is exceeded.
228fn assert_access_keys_length(
229    current_access_keys: &AccessKeys,
230    access_keys_ids: &[AccessKeyId],
231    max_access_keys: usize,
232) -> Result<(), String> {
233    let current_access_key_ids = into_access_key_ids(current_access_keys);
234
235    let new_access_key_ids = access_keys_ids.iter().copied().filter(|id| {
236        !current_access_key_ids
237            .iter()
238            .any(|current_id| current_id == id)
239    });
240
241    if current_access_key_ids.len() + new_access_key_ids.count() > max_access_keys {
242        return Err(format!(
243            "{JUNO_ERROR_CONTROLLERS_MAX_NUMBER} ({max_access_keys})"
244        ));
245    }
246
247    Ok(())
248}
249
250/// Asserts that the access key IDs are not anonymous and not revoked.
251///
252/// # Arguments
253/// - `access_keys_ids`: Slice of `AccessKeyId` to validate.
254///
255/// # Returns
256/// `Ok(())` if no anonymous and no revoked IDs are present, or `Err(String)` if any are found.
257pub fn assert_controllers(access_keys_ids: &[AccessKeyId]) -> Result<(), String> {
258    assert_no_anonymous_access_key(access_keys_ids)?;
259    assert_no_revoked_controller(access_keys_ids)?;
260
261    Ok(())
262}
263
264/// Asserts that no access key IDs are anonymous.
265///
266/// # Arguments
267/// - `access_keys_ids`: Slice of `AccessKeyId` to validate.
268///
269/// # Returns
270/// `Ok(())` if no anonymous IDs are present, or `Err(String)` if any are found.
271fn assert_no_anonymous_access_key(access_keys_ids: &[AccessKeyId]) -> Result<(), String> {
272    let has_anonymous = access_keys_ids
273        .iter()
274        .any(|controller_id| principal_anonymous(*controller_id));
275
276    match has_anonymous {
277        true => Err(JUNO_ERROR_CONTROLLERS_ANONYMOUS_NOT_ALLOWED.to_string()),
278        false => Ok(()),
279    }
280}
281
282/// Asserts that no access key IDs are revoked for security reason.
283///
284/// # Arguments
285/// - `access_keys_ids`: Slice of `AccessKeyId` to validate.
286///
287/// # Returns
288/// `Ok(())` if no revoked IDs are present, or `Err(String)` if any are found.
289fn assert_no_revoked_controller(access_keys_ids: &[AccessKeyId]) -> Result<(), String> {
290    // We treat revoked controllers as anonymous controllers.
291    let has_revoked = access_keys_ids.iter().any(controller_revoked);
292
293    match has_revoked {
294        true => Err(JUNO_ERROR_CONTROLLERS_REVOKED_NOT_ALLOWED.to_string()),
295        false => Ok(()),
296    }
297}
298
299/// Validates access key expiration settings.
300///
301/// Ensures that:
302/// - Admin access keys do not have an expiration date set
303/// - If an expiration is set, it's not in the past
304///
305/// # Arguments
306/// - `access_key`: The access key configuration to validate
307///
308/// # Returns
309/// `Ok(())` if validation passes, or `Err(String)` with error message if validation fails
310pub fn assert_access_key_expiration(access_key: &SetAccessKey) -> Result<(), String> {
311    if matches!(access_key.scope, AccessKeyScope::Admin) && access_key.expires_at.is_some() {
312        return Err(JUNO_ERROR_CONTROLLERS_ADMIN_NO_EXPIRY.to_string());
313    }
314
315    if let Some(expires_at) = access_key.expires_at {
316        if expires_at < time() {
317            return Err(JUNO_ERROR_CONTROLLERS_EXPIRY_IN_PAST.to_string());
318        }
319    }
320
321    Ok(())
322}
323
324/// Checks if the id is the console.
325///
326/// # Arguments
327/// - `id`: `Principal` e.g. of the caller.
328///
329/// # Returns
330/// `true` if the id matches the console's principal, otherwise `false`.
331pub fn is_console(id: Principal) -> bool {
332    let console = Principal::from_text(CONSOLE).unwrap();
333
334    principal_equal(id, console)
335}
336
337/// Checks if the id is the observatory.
338///
339/// # Arguments
340/// - `id`: `Principal` e.g. of the caller.
341///
342/// # Returns
343/// `true` if the id matches the observatory's principal, otherwise `false`.
344pub fn is_observatory(id: Principal) -> bool {
345    let observatory = Principal::from_text(OBSERVATORY).unwrap();
346
347    principal_equal(id, observatory)
348}
349
350/// Checks if an id (e.g. a caller, a user or a principal) is the canister itself.
351///
352/// # Arguments
353/// - `id`: `Principal` of the id.
354///
355/// # Returns
356/// `true` if the id is calling itself, if the canister is the caller, otherwise `false`.
357pub fn is_self(provided_id: Principal) -> bool {
358    let itself = id();
359
360    principal_equal(provided_id, itself)
361}
362
363/// Filters the set of access keys, returning only those with administrative scope.
364///
365/// # Arguments
366/// - `access_keys`: Reference to the current set of access keys.
367///
368/// # Returns
369/// An `AccessKeys` collection containing only admin access keys.
370pub fn filter_admin_access_keys(access_keys: &AccessKeys) -> AccessKeys {
371    filter_access_keys(access_keys, &AccessKeyScope::Admin)
372}
373
374/// Filters the set of access keys, returning only those matching the provided scope.
375///
376/// # Arguments
377/// - `access_keys`: Reference to the current set of access keys.
378/// - `scope`: The expected scope.
379///
380/// # Returns
381/// An `AccessKeys` collection containing only matching access keys.
382fn filter_access_keys(controllers: &AccessKeys, scope: &AccessKeyScope) -> AccessKeys {
383    #[allow(clippy::match_like_matches_macro)]
384    controllers
385        .clone()
386        .into_iter()
387        .filter(|(_, controller)| match scope {
388            AccessKeyScope::Write => matches!(controller.scope, AccessKeyScope::Write),
389            AccessKeyScope::Admin => matches!(controller.scope, AccessKeyScope::Admin),
390            AccessKeyScope::Submit => matches!(controller.scope, AccessKeyScope::Submit),
391        })
392        .collect()
393}
394
395fn controller_revoked(access_key_id: &AccessKeyId) -> bool {
396    REVOKED_CONTROLLERS.iter().any(|revoked_controller_id| {
397        principal_equal(
398            Principal::from_text(revoked_controller_id).unwrap(),
399            *access_key_id,
400        )
401    })
402}
403
404#[cfg(test)]
405mod tests {
406    use super::*;
407    use crate::types::state::{AccessKey, AccessKeyKind, AccessKeyScope, AccessKeys};
408    use candid::Principal;
409    use std::collections::HashMap;
410
411    fn mock_time() -> u64 {
412        1_000_000_000_000
413    }
414
415    fn test_principal(id: u8) -> Principal {
416        Principal::from_slice(&[id])
417    }
418
419    fn create_access_key(
420        scope: AccessKeyScope,
421        expires_at: Option<u64>,
422        kind: Option<AccessKeyKind>,
423    ) -> AccessKey {
424        AccessKey {
425            metadata: HashMap::new(),
426            created_at: mock_time(),
427            updated_at: mock_time(),
428            expires_at,
429            scope,
430            kind,
431        }
432    }
433
434    #[test]
435    fn test_is_access_key_expired_admin_never_expires() {
436        let admin = create_access_key(AccessKeyScope::Admin, Some(mock_time() - 1000), None);
437        assert!(!is_access_key_expired(&admin));
438    }
439
440    #[test]
441    fn test_is_access_key_expired_no_expiration() {
442        let access_key = create_access_key(AccessKeyScope::Write, None, None);
443        assert!(!is_access_key_expired(&access_key));
444    }
445
446    #[test]
447    fn test_is_access_key_expired_future_expiration() {
448        let access_key = create_access_key(AccessKeyScope::Write, Some(time() + 1_000_000), None);
449        assert!(!is_access_key_expired(&access_key));
450    }
451
452    #[test]
453    fn test_is_access_key_not_expired() {
454        let admin = create_access_key(AccessKeyScope::Admin, Some(time() - 1000), None);
455        assert!(is_access_key_not_expired(&admin));
456
457        let expired = create_access_key(AccessKeyScope::Write, Some(time() - 1), None);
458        assert!(!is_access_key_not_expired(&expired));
459
460        let valid = create_access_key(AccessKeyScope::Write, Some(time() + 1000), None);
461        assert!(is_access_key_not_expired(&valid));
462    }
463
464    #[test]
465    fn test_is_access_key_expired_past_expiration() {
466        let access_key = create_access_key(AccessKeyScope::Write, Some(mock_time() - 1), None);
467        assert!(is_access_key_expired(&access_key));
468    }
469
470    #[test]
471    fn test_access_key_can_write_anonymous_rejected() {
472        let access_keys = AccessKeys::new();
473        assert!(!is_write_access_key(Principal::anonymous(), &access_keys));
474    }
475
476    #[test]
477    fn test_access_key_can_write_self_allowed() {
478        let access_keys = AccessKeys::new();
479        let self_principal = id();
480        assert!(is_write_access_key(self_principal, &access_keys));
481    }
482
483    #[test]
484    fn test_access_key_can_write_admin_allowed() {
485        let mut access_keys = AccessKeys::new();
486        let admin_principal = test_principal(1);
487
488        access_keys.insert(
489            admin_principal,
490            create_access_key(AccessKeyScope::Admin, None, None),
491        );
492
493        assert!(is_write_access_key(admin_principal, &access_keys));
494    }
495
496    #[test]
497    fn test_access_key_can_write_write_scope_allowed() {
498        let mut access_keys = AccessKeys::new();
499        let write_principal = test_principal(2);
500
501        access_keys.insert(
502            write_principal,
503            create_access_key(AccessKeyScope::Write, None, None),
504        );
505
506        assert!(is_write_access_key(write_principal, &access_keys));
507    }
508
509    #[test]
510    fn test_access_key_can_write_submit_rejected() {
511        let mut access_keys = AccessKeys::new();
512        let submit_principal = test_principal(3);
513
514        access_keys.insert(
515            submit_principal,
516            create_access_key(AccessKeyScope::Submit, None, None),
517        );
518
519        assert!(!is_write_access_key(submit_principal, &access_keys));
520    }
521
522    #[test]
523    fn test_access_key_can_write_expired_rejected() {
524        let mut access_keys = AccessKeys::new();
525        let expired_principal = test_principal(4);
526
527        access_keys.insert(
528            expired_principal,
529            create_access_key(AccessKeyScope::Write, Some(mock_time() - 1), None),
530        );
531
532        assert!(!is_write_access_key(expired_principal, &access_keys));
533    }
534
535    #[test]
536    fn test_is_valid_access_key_anonymous_rejected() {
537        let access_keys = AccessKeys::new();
538        assert!(!is_valid_access_key(Principal::anonymous(), &access_keys));
539    }
540
541    #[test]
542    fn test_is_valid_access_key_self_allowed() {
543        let access_keys = AccessKeys::new();
544        let self_principal = id();
545        assert!(is_valid_access_key(self_principal, &access_keys));
546    }
547
548    #[test]
549    fn test_is_valid_access_key_all_scopes_allowed() {
550        let mut access_keys = AccessKeys::new();
551
552        let admin = test_principal(1);
553        let write = test_principal(2);
554        let submit = test_principal(3);
555
556        access_keys.insert(admin, create_access_key(AccessKeyScope::Admin, None, None));
557        access_keys.insert(write, create_access_key(AccessKeyScope::Write, None, None));
558        access_keys.insert(
559            submit,
560            create_access_key(AccessKeyScope::Submit, None, None),
561        );
562
563        assert!(is_valid_access_key(admin, &access_keys));
564        assert!(is_valid_access_key(write, &access_keys));
565        assert!(is_valid_access_key(submit, &access_keys));
566    }
567
568    #[test]
569    fn test_is_valid_access_key_expired_rejected() {
570        let mut access_keys = AccessKeys::new();
571        let expired_principal = test_principal(4);
572
573        access_keys.insert(
574            expired_principal,
575            create_access_key(AccessKeyScope::Write, Some(mock_time() - 1), None),
576        );
577
578        assert!(!is_valid_access_key(expired_principal, &access_keys));
579    }
580
581    #[test]
582    fn test_is_valid_access_key_admin_expired_still_valid() {
583        let mut access_keys = AccessKeys::new();
584        let admin_principal = test_principal(5);
585
586        // Admin with past expiration should still be valid (admins never expire)
587        access_keys.insert(
588            admin_principal,
589            create_access_key(AccessKeyScope::Admin, Some(mock_time() - 1000), None),
590        );
591
592        assert!(is_valid_access_key(admin_principal, &access_keys));
593    }
594
595    #[test]
596    fn test_init_admin_access_keys() {
597        let principals = vec![test_principal(1), test_principal(2)];
598        let access_keys = init_admin_access_keys(&principals);
599
600        assert_eq!(access_keys.len(), 2);
601
602        for principal in principals {
603            let access_key = access_keys.get(&principal).unwrap();
604            assert!(matches!(access_key.scope, AccessKeyScope::Admin));
605            assert!(access_key.expires_at.is_none());
606            assert!(access_key.kind.is_none());
607        }
608    }
609
610    #[test]
611    fn test_set_access_keys_new() {
612        let mut access_keys = AccessKeys::new();
613        let principals = vec![test_principal(1)];
614
615        let access_key_data = SetAccessKey {
616            metadata: HashMap::new(),
617            expires_at: Some(mock_time() + 1000),
618            scope: AccessKeyScope::Write,
619            kind: Some(AccessKeyKind::Automation),
620        };
621
622        set_access_keys(&principals, &access_key_data, &mut access_keys);
623
624        assert_eq!(access_keys.len(), 1);
625        let access_key = access_keys.get(&principals[0]).unwrap();
626        assert!(matches!(access_key.scope, AccessKeyScope::Write));
627        assert_eq!(access_key.expires_at, Some(mock_time() + 1000));
628        assert!(matches!(access_key.kind, Some(AccessKeyKind::Automation)));
629    }
630
631    #[test]
632    fn test_set_access_keys_update_preserves_created_at() {
633        let mut access_keys = AccessKeys::new();
634        let principal = test_principal(1);
635
636        // First insert
637        let initial_data = SetAccessKey {
638            metadata: HashMap::new(),
639            expires_at: None,
640            scope: AccessKeyScope::Write,
641            kind: None,
642        };
643        set_access_keys(&[principal], &initial_data, &mut access_keys);
644        let original_created_at = access_keys.get(&principal).unwrap().created_at;
645
646        // Update
647        let update_data = SetAccessKey {
648            metadata: HashMap::new(),
649            expires_at: Some(mock_time() + 1000),
650            scope: AccessKeyScope::Admin,
651            kind: Some(AccessKeyKind::Automation),
652        };
653        set_access_keys(&[principal], &update_data, &mut access_keys);
654
655        let updated = access_keys.get(&principal).unwrap();
656        assert_eq!(updated.created_at, original_created_at);
657        assert!(matches!(updated.scope, AccessKeyScope::Admin));
658    }
659
660    #[test]
661    fn test_delete_access_keys() {
662        let mut access_keys = AccessKeys::new();
663        let principals = vec![test_principal(1), test_principal(2), test_principal(3)];
664
665        for principal in &principals {
666            access_keys.insert(
667                *principal,
668                create_access_key(AccessKeyScope::Write, None, None),
669            );
670        }
671
672        delete_access_keys(&principals[0..2], &mut access_keys);
673
674        assert_eq!(access_keys.len(), 1);
675        assert!(access_keys.contains_key(&principals[2]));
676        assert!(!access_keys.contains_key(&principals[0]));
677        assert!(!access_keys.contains_key(&principals[1]));
678    }
679
680    #[test]
681    fn test_filter_admin_access_keys() {
682        let mut access_keys = AccessKeys::new();
683
684        access_keys.insert(
685            test_principal(1),
686            create_access_key(AccessKeyScope::Admin, None, None),
687        );
688        access_keys.insert(
689            test_principal(2),
690            create_access_key(AccessKeyScope::Write, None, None),
691        );
692        access_keys.insert(
693            test_principal(3),
694            create_access_key(AccessKeyScope::Admin, None, None),
695        );
696        access_keys.insert(
697            test_principal(4),
698            create_access_key(AccessKeyScope::Submit, None, None),
699        );
700
701        let admin_only = filter_admin_access_keys(&access_keys);
702
703        assert_eq!(admin_only.len(), 2);
704        assert!(admin_only.contains_key(&test_principal(1)));
705        assert!(admin_only.contains_key(&test_principal(3)));
706    }
707
708    #[test]
709    fn test_filter_access_keys_write() {
710        let mut access_keys = AccessKeys::new();
711
712        access_keys.insert(
713            test_principal(1),
714            create_access_key(AccessKeyScope::Admin, None, None),
715        );
716        access_keys.insert(
717            test_principal(2),
718            create_access_key(AccessKeyScope::Write, None, None),
719        );
720        access_keys.insert(
721            test_principal(3),
722            create_access_key(AccessKeyScope::Write, None, None),
723        );
724
725        let write_only = filter_access_keys(&access_keys, &AccessKeyScope::Write);
726
727        assert_eq!(write_only.len(), 2);
728        assert!(write_only.contains_key(&test_principal(2)));
729        assert!(write_only.contains_key(&test_principal(3)));
730    }
731
732    #[test]
733    fn test_filter_access_keys_submit() {
734        let mut access_keys = AccessKeys::new();
735
736        access_keys.insert(
737            test_principal(1),
738            create_access_key(AccessKeyScope::Submit, None, None),
739        );
740        access_keys.insert(
741            test_principal(2),
742            create_access_key(AccessKeyScope::Write, None, None),
743        );
744
745        let submit_only = filter_access_keys(&access_keys, &AccessKeyScope::Submit);
746
747        assert_eq!(submit_only.len(), 1);
748        assert!(submit_only.contains_key(&test_principal(1)));
749    }
750
751    #[test]
752    fn test_is_admin_controller() {
753        let mut access_keys = AccessKeys::new();
754        let admin_principal = id(); // Self canister in test
755
756        access_keys.insert(
757            admin_principal,
758            create_access_key(AccessKeyScope::Admin, None, None),
759        );
760
761        assert!(is_admin_controller(admin_principal, &access_keys));
762    }
763
764    #[test]
765    fn test_is_admin_controller_not_canister_access_key() {
766        let mut access_keys = AccessKeys::new();
767        let not_canister = test_principal(99);
768
769        access_keys.insert(
770            not_canister,
771            create_access_key(AccessKeyScope::Admin, None, None),
772        );
773
774        // Will fail because test_principal(99) is not the canister access_key
775        assert!(!is_admin_controller(not_canister, &access_keys));
776    }
777
778    #[test]
779    fn test_assert_expiration_access_key_admin_with_expiry_rejected() {
780        let access_key = SetAccessKey {
781            metadata: HashMap::new(),
782            expires_at: Some(time() + 1000),
783            scope: AccessKeyScope::Admin,
784            kind: None,
785        };
786
787        let result = assert_access_key_expiration(&access_key);
788        assert!(result.is_err());
789        assert_eq!(
790            result.unwrap_err(),
791            JUNO_ERROR_CONTROLLERS_ADMIN_NO_EXPIRY.to_string()
792        );
793    }
794
795    #[test]
796    fn test_assert_expiration_access_key_admin_without_expiry_allowed() {
797        let access_key = SetAccessKey {
798            metadata: HashMap::new(),
799            expires_at: None,
800            scope: AccessKeyScope::Admin,
801            kind: None,
802        };
803
804        let result = assert_access_key_expiration(&access_key);
805        assert!(result.is_ok());
806    }
807
808    #[test]
809    fn test_assert_expiration_access_key_past_expiry_rejected() {
810        let access_key = SetAccessKey {
811            metadata: HashMap::new(),
812            expires_at: Some(time() - 1000),
813            scope: AccessKeyScope::Write,
814            kind: None,
815        };
816
817        let result = assert_access_key_expiration(&access_key);
818        assert!(result.is_err());
819        assert_eq!(
820            result.unwrap_err(),
821            JUNO_ERROR_CONTROLLERS_EXPIRY_IN_PAST.to_string()
822        );
823    }
824
825    #[test]
826    fn test_assert_expiration_access_key_future_expiry_allowed() {
827        let access_key = SetAccessKey {
828            metadata: HashMap::new(),
829            expires_at: Some(time() + 1000),
830            scope: AccessKeyScope::Write,
831            kind: None,
832        };
833
834        let result = assert_access_key_expiration(&access_key);
835        assert!(result.is_ok());
836    }
837
838    #[test]
839    fn test_assert_expiration_access_key_no_expiry_allowed() {
840        let access_key = SetAccessKey {
841            metadata: HashMap::new(),
842            expires_at: None,
843            scope: AccessKeyScope::Write,
844            kind: None,
845        };
846
847        let result = assert_access_key_expiration(&access_key);
848        assert!(result.is_ok());
849    }
850
851    #[test]
852    fn test_assert_expiration_access_key_submit_scope_with_future_expiry() {
853        let access_key = SetAccessKey {
854            metadata: HashMap::new(),
855            expires_at: Some(time() + 1000),
856            scope: AccessKeyScope::Submit,
857            kind: None,
858        };
859
860        let result = assert_access_key_expiration(&access_key);
861        assert!(result.is_ok());
862    }
863
864    #[test]
865    fn test_assert_max_number_of_access_keys_admin_default() {
866        let mut access_keys = AccessKeys::new();
867
868        // Add MAX_NUMBER_OF_ADMIN_CONTROLLERS - 1 admin access_keys
869        for i in 0..(MAX_NUMBER_OF_ADMIN_CONTROLLERS - 1) {
870            access_keys.insert(
871                test_principal(i as u8),
872                create_access_key(AccessKeyScope::Admin, None, None),
873            );
874        }
875
876        // Adding one more should succeed
877        let new_access_keys = vec![test_principal(99)];
878        let result = assert_max_number_of_access_keys(
879            &access_keys,
880            &new_access_keys,
881            &AccessKeyScope::Admin,
882            None,
883        );
884        assert!(result.is_ok());
885
886        // Adding two more should fail
887        let new_access_keys = vec![test_principal(98), test_principal(99)];
888        let result = assert_max_number_of_access_keys(
889            &access_keys,
890            &new_access_keys,
891            &AccessKeyScope::Admin,
892            None,
893        );
894        assert!(result.is_err());
895    }
896
897    #[test]
898    fn test_assert_max_number_of_access_keys_write_default() {
899        let mut access_keys = AccessKeys::new();
900
901        // Add MAX_NUMBER_OF_ACCESS_KEYS - 1 write access_keys
902        for i in 0..(MAX_NUMBER_OF_ACCESS_KEYS - 1) {
903            access_keys.insert(
904                test_principal(i as u8),
905                create_access_key(AccessKeyScope::Write, None, None),
906            );
907        }
908
909        // Adding one more should succeed
910        let new_access_keys = vec![test_principal(255)];
911        let result = assert_max_number_of_access_keys(
912            &access_keys,
913            &new_access_keys,
914            &AccessKeyScope::Write,
915            None,
916        );
917        assert!(result.is_ok());
918
919        // Adding two more should fail
920        let new_access_keys = vec![test_principal(254), test_principal(255)];
921        let result = assert_max_number_of_access_keys(
922            &access_keys,
923            &new_access_keys,
924            &AccessKeyScope::Write,
925            None,
926        );
927        assert!(result.is_err());
928    }
929
930    #[test]
931    fn test_assert_max_number_of_access_keys_custom_limit() {
932        let mut access_keys = AccessKeys::new();
933
934        // Add 2 admin access_keys
935        for i in 0..2 {
936            access_keys.insert(
937                test_principal(i),
938                create_access_key(AccessKeyScope::Admin, None, None),
939            );
940        }
941
942        // With custom limit of 3, adding one more should succeed
943        let new_access_keys = vec![test_principal(3)];
944        let result = assert_max_number_of_access_keys(
945            &access_keys,
946            &new_access_keys,
947            &AccessKeyScope::Admin,
948            Some(3),
949        );
950        assert!(result.is_ok());
951
952        // Adding two more should fail
953        let new_access_keys = vec![test_principal(3), test_principal(4)];
954        let result = assert_max_number_of_access_keys(
955            &access_keys,
956            &new_access_keys,
957            &AccessKeyScope::Admin,
958            Some(3),
959        );
960        assert!(result.is_err());
961    }
962
963    #[test]
964    fn test_assert_max_number_of_access_keys_filters_by_scope() {
965        let mut access_keys = AccessKeys::new();
966
967        // Add admin access_keys
968        for i in 0..5 {
969            access_keys.insert(
970                test_principal(i),
971                create_access_key(AccessKeyScope::Admin, None, None),
972            );
973        }
974
975        // Add write access_keys
976        for i in 10..15 {
977            access_keys.insert(
978                test_principal(i),
979                create_access_key(AccessKeyScope::Write, None, None),
980            );
981        }
982
983        // Adding a write access_key should only count against write limit, not admin
984        let new_access_keys = vec![test_principal(20)];
985        let result = assert_max_number_of_access_keys(
986            &access_keys,
987            &new_access_keys,
988            &AccessKeyScope::Write,
989            Some(6), // 5 existing + 1 new = 6, should succeed
990        );
991        assert!(result.is_ok());
992
993        // Adding an admin access_key should only count against admin limit
994        let new_access_keys = vec![test_principal(21)];
995        let result = assert_max_number_of_access_keys(
996            &access_keys,
997            &new_access_keys,
998            &AccessKeyScope::Admin,
999            Some(6), // 5 existing + 1 new = 6, should succeed
1000        );
1001        assert!(result.is_ok());
1002    }
1003}