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

junobuild_shared/segments/
controllers.rs
