#![allow(dead_code, unused_variables, unused_imports)] use alloc::vec::Vec; use core::{hash::Hash, sync::atomic::AtomicU32}; use system_error::SystemError; use alloc::sync::Arc; use hashbrown::HashMap; use log::warn; use super::user_namespace::UserNamespace; use crate::libs::mutex::Mutex; #[derive(Clone, Copy)] pub enum Ucount { UserNamespaces = 1, PidNamespaces = 2, UtsNamespaces = 3, IpcNamespaces = 4, NetNamespaces = 5, MntNamespaces = 6, CgroupNamespaces = 7, TimeNamespaces = 8, Counts = 9, } pub enum UcountRlimit { Nproc = 1, Msgqueue = 2, Sigpending = 3, Memlock = 4, Counts = 5, } lazy_static! { static ref COUNT_MANAGER: Arc = Arc::new(CountManager::new()); } #[derive(Debug)] pub struct UCounts { /// 对应的user_namespace ns: Arc, /// 用户标识符 uid: usize, count: AtomicU32, ucount: Vec, //[AtomicU32; UCOUNT_COUNTS as usize], rlimit: Vec, //[AtomicU32; UCOUNT_RLIMIT_COUNTS as usize], } impl Default for UCounts { fn default() -> Self { Self::new() } } impl UCounts { pub fn new() -> Self { Self { ns: Arc::new(UserNamespace::new()), uid: 0, count: AtomicU32::new(1), ucount: (0..Ucount::Counts as usize) .map(|_| AtomicU32::new(0)) .collect(), rlimit: (0..UcountRlimit::Counts as usize) .map(|_| AtomicU32::new(0)) .collect(), } } fn alloc_ucounts(&self, ns: Arc, uid: usize) -> Arc { let mut counts = COUNT_MANAGER.counts.lock(); let key = UKey { user_ns: ns.clone(), uid, }; let uc = if let Some(uc) = counts.get(&key) { self.count .fetch_add(1, core::sync::atomic::Ordering::SeqCst); uc.clone() } else { Arc::new(Self { ns, uid, count: AtomicU32::new(1), ucount: (0..Ucount::Counts as usize) .map(|_| AtomicU32::new(0)) .collect(), rlimit: (0..UcountRlimit::Counts as usize) .map(|_| AtomicU32::new(0)) .collect(), }) }; counts.insert(key, uc.clone()); uc } pub fn inc_ucounts( &self, user_ns: Arc, uid: usize, ucount_type: Ucount, ) -> Option> { let uc_type = ucount_type as usize; let uc = self.alloc_ucounts(user_ns, uid); let mut uc_iter = Some(uc.clone()); let mut ucounts_add = vec![]; while let Some(iter) = uc_iter { let num = iter.ucount[uc_type].fetch_add(1, core::sync::atomic::Ordering::SeqCst); ucounts_add.push(iter.clone()); // 分配失败回滚 if num > iter.ns.ucount_max[uc_type] { for add_iter in &ucounts_add { add_iter.ucount[uc_type].fetch_sub(1, core::sync::atomic::Ordering::SeqCst); } return None; } uc_iter = iter.ns.ucounts.clone(); } return Some(uc); } fn find_ucounts(user_ns: Arc, uid: usize) -> Option> { let counts = COUNT_MANAGER.counts.lock(); let key = UKey { user_ns, uid }; counts.get(&key).cloned() } fn get_ucounts(uc: Arc) { let mut counts = COUNT_MANAGER.counts.lock(); let ukey = UKey { user_ns: uc.ns.clone(), uid: uc.uid, }; counts.insert(ukey, uc); } pub fn dec_ucount(uc: Arc, ucount_type: Ucount) { let mut uc_iter = Some(uc.clone()); let uc_type = ucount_type as usize; while let Some(iter) = uc_iter { let num = iter.ucount[uc_type].fetch_sub(1, core::sync::atomic::Ordering::SeqCst); if num == 0 { warn!("count has reached zero"); } uc_iter = iter.ns.ucounts.clone(); } Self::put_ucounts(uc); } fn put_ucounts(uc: Arc) { let mut counts = COUNT_MANAGER.counts.lock(); let key = UKey { user_ns: uc.ns.clone(), uid: uc.uid, }; counts.remove(&key); } } struct UKey { user_ns: Arc, uid: usize, } impl Hash for UKey { fn hash(&self, state: &mut H) { let user_ns_ptr = Arc::as_ptr(&self.user_ns); user_ns_ptr.hash(state); self.uid.hash(state) } } impl Eq for UKey {} impl PartialEq for UKey { fn eq(&self, other: &Self) -> bool { Arc::ptr_eq(&self.user_ns, &other.user_ns) && self.uid == other.uid } } struct CountManager { counts: Mutex>>, } impl CountManager { fn new() -> Self { Self { counts: Mutex::new(HashMap::new()), } } }