Skip to main content

cougr_core/accounts/
multi_device.rs

1use soroban_sdk::{contracttype, Address, BytesN, Env, Symbol, Vec};
2
3use super::device_storage::DeviceStorage;
4use super::error::AccountError;
5
6/// A registered device key with metadata.
7#[contracttype]
8#[derive(Clone, Debug, PartialEq, Eq)]
9pub struct DeviceKey {
10    /// Unique identifier for this device key.
11    pub key_id: BytesN<32>,
12    /// Human-readable device name (e.g., "phone", "laptop").
13    pub device_name: Symbol,
14    /// Ledger timestamp when the device was registered.
15    pub registered_at: u64,
16    /// Ledger timestamp of the last use.
17    pub last_used: u64,
18    /// Whether this device key is currently active.
19    pub is_active: bool,
20}
21
22/// Policy for multi-device management.
23#[contracttype]
24#[derive(Clone, Debug)]
25pub struct DevicePolicy {
26    /// Maximum number of devices that can be registered.
27    pub max_devices: u32,
28    /// Number of ledger slots of inactivity before auto-revoke.
29    /// Set to 0 to disable auto-revoke.
30    pub auto_revoke_after: u64,
31}
32
33/// Trait for account types that support multi-device key management.
34pub trait MultiDeviceProvider {
35    /// Register a new device key.
36    fn register_device(
37        &mut self,
38        env: &Env,
39        key_id: BytesN<32>,
40        device_name: Symbol,
41    ) -> Result<DeviceKey, AccountError>;
42
43    /// Revoke a device key by its ID.
44    fn revoke_device(&mut self, env: &Env, key_id: &BytesN<32>) -> Result<(), AccountError>;
45
46    /// List all registered device keys (active and inactive).
47    fn list_devices(&self, env: &Env) -> Vec<DeviceKey>;
48
49    /// Returns the number of active devices.
50    fn active_device_count(&self, env: &Env) -> usize;
51
52    /// Update the last_used timestamp for a device.
53    fn update_last_used(&mut self, env: &Env, key_id: &BytesN<32>) -> Result<(), AccountError>;
54
55    /// Set the device management policy.
56    fn set_policy(&mut self, env: &Env, policy: DevicePolicy);
57
58    /// Get the current device policy.
59    fn policy(&self, env: &Env) -> DevicePolicy;
60
61    /// Revoke devices that have been inactive beyond the policy's auto_revoke_after.
62    /// Returns the number of devices revoked.
63    fn cleanup_inactive(&mut self, env: &Env) -> u32;
64}
65
66/// Persistent implementation of multi-device management.
67///
68/// Device keys and policy are stored in Soroban persistent storage
69/// via [`DeviceStorage`], surviving across contract invocations.
70pub struct DeviceManager {
71    address: Address,
72}
73
74impl DeviceManager {
75    /// Create a new device manager with the given policy, persisting it.
76    pub fn new(address: Address, policy: DevicePolicy, env: &Env) -> Self {
77        DeviceStorage::store_policy(env, &address, &policy);
78        DeviceStorage::store_devices(env, &address, &Vec::new(env));
79        Self { address }
80    }
81
82    /// Load an existing device manager (policy must already be stored).
83    pub fn load(address: Address) -> Self {
84        Self { address }
85    }
86
87    /// Create a device manager with a default policy.
88    pub fn with_defaults(address: Address, env: &Env) -> Self {
89        let policy = DevicePolicy {
90            max_devices: 5,
91            auto_revoke_after: 0,
92        };
93        Self::new(address, policy, env)
94    }
95}
96
97impl MultiDeviceProvider for DeviceManager {
98    fn register_device(
99        &mut self,
100        env: &Env,
101        key_id: BytesN<32>,
102        device_name: Symbol,
103    ) -> Result<DeviceKey, AccountError> {
104        let policy =
105            DeviceStorage::load_policy(env, &self.address).ok_or(AccountError::StorageError)?;
106        let devices = DeviceStorage::load_devices(env, &self.address);
107
108        // Check device limit (active only)
109        let mut active_count: u32 = 0;
110        for i in 0..devices.len() {
111            if let Some(d) = devices.get(i) {
112                if d.is_active {
113                    active_count += 1;
114                }
115                // Check for duplicate active key_id
116                if d.key_id == key_id && d.is_active {
117                    return Err(AccountError::DeviceLimitReached);
118                }
119            }
120        }
121
122        if active_count >= policy.max_devices {
123            return Err(AccountError::DeviceLimitReached);
124        }
125
126        let now = env.ledger().timestamp();
127        let device = DeviceKey {
128            key_id,
129            device_name,
130            registered_at: now,
131            last_used: now,
132            is_active: true,
133        };
134
135        let mut new_devices = devices;
136        new_devices.push_back(device.clone());
137        DeviceStorage::store_devices(env, &self.address, &new_devices);
138        Ok(device)
139    }
140
141    fn revoke_device(&mut self, env: &Env, key_id: &BytesN<32>) -> Result<(), AccountError> {
142        DeviceStorage::update_device(env, &self.address, key_id, |d| {
143            if !d.is_active {
144                // Will still succeed but we need custom handling
145            }
146            d.is_active = false;
147        })
148    }
149
150    fn list_devices(&self, env: &Env) -> Vec<DeviceKey> {
151        DeviceStorage::load_devices(env, &self.address)
152    }
153
154    fn active_device_count(&self, env: &Env) -> usize {
155        let devices = DeviceStorage::load_devices(env, &self.address);
156        let mut count: usize = 0;
157        for i in 0..devices.len() {
158            if let Some(d) = devices.get(i) {
159                if d.is_active {
160                    count += 1;
161                }
162            }
163        }
164        count
165    }
166
167    fn update_last_used(&mut self, env: &Env, key_id: &BytesN<32>) -> Result<(), AccountError> {
168        let now = env.ledger().timestamp();
169        DeviceStorage::update_device(env, &self.address, key_id, |d| {
170            d.last_used = now;
171        })
172    }
173
174    fn set_policy(&mut self, env: &Env, policy: DevicePolicy) {
175        DeviceStorage::store_policy(env, &self.address, &policy);
176    }
177
178    fn policy(&self, env: &Env) -> DevicePolicy {
179        DeviceStorage::load_policy(env, &self.address).unwrap_or(DevicePolicy {
180            max_devices: 5,
181            auto_revoke_after: 0,
182        })
183    }
184
185    fn cleanup_inactive(&mut self, env: &Env) -> u32 {
186        let policy = DeviceStorage::load_policy(env, &self.address).unwrap_or(DevicePolicy {
187            max_devices: 5,
188            auto_revoke_after: 0,
189        });
190
191        if policy.auto_revoke_after == 0 {
192            return 0;
193        }
194
195        let now = env.ledger().timestamp();
196        let threshold = policy.auto_revoke_after;
197        let devices = DeviceStorage::load_devices(env, &self.address);
198        let mut new_devices: Vec<DeviceKey> = Vec::new(env);
199        let mut revoked: u32 = 0;
200
201        for i in 0..devices.len() {
202            if let Some(mut d) = devices.get(i) {
203                if d.is_active && now.saturating_sub(d.last_used) > threshold {
204                    d.is_active = false;
205                    revoked += 1;
206                }
207                new_devices.push_back(d);
208            }
209        }
210
211        if revoked > 0 {
212            DeviceStorage::store_devices(env, &self.address, &new_devices);
213        }
214
215        revoked
216    }
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222    use soroban_sdk::{contract, contractimpl, symbol_short, testutils::Address as _, Env};
223
224    #[contract]
225    pub struct TestContract;
226
227    #[contractimpl]
228    impl TestContract {}
229
230    fn default_policy() -> DevicePolicy {
231        DevicePolicy {
232            max_devices: 3,
233            auto_revoke_after: 0,
234        }
235    }
236
237    #[test]
238    fn test_register_device() {
239        let env = Env::default();
240        let contract_id = env.register(TestContract, ());
241        let addr = Address::generate(&env);
242
243        env.as_contract(&contract_id, || {
244            let mut manager = DeviceManager::new(addr, default_policy(), &env);
245            let key_id = BytesN::from_array(&env, &[1u8; 32]);
246            let device = manager
247                .register_device(&env, key_id, symbol_short!("phone"))
248                .unwrap();
249
250            assert!(device.is_active);
251            assert_eq!(manager.active_device_count(&env), 1);
252        });
253    }
254
255    #[test]
256    fn test_device_limit() {
257        let env = Env::default();
258        let contract_id = env.register(TestContract, ());
259        let addr = Address::generate(&env);
260
261        env.as_contract(&contract_id, || {
262            let policy = DevicePolicy {
263                max_devices: 2,
264                auto_revoke_after: 0,
265            };
266            let mut manager = DeviceManager::new(addr, policy, &env);
267
268            manager
269                .register_device(
270                    &env,
271                    BytesN::from_array(&env, &[1u8; 32]),
272                    symbol_short!("dev1"),
273                )
274                .unwrap();
275            manager
276                .register_device(
277                    &env,
278                    BytesN::from_array(&env, &[2u8; 32]),
279                    symbol_short!("dev2"),
280                )
281                .unwrap();
282
283            let result = manager.register_device(
284                &env,
285                BytesN::from_array(&env, &[3u8; 32]),
286                symbol_short!("dev3"),
287            );
288            assert_eq!(result, Err(AccountError::DeviceLimitReached));
289        });
290    }
291
292    #[test]
293    fn test_revoke_device() {
294        let env = Env::default();
295        let contract_id = env.register(TestContract, ());
296        let addr = Address::generate(&env);
297
298        env.as_contract(&contract_id, || {
299            let mut manager = DeviceManager::new(addr, default_policy(), &env);
300            let key_id = BytesN::from_array(&env, &[1u8; 32]);
301
302            manager
303                .register_device(&env, key_id.clone(), symbol_short!("phone"))
304                .unwrap();
305            manager.revoke_device(&env, &key_id).unwrap();
306
307            assert_eq!(manager.active_device_count(&env), 0);
308            assert_eq!(manager.list_devices(&env).len(), 1); // still in list, just inactive
309        });
310    }
311
312    #[test]
313    fn test_revoke_nonexistent() {
314        let env = Env::default();
315        let contract_id = env.register(TestContract, ());
316        let addr = Address::generate(&env);
317
318        env.as_contract(&contract_id, || {
319            let mut manager = DeviceManager::new(addr, default_policy(), &env);
320            let key_id = BytesN::from_array(&env, &[99u8; 32]);
321
322            let result = manager.revoke_device(&env, &key_id);
323            assert_eq!(result, Err(AccountError::DeviceNotFound));
324        });
325    }
326
327    #[test]
328    fn test_update_last_used() {
329        let env = Env::default();
330        let contract_id = env.register(TestContract, ());
331        let addr = Address::generate(&env);
332
333        env.as_contract(&contract_id, || {
334            let mut manager = DeviceManager::new(addr, default_policy(), &env);
335            let key_id = BytesN::from_array(&env, &[1u8; 32]);
336
337            manager
338                .register_device(&env, key_id.clone(), symbol_short!("phone"))
339                .unwrap();
340            manager.update_last_used(&env, &key_id).unwrap();
341
342            let devices = manager.list_devices(&env);
343            assert_eq!(devices.len(), 1);
344        });
345    }
346
347    #[test]
348    fn test_update_last_used_nonexistent() {
349        let env = Env::default();
350        let contract_id = env.register(TestContract, ());
351        let addr = Address::generate(&env);
352
353        env.as_contract(&contract_id, || {
354            let mut manager = DeviceManager::new(addr, default_policy(), &env);
355            let key_id = BytesN::from_array(&env, &[99u8; 32]);
356
357            let result = manager.update_last_used(&env, &key_id);
358            assert_eq!(result, Err(AccountError::DeviceNotFound));
359        });
360    }
361
362    #[test]
363    fn test_cleanup_inactive_disabled() {
364        let env = Env::default();
365        let contract_id = env.register(TestContract, ());
366        let addr = Address::generate(&env);
367
368        env.as_contract(&contract_id, || {
369            let mut manager = DeviceManager::new(addr, default_policy(), &env);
370
371            manager
372                .register_device(
373                    &env,
374                    BytesN::from_array(&env, &[1u8; 32]),
375                    symbol_short!("phone"),
376                )
377                .unwrap();
378
379            let revoked = manager.cleanup_inactive(&env);
380            assert_eq!(revoked, 0);
381            assert_eq!(manager.active_device_count(&env), 1);
382        });
383    }
384
385    #[test]
386    fn test_set_policy() {
387        let env = Env::default();
388        let contract_id = env.register(TestContract, ());
389        let addr = Address::generate(&env);
390
391        env.as_contract(&contract_id, || {
392            let mut manager = DeviceManager::new(addr, default_policy(), &env);
393            let new_policy = DevicePolicy {
394                max_devices: 10,
395                auto_revoke_after: 500,
396            };
397            manager.set_policy(&env, new_policy);
398            assert_eq!(manager.policy(&env).max_devices, 10);
399            assert_eq!(manager.policy(&env).auto_revoke_after, 500);
400        });
401    }
402
403    #[test]
404    fn test_with_defaults() {
405        let env = Env::default();
406        let contract_id = env.register(TestContract, ());
407        let addr = Address::generate(&env);
408
409        env.as_contract(&contract_id, || {
410            let manager = DeviceManager::with_defaults(addr, &env);
411            assert_eq!(manager.policy(&env).max_devices, 5);
412            assert_eq!(manager.active_device_count(&env), 0);
413        });
414    }
415
416    #[test]
417    fn test_revoked_device_allows_new_registration() {
418        let env = Env::default();
419        let contract_id = env.register(TestContract, ());
420        let addr = Address::generate(&env);
421
422        env.as_contract(&contract_id, || {
423            let policy = DevicePolicy {
424                max_devices: 1,
425                auto_revoke_after: 0,
426            };
427            let mut manager = DeviceManager::new(addr, policy, &env);
428
429            let key1 = BytesN::from_array(&env, &[1u8; 32]);
430            manager
431                .register_device(&env, key1.clone(), symbol_short!("old"))
432                .unwrap();
433            manager.revoke_device(&env, &key1).unwrap();
434
435            // Should be able to register again since active count is 0
436            let key2 = BytesN::from_array(&env, &[2u8; 32]);
437            manager
438                .register_device(&env, key2, symbol_short!("new"))
439                .unwrap();
440            assert_eq!(manager.active_device_count(&env), 1);
441        });
442    }
443}