use std::path::PathBuf; use std::mem::MaybeUninit; use std::time::{Instant, Duration}; use anyhow::{bail, Error}; use nix::sys::stat::Mode; use proxmox::tools::fs::{create_path, CreateOptions}; use proxmox_http::client::{RateLimit, RateLimiter, ShareableRateLimit}; use proxmox_shared_memory::{Init, SharedMemory, SharedMutex}; use proxmox_shared_memory::{check_subtype, initialize_subtype}; // openssl::sha::sha256(b"Proxmox Backup SharedRateLimiter v1.0")[0..8]; pub const PROXMOX_BACKUP_SHARED_RATE_LIMITER_MAGIC_1_0: [u8; 8] = [6, 58, 213, 96, 161, 122, 130, 117]; const BASE_PATH: &str = pbs_buildcfg::rundir!("/shmem/tbf"); // Wrap RateLimiter, so that we can provide an Init impl #[repr(C)] struct WrapLimiter(RateLimiter); impl Init for WrapLimiter { fn initialize(this: &mut MaybeUninit) { // default does not matter here, because we override later this.write(WrapLimiter(RateLimiter::new(1_000_000, 1_000_000))); } } #[repr(C)] struct SharedRateLimiterData { magic: [u8; 8], tbf: SharedMutex, padding: [u8; 4096 - 120], } impl Init for SharedRateLimiterData { fn initialize(this: &mut MaybeUninit) { unsafe { let me = &mut *this.as_mut_ptr(); me.magic = PROXMOX_BACKUP_SHARED_RATE_LIMITER_MAGIC_1_0; initialize_subtype(&mut me.tbf); } } fn check_type_magic(this: &MaybeUninit) -> Result<(), Error> { unsafe { let me = &*this.as_ptr(); if me.magic != PROXMOX_BACKUP_SHARED_RATE_LIMITER_MAGIC_1_0 { bail!("SharedRateLimiterData: wrong magic number"); } check_subtype(&me.tbf)?; Ok(()) } } } pub struct SharedRateLimiter { shmem: SharedMemory } impl SharedRateLimiter { pub fn mmap_shmem(name: &str, rate: u64, burst: u64) -> Result { let mut path = PathBuf::from(BASE_PATH); let user = pbs_config::backup_user()?; let dir_opts = CreateOptions::new() .perm(Mode::from_bits_truncate(0o770)) .owner(user.uid) .group(user.gid); create_path( &path, Some(dir_opts.clone()), Some(dir_opts))?; path.push(name); let file_opts = CreateOptions::new() .perm(Mode::from_bits_truncate(0o660)) .owner(user.uid) .group(user.gid); let shmem: SharedMemory = SharedMemory::open(&path, file_opts)?; shmem.data().tbf.lock().0.update_rate(rate, burst); Ok(Self { shmem }) } } impl ShareableRateLimit for SharedRateLimiter { fn update_rate(&self, rate: u64, bucket_size: u64) { self.shmem.data().tbf.lock().0 .update_rate(rate, bucket_size); } fn average_rate(&self, current_time: Instant) -> f64 { self.shmem.data().tbf.lock().0 .average_rate(current_time) } fn register_traffic(&self, current_time: Instant, data_len: u64) -> Duration { self.shmem.data().tbf.lock().0 .register_traffic(current_time, data_len) } }