1 #![allow(dead_code, unused_variables, unused_imports)] 2 use alloc::vec::Vec; 3 use core::{hash::Hash, sync::atomic::AtomicU32}; 4 use system_error::SystemError; 5 6 use alloc::sync::Arc; 7 use hashbrown::HashMap; 8 use log::warn; 9 10 use super::user_namespace::UserNamespace; 11 use crate::libs::mutex::Mutex; 12 13 #[derive(Clone, Copy)] 14 pub enum Ucount { 15 UserNamespaces = 1, 16 PidNamespaces = 2, 17 UtsNamespaces = 3, 18 IpcNamespaces = 4, 19 NetNamespaces = 5, 20 MntNamespaces = 6, 21 CgroupNamespaces = 7, 22 TimeNamespaces = 8, 23 Counts = 9, 24 } 25 26 pub enum UcountRlimit { 27 Nproc = 1, 28 Msgqueue = 2, 29 Sigpending = 3, 30 Memlock = 4, 31 Counts = 5, 32 } 33 34 lazy_static! { 35 static ref COUNT_MANAGER: Arc<CountManager> = Arc::new(CountManager::new()); 36 } 37 38 #[derive(Debug)] 39 pub struct UCounts { 40 /// 对应的user_namespace 41 ns: Arc<UserNamespace>, 42 /// 用户标识符 43 uid: usize, 44 count: AtomicU32, 45 ucount: Vec<AtomicU32>, //[AtomicU32; UCOUNT_COUNTS as usize], 46 rlimit: Vec<AtomicU32>, //[AtomicU32; UCOUNT_RLIMIT_COUNTS as usize], 47 } 48 49 impl Default for UCounts { 50 fn default() -> Self { 51 Self::new() 52 } 53 } 54 impl UCounts { 55 pub fn new() -> Self { 56 Self { 57 ns: Arc::new(UserNamespace::new()), 58 uid: 0, 59 count: AtomicU32::new(1), 60 ucount: (0..Ucount::Counts as usize) 61 .map(|_| AtomicU32::new(0)) 62 .collect(), 63 rlimit: (0..UcountRlimit::Counts as usize) 64 .map(|_| AtomicU32::new(0)) 65 .collect(), 66 } 67 } 68 69 fn alloc_ucounts(&self, ns: Arc<UserNamespace>, uid: usize) -> Arc<Self> { 70 let mut counts = COUNT_MANAGER.counts.lock(); 71 let key = UKey { 72 user_ns: ns.clone(), 73 uid, 74 }; 75 let uc = if let Some(uc) = counts.get(&key) { 76 self.count 77 .fetch_add(1, core::sync::atomic::Ordering::SeqCst); 78 uc.clone() 79 } else { 80 Arc::new(Self { 81 ns, 82 uid, 83 count: AtomicU32::new(1), 84 ucount: (0..Ucount::Counts as usize) 85 .map(|_| AtomicU32::new(0)) 86 .collect(), 87 rlimit: (0..UcountRlimit::Counts as usize) 88 .map(|_| AtomicU32::new(0)) 89 .collect(), 90 }) 91 }; 92 counts.insert(key, uc.clone()); 93 uc 94 } 95 96 pub fn inc_ucounts( 97 &self, 98 user_ns: Arc<UserNamespace>, 99 uid: usize, 100 ucount_type: Ucount, 101 ) -> Option<Arc<UCounts>> { 102 let uc_type = ucount_type as usize; 103 let uc = self.alloc_ucounts(user_ns, uid); 104 let mut uc_iter = Some(uc.clone()); 105 let mut ucounts_add = vec![]; 106 while let Some(iter) = uc_iter { 107 let num = iter.ucount[uc_type].fetch_add(1, core::sync::atomic::Ordering::SeqCst); 108 ucounts_add.push(iter.clone()); 109 // 分配失败回滚 110 if num > iter.ns.ucount_max[uc_type] { 111 for add_iter in &ucounts_add { 112 add_iter.ucount[uc_type].fetch_sub(1, core::sync::atomic::Ordering::SeqCst); 113 } 114 return None; 115 } 116 uc_iter = iter.ns.ucounts.clone(); 117 } 118 return Some(uc); 119 } 120 121 fn find_ucounts(user_ns: Arc<UserNamespace>, uid: usize) -> Option<Arc<UCounts>> { 122 let counts = COUNT_MANAGER.counts.lock(); 123 let key = UKey { user_ns, uid }; 124 counts.get(&key).cloned() 125 } 126 127 fn get_ucounts(uc: Arc<UCounts>) { 128 let mut counts = COUNT_MANAGER.counts.lock(); 129 let ukey = UKey { 130 user_ns: uc.ns.clone(), 131 uid: uc.uid, 132 }; 133 counts.insert(ukey, uc); 134 } 135 136 pub fn dec_ucount(uc: Arc<UCounts>, ucount_type: Ucount) { 137 let mut uc_iter = Some(uc.clone()); 138 let uc_type = ucount_type as usize; 139 while let Some(iter) = uc_iter { 140 let num = iter.ucount[uc_type].fetch_sub(1, core::sync::atomic::Ordering::SeqCst); 141 if num == 0 { 142 warn!("count has reached zero"); 143 } 144 uc_iter = iter.ns.ucounts.clone(); 145 } 146 Self::put_ucounts(uc); 147 } 148 149 fn put_ucounts(uc: Arc<UCounts>) { 150 let mut counts = COUNT_MANAGER.counts.lock(); 151 let key = UKey { 152 user_ns: uc.ns.clone(), 153 uid: uc.uid, 154 }; 155 counts.remove(&key); 156 } 157 } 158 struct UKey { 159 user_ns: Arc<UserNamespace>, 160 uid: usize, 161 } 162 163 impl Hash for UKey { 164 fn hash<H: core::hash::Hasher>(&self, state: &mut H) { 165 let user_ns_ptr = Arc::as_ptr(&self.user_ns); 166 user_ns_ptr.hash(state); 167 self.uid.hash(state) 168 } 169 } 170 impl Eq for UKey {} 171 impl PartialEq for UKey { 172 fn eq(&self, other: &Self) -> bool { 173 Arc::ptr_eq(&self.user_ns, &other.user_ns) && self.uid == other.uid 174 } 175 } 176 177 struct CountManager { 178 counts: Mutex<HashMap<UKey, Arc<UCounts>>>, 179 } 180 181 impl CountManager { 182 fn new() -> Self { 183 Self { 184 counts: Mutex::new(HashMap::new()), 185 } 186 } 187 } 188