1use soroban_sdk::{contracttype, Address, BytesN, Env, Symbol, Vec};
2
3use super::device_storage::DeviceStorage;
4use super::error::AccountError;
5
6#[contracttype]
8#[derive(Clone, Debug, PartialEq, Eq)]
9pub struct DeviceKey {
10 pub key_id: BytesN<32>,
12 pub device_name: Symbol,
14 pub registered_at: u64,
16 pub last_used: u64,
18 pub is_active: bool,
20}
21
22#[contracttype]
24#[derive(Clone, Debug)]
25pub struct DevicePolicy {
26 pub max_devices: u32,
28 pub auto_revoke_after: u64,
31}
32
33pub trait MultiDeviceProvider {
35 fn register_device(
37 &mut self,
38 env: &Env,
39 key_id: BytesN<32>,
40 device_name: Symbol,
41 ) -> Result<DeviceKey, AccountError>;
42
43 fn revoke_device(&mut self, env: &Env, key_id: &BytesN<32>) -> Result<(), AccountError>;
45
46 fn list_devices(&self, env: &Env) -> Vec<DeviceKey>;
48
49 fn active_device_count(&self, env: &Env) -> usize;
51
52 fn update_last_used(&mut self, env: &Env, key_id: &BytesN<32>) -> Result<(), AccountError>;
54
55 fn set_policy(&mut self, env: &Env, policy: DevicePolicy);
57
58 fn policy(&self, env: &Env) -> DevicePolicy;
60
61 fn cleanup_inactive(&mut self, env: &Env) -> u32;
64}
65
66pub struct DeviceManager {
71 address: Address,
72}
73
74impl DeviceManager {
75 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 pub fn load(address: Address) -> Self {
84 Self { address }
85 }
86
87 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 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 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 }
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); });
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 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}