diff --git a/crates/cardwire-core/src/gpu/ebpf.rs b/crates/cardwire-core/src/gpu/ebpf.rs index 0e3bd76..a7c6f75 100644 --- a/crates/cardwire-core/src/gpu/ebpf.rs +++ b/crates/cardwire-core/src/gpu/ebpf.rs @@ -2,7 +2,7 @@ use std::collections::BTreeMap; use crate::{errors::Error as CardwireError, gpu::models::GpuDevice, pci::PciDevice}; -use cardwire_ebpf::EbpfBlocker; +use cardwire_ebpf::{BlockKind, EbpfBlocker}; use log::{info, warn}; pub struct GpuBlocker { @@ -18,20 +18,15 @@ impl GpuBlocker { pub fn set_nvidia_setting(&mut self, block: bool) -> Result<(), CardwireError> { self.inner - .set_nvidia_block(block) - .map_err(|err| CardwireError::UnknownBlockState(err.to_string()))?; + .block_kind(&block.to_string(), BlockKind::NvidiaSetting)?; Ok(()) } pub fn set_file_block(&mut self, file: &str) -> Result<(), CardwireError> { - self.inner - .set_file_block(file) - .map_err(|err| CardwireError::UnknownBlockState(err.to_string()))?; + self.inner.block_kind(file, BlockKind::File)?; Ok(()) } pub fn set_nvidia_file_block(&mut self, file: &str) -> Result<(), CardwireError> { - self.inner - .set_nvidia_file_block(file) - .map_err(|err| CardwireError::UnknownBlockState(err.to_string()))?; + self.inner.block_kind(file, BlockKind::NvidiaFile)?; Ok(()) } } @@ -39,25 +34,20 @@ impl GpuBlocker { pub fn is_gpu_blocked(blocker: &GpuBlocker, gpu: &GpuDevice) -> Result { let card_id = *gpu.card(); let render_id = *gpu.render(); + // PCI -> Card -> Render -> Nvidia Ok(blocker .inner - .is_pci_blocked(gpu.pci.pci_address()) - .map_err(|err| CardwireError::UnknownBlockState(err.to_string()))? + .is_kind_blocked(gpu.pci.pci_address(), BlockKind::Pci)? && blocker .inner - .is_card_blocked(card_id) - .map_err(|err| CardwireError::UnknownBlockState(err.to_string()))? + .is_kind_blocked(&card_id.to_string(), BlockKind::Card)? && blocker .inner - .is_render_blocked(render_id) - .map_err(|err| CardwireError::UnknownBlockState(err.to_string()))? - && if gpu.nvidia() { - // unwrap because it should be Some if it's an nvidia gpu, if not it's a bug and should - // be reported + .is_kind_blocked(&render_id.to_string(), BlockKind::Render)? + && if let Some(minor) = gpu.nvidia_minor() { blocker .inner - .is_nvidia_blocked(gpu.nvidia_minor().unwrap()) - .map_err(|err| CardwireError::UnknownBlockState(err.to_string()))? + .is_kind_blocked(&minor.to_string(), BlockKind::Nvidia)? } else { true }) @@ -73,19 +63,40 @@ pub fn block_gpu( let render_id = *gpu.render(); if block { - blocker.inner.block_card(card_id)?; - blocker.inner.block_render(render_id)?; + //blocker.inner.block_card(card_id)?; + // block card + blocker + .inner + .block_kind(&card_id.to_string(), BlockKind::Card)?; + // block render + blocker + .inner + .block_kind(&render_id.to_string(), BlockKind::Render)?; + // block pci chain_block_pci(blocker, gpu, pci_list)?; + // block nvidia if gpu.nvidia() { - blocker.inner.block_nvidia(gpu.nvidia_minor().unwrap())? + blocker + .inner + .block_kind(&gpu.nvidia_minor().unwrap().to_string(), BlockKind::Nvidia)?; } Ok(()) } else { - blocker.inner.unblock_card(card_id)?; - blocker.inner.unblock_render(render_id)?; + // unblock card + blocker + .inner + .unblock_kind(&card_id.to_string(), BlockKind::Card)?; + // unblock render + blocker + .inner + .unblock_kind(&render_id.to_string(), BlockKind::Render)?; + // unblock pci chain_unblock_pci(blocker, gpu, pci_list)?; + // unblock nvidia if gpu.nvidia() { - blocker.inner.unblock_nvidia(gpu.nvidia_minor().unwrap())? + blocker + .inner + .unblock_kind(&gpu.nvidia_minor().unwrap().to_string(), BlockKind::Nvidia)?; } Ok(()) } @@ -97,12 +108,16 @@ fn chain_block_pci( pci_list: &BTreeMap, ) -> Result<(), CardwireError> { // Block the gpu pci - blocker.inner.block_pci(gpu.pci.pci_address())?; + blocker + .inner + .block_kind(gpu.pci.pci_address(), BlockKind::Pci)?; info!("blocking pci: {}", gpu.pci.pci_address()); // also block audio card if gpu.pci.pci_address().ends_with(".0") { let gpu_audio_adress = gpu.pci.pci_address().replace(".0", ".1"); - blocker.inner.block_pci(&gpu_audio_adress)?; + blocker + .inner + .block_kind(&gpu_audio_adress, BlockKind::Pci)?; } // Check if gpu has a parent pci // first pci to block @@ -111,7 +126,9 @@ fn chain_block_pci( while let Some(parent_pci) = current_parent { if let Some(pci_device) = pci_list.get(&parent_pci) { info!("chain blocking pci: {}", pci_device.pci_address()); - blocker.inner.block_pci(pci_device.pci_address())?; + blocker + .inner + .block_kind(pci_device.pci_address(), BlockKind::Pci)?; current_parent = pci_device.parent_pci().clone(); } else { warn!("expected parent pci {} not found in pci_list", parent_pci); @@ -127,11 +144,16 @@ fn chain_unblock_pci( ) -> Result<(), CardwireError> { // Unblock the gpu pci info!("unblocking pci: {}", gpu.pci.pci_address()); - blocker.inner.unblock_pci(gpu.pci.pci_address())?; + blocker + .inner + .unblock_kind(gpu.pci.pci_address(), BlockKind::Pci)?; + // also unblock audio card if gpu.pci.pci_address().ends_with(".0") { let gpu_audio_adress = gpu.pci.pci_address().to_string().replace(".0", ".1"); - blocker.inner.unblock_pci(&gpu_audio_adress)?; + blocker + .inner + .unblock_kind(&gpu_audio_adress, BlockKind::Pci)?; } // Check if gpu has a parent pci // first pci to block @@ -140,7 +162,9 @@ fn chain_unblock_pci( while let Some(parent_pci) = current_parent { if let Some(pci_device) = pci_list.get(&parent_pci) { info!("chain unblocking pci: {}", pci_device.pci_address()); - blocker.inner.unblock_pci(pci_device.pci_address())?; + blocker + .inner + .unblock_kind(pci_device.pci_address(), BlockKind::Pci)?; current_parent = pci_device.parent_pci().clone(); } else { warn!("expected parent pci {} not found in pci_list", parent_pci); diff --git a/crates/cardwire-daemon/src/models.rs b/crates/cardwire-daemon/src/models.rs index 9f0708d..24b9da7 100644 --- a/crates/cardwire-daemon/src/models.rs +++ b/crates/cardwire-daemon/src/models.rs @@ -57,19 +57,15 @@ impl Modes { } pub struct DaemonState { + // these are file related pub config: RwLock, pub gpu_state: RwLock, pub mode_state: RwLock, + // temp data pub gpu_list: BTreeMap, pub ebpf_blocker: RwLock, // for future uses, related to vfio pub pci_devices: BTreeMap, - pub _iommu: bool, -} -impl DaemonState { - pub async fn _iommu(&self) -> bool { - self._iommu - } } pub struct Daemon { @@ -78,7 +74,6 @@ pub struct Daemon { impl Daemon { pub async fn new() -> Result { - let iommu: bool = pci::is_iommu_enabled(); let config = CardwireConfig::build().context("Error building config")?; let mut gpu_state = CardwireGpuState::build().context("Error building gpu_state")?; let mode_state = CardwireModeState::build().context("Error building mode")?; @@ -90,17 +85,17 @@ impl Daemon { if let Err(err) = check_default_drm_class(&mut gpu_list) { warn!("Failed to determine default GPU: {}", err); } - // TODO: Exit if ebpf returns an error, or try to recover from it? + // Exit if ebpf returns an error let ebpf_blocker = GpuBlocker::new()?; - // Do not stop the program if there is no gpu, cardwire will also be usable as a pci manager - // in a near future + if !gpu_list.is_empty() && gpu_state.is_default_state() { gpu_state .save_state(&gpu_list, &ebpf_blocker) .await .context("Could not save gpu state")?; } else if gpu_list.is_empty() { - warn!("could not detect gpus, daemon is still running for pci management usage") + // the daemon needs to be running to print out the pci list + warn!("could not detect gpus, daemon is still running for debugging") } Ok(Self { @@ -109,7 +104,6 @@ impl Daemon { gpu_state: RwLock::new(gpu_state), mode_state: RwLock::new(mode_state), pci_devices, - _iommu: iommu, gpu_list, ebpf_blocker: RwLock::new(ebpf_blocker), }, diff --git a/crates/cardwire-ebpf/src/bpf.c b/crates/cardwire-ebpf/src/bpf.c index ed966a4..d4e6847 100644 --- a/crates/cardwire-ebpf/src/bpf.c +++ b/crates/cardwire-ebpf/src/bpf.c @@ -36,7 +36,7 @@ struct qstr { const unsigned char *name; } __attribute__((preserve_access_index)); -struct dentry { +struct dentry___old { struct qstr d_name; struct dentry *d_parent; struct inode *d_inode; @@ -45,6 +45,13 @@ struct dentry { } d_u; } __attribute__((preserve_access_index)); +struct dentry { + struct qstr d_name; + struct dentry *d_parent; + struct inode *d_inode; + struct hlist_node d_alias; +} __attribute__((preserve_access_index)); + struct path { struct dentry *dentry; } __attribute__((preserve_access_index)); @@ -250,8 +257,13 @@ int BPF_PROG(inode_permission, struct inode *inode, int mask) return 0; } - unsigned long offset = - bpf_core_field_offset(struct dentry, d_u.d_alias); + unsigned long offset; + if (bpf_core_field_exists(((struct dentry___old *)0)->d_u.d_alias)) { + offset = + bpf_core_field_offset(struct dentry___old, d_u.d_alias); + } else { + offset = bpf_core_field_offset(struct dentry, d_alias); + } struct dentry *d = (struct dentry *)((void *)first - offset); // return is_blocked_device(d); diff --git a/crates/cardwire-ebpf/src/errors.rs b/crates/cardwire-ebpf/src/errors.rs index 345431d..1169aa3 100644 --- a/crates/cardwire-ebpf/src/errors.rs +++ b/crates/cardwire-ebpf/src/errors.rs @@ -10,22 +10,30 @@ pub enum CardwireEbpfError { Io(#[from] io::Error), #[error("couldn't load ebpf: {0}")] EbpfLoadError(String), - #[error("missing {kind}: {name}")] - MissingEntity { kind: String, name: String }, - #[error("aya error: {0}")] + #[error("missing lsm: {name}")] + MissingLsm { name: String }, + #[error("missing map: {name}")] + MissingMap { name: String }, + // for block/unblock, used if passed String is not in a pci format for example + #[error("wrong format, expected {kind} got: {input}")] + WrongFormat { kind: String, input: String }, + #[error("{0}")] Aya(String), #[error("{0}")] Other(String), } impl CardwireEbpfError { - pub fn missing_entity(kind: &str, name: &str) -> Self { - Self::MissingEntity { - kind: kind.to_string(), + pub fn missing_lsm(name: &str) -> Self { + Self::MissingLsm { + name: name.to_string(), + } + } + pub fn missing_map(name: &str) -> Self { + Self::MissingMap { name: name.to_string(), } } - pub fn aya(err: E) -> Self { Self::Aya(err.to_string()) } diff --git a/crates/cardwire-ebpf/src/lib.rs b/crates/cardwire-ebpf/src/lib.rs index 6e698f5..86adbad 100644 --- a/crates/cardwire-ebpf/src/lib.rs +++ b/crates/cardwire-ebpf/src/lib.rs @@ -1,6 +1,8 @@ //! main lib code of cardwire-ebpf mod errors; +use std::fmt; + pub use crate::errors::{CardwireEbpfError, CardwireEbpfResult}; use aya::{ Btf, Ebpf, maps::{HashMap, MapError}, programs::Lsm @@ -9,11 +11,37 @@ pub struct EbpfBlocker { ebpf: Ebpf, } +#[derive(PartialEq)] +pub enum BlockKind { + Card, + Render, + Pci, + Nvidia, + NvidiaSetting, + NvidiaFile, + File, +} +impl fmt::Display for BlockKind { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + match self { + BlockKind::Card => write!(f, "BLOCKED_CARDID"), + BlockKind::Render => write!(f, "BLOCKED_RENDERID"), + BlockKind::Pci => write!(f, "BLOCKED_PCI"), + BlockKind::Nvidia => write!(f, "BLOCKED_NVIDIAID"), + BlockKind::NvidiaSetting => write!(f, "SETTINGS"), + BlockKind::NvidiaFile => write!(f, "BLOCKED_NVIDIA_FILES"), + BlockKind::File => write!(f, "BLOCKED_PCI_FILES"), + } + } +} + impl EbpfBlocker { pub fn new() -> CardwireEbpfResult { + // quit if bpf is not enabled if !Self::is_bpf_enabled() { return Err(CardwireEbpfError::LSMNotEnabled); } + // load the program from the .o let mut ebpf = match Ebpf::load(aya::include_bytes_aligned!(concat!( env!("OUT_DIR"), "/bpf.o" @@ -24,42 +52,16 @@ impl EbpfBlocker { let btf = Btf::from_sys_fs().map_err(CardwireEbpfError::aya)?; - let program_file_open: &mut Lsm = ebpf - .program_mut("file_open") - .ok_or_else(|| Self::missing_entity("program", "file_open"))? - .try_into() - .map_err(CardwireEbpfError::aya)?; - program_file_open - .load("file_open", &btf) - .map_err(CardwireEbpfError::aya)?; - program_file_open.attach().map_err(CardwireEbpfError::aya)?; - - // For lsm/inode_permission - let program_inode_permission: &mut Lsm = ebpf - .program_mut("inode_permission") - .ok_or_else(|| Self::missing_entity("program", "inode_permission"))? - .try_into() - .map_err(CardwireEbpfError::aya)?; - program_inode_permission - .load("inode_permission", &btf) - .map_err(CardwireEbpfError::aya)?; - program_inode_permission - .attach() - .map_err(CardwireEbpfError::aya)?; - - // For lsm/inode_getattr - let program_inode_getattr: &mut Lsm = ebpf - .program_mut("inode_getattr") - .ok_or_else(|| Self::missing_entity("program", "inode_getattr"))? - .try_into() - .map_err(CardwireEbpfError::aya)?; - program_inode_getattr - .load("inode_getattr", &btf) - .map_err(CardwireEbpfError::aya)?; - program_inode_getattr - .attach() - .map_err(CardwireEbpfError::aya)?; - + let load_list: [&str; 3] = ["file_open", "inode_permission", "inode_getattr"]; + for entity in load_list { + let program: &mut Lsm = ebpf + .program_mut(entity) + .ok_or_else(|| CardwireEbpfError::missing_lsm(entity))? + .try_into() + .map_err(CardwireEbpfError::aya)?; + program.load(entity, &btf).map_err(CardwireEbpfError::aya)?; + program.attach().map_err(CardwireEbpfError::aya)?; + } Ok(Self { ebpf }) } @@ -83,11 +85,6 @@ impl EbpfBlocker { key[len] = 0; key } - - fn missing_entity(kind: &str, name: &str) -> CardwireEbpfError { - CardwireEbpfError::missing_entity(kind, name) - } - /* Checks if bpf/lsm is enabled in the kernel */ @@ -98,201 +95,184 @@ impl EbpfBlocker { } } + fn is_format_valid(entity: &str, kind: &BlockKind) -> bool { + match kind { + BlockKind::Render => entity.parse::().is_ok(), + BlockKind::Card => entity.parse::().is_ok(), + BlockKind::Nvidia => entity.parse::().is_ok(), + // either 0 or 1 + BlockKind::NvidiaSetting => entity.parse::().is_ok(), + // just a string + BlockKind::NvidiaFile | BlockKind::File => true, + // Only the Pci need a real check + BlockKind::Pci => entity.starts_with("0000:") && !entity.contains("pcie"), + } + } + /* - This part is for blocking a specific CardID + Block a kind */ + pub fn block_kind(&mut self, entity: &str, kind: BlockKind) -> CardwireEbpfResult<()> { + // validate input format for the bpf map, else return Err + if !Self::is_format_valid(entity, &kind) { + return Err(CardwireEbpfError::WrongFormat { + kind: kind.to_string(), + input: entity.to_string(), + }); + } - pub fn block_card(&mut self, id: u32) -> CardwireEbpfResult<()> { - let mut map: HashMap<_, u32, u8> = HashMap::try_from( - self.ebpf - .map_mut("BLOCKED_CARDID") - .ok_or_else(|| Self::missing_entity("map", "BLOCKED_CARDID"))?, - ) - .map_err(CardwireEbpfError::aya)?; - map.insert(id, 1, 0).map_err(CardwireEbpfError::aya)?; - Ok(()) - } + let kind_string = kind.to_string(); - pub fn unblock_card(&mut self, id: u32) -> CardwireEbpfResult<()> { - let mut map: HashMap<_, u32, u8> = HashMap::try_from( - self.ebpf - .map_mut("BLOCKED_CARDID") - .ok_or_else(|| Self::missing_entity("map", "BLOCKED_CARDID"))?, - ) - .map_err(CardwireEbpfError::aya)?; - let _ = map.remove(&id); - Ok(()) - } + match kind { + BlockKind::Pci => { + let mut map: HashMap<_, [u8; 16], u8> = HashMap::try_from( + self.ebpf + .map_mut(&kind_string) + .ok_or_else(|| CardwireEbpfError::missing_map(&kind_string))?, + ) + .map_err(CardwireEbpfError::aya)?; - pub fn is_card_blocked(&self, id: u32) -> CardwireEbpfResult { - let map: HashMap<_, u32, u8> = HashMap::try_from( - self.ebpf - .map("BLOCKED_CARDID") - .ok_or_else(|| Self::missing_entity("map", "BLOCKED_CARDID"))?, - ) - .map_err(CardwireEbpfError::aya)?; - match map.get(&id, 0) { - Ok(_) => Ok(true), - Err(MapError::KeyNotFound) => Ok(false), - Err(err) => Err(CardwireEbpfError::aya(err)), - } - } - /* - This part is for blocking a specific RenderID - */ + let key = Self::pci_key(entity); + map.insert(key, 1, 0).map_err(CardwireEbpfError::aya)?; + } + // set file blocklist + BlockKind::NvidiaFile | BlockKind::File => { + let mut map: HashMap<_, [u8; 30], u8> = HashMap::try_from( + self.ebpf + .map_mut(&kind_string) + .ok_or_else(|| CardwireEbpfError::missing_map(&kind_string))?, + ) + .map_err(CardwireEbpfError::aya)?; + let key = Self::file_key(entity); + map.insert(key, 1, 0).map_err(CardwireEbpfError::aya)?; + } + BlockKind::NvidiaSetting => { + if let Ok(block) = entity.parse::() { + let mut map: HashMap<_, u32, u8> = HashMap::try_from( + self.ebpf + .map_mut(&kind_string) + .ok_or_else(|| CardwireEbpfError::missing_map(&kind_string))?, + ) + .map_err(CardwireEbpfError::aya)?; + if block { + map.insert(0, 1, 0).map_err(CardwireEbpfError::aya)?; + } else { + let _ = map.remove(&0); + } + } + } + BlockKind::Render | BlockKind::Card | BlockKind::Nvidia => { + let mut map: HashMap<_, u32, u8> = HashMap::try_from( + self.ebpf + .map_mut(&kind_string) + .ok_or_else(|| CardwireEbpfError::missing_map(&kind_string))?, + ) + .map_err(CardwireEbpfError::aya)?; - pub fn block_render(&mut self, id: u32) -> CardwireEbpfResult<()> { - let mut map: HashMap<_, u32, u8> = HashMap::try_from( - self.ebpf - .map_mut("BLOCKED_RENDERID") - .ok_or_else(|| Self::missing_entity("map", "BLOCKED_RENDERID"))?, - ) - .map_err(CardwireEbpfError::aya)?; - map.insert(id, 1, 0).map_err(CardwireEbpfError::aya)?; - Ok(()) - } + if let Ok(value) = entity.parse::() { + map.insert(value, 1, 0).map_err(CardwireEbpfError::aya)?; + } + } + } - pub fn unblock_render(&mut self, id: u32) -> CardwireEbpfResult<()> { - let mut map: HashMap<_, u32, u8> = HashMap::try_from( - self.ebpf - .map_mut("BLOCKED_RENDERID") - .ok_or_else(|| Self::missing_entity("map", "BLOCKED_RENDERID"))?, - ) - .map_err(CardwireEbpfError::aya)?; - let _ = map.remove(&id); Ok(()) } - pub fn is_render_blocked(&self, id: u32) -> CardwireEbpfResult { - let map: HashMap<_, u32, u8> = HashMap::try_from( - self.ebpf - .map("BLOCKED_RENDERID") - .ok_or_else(|| Self::missing_entity("map", "BLOCKED_RENDERID"))?, - ) - .map_err(CardwireEbpfError::aya)?; - match map.get(&id, 0) { - Ok(_) => Ok(true), - Err(MapError::KeyNotFound) => Ok(false), - Err(err) => Err(CardwireEbpfError::aya(err)), - } - } /* - This part is for blocking a specific NvidiaID + Unblock a kind */ + pub fn unblock_kind(&mut self, entity: &str, kind: BlockKind) -> CardwireEbpfResult<()> { + // validate input format for the bpf map, else return Err + if !Self::is_format_valid(entity, &kind) { + return Err(CardwireEbpfError::WrongFormat { + kind: kind.to_string(), + input: entity.to_string(), + }); + } - pub fn block_nvidia(&mut self, id: u32) -> CardwireEbpfResult<()> { - let mut map: HashMap<_, u32, u8> = HashMap::try_from( - self.ebpf - .map_mut("BLOCKED_NVIDIAID") - .ok_or_else(|| Self::missing_entity("map", "BLOCKED_NVIDIAID"))?, - ) - .map_err(CardwireEbpfError::aya)?; - map.insert(id, 1, 0).map_err(CardwireEbpfError::aya)?; - Ok(()) - } + let kind_string = kind.to_string(); - pub fn unblock_nvidia(&mut self, id: u32) -> CardwireEbpfResult<()> { - let mut map: HashMap<_, u32, u8> = HashMap::try_from( - self.ebpf - .map_mut("BLOCKED_NVIDIAID") - .ok_or_else(|| Self::missing_entity("map", "BLOCKED_NVIDIAID"))?, - ) - .map_err(CardwireEbpfError::aya)?; - let _ = map.remove(&id); - Ok(()) - } + match kind { + BlockKind::Pci => { + let mut map: HashMap<_, [u8; 16], u8> = HashMap::try_from( + self.ebpf + .map_mut(&kind_string) + .ok_or_else(|| CardwireEbpfError::missing_map(&kind_string))?, + ) + .map_err(CardwireEbpfError::aya)?; + + let value = Self::pci_key(entity); + let _ = map.remove(&value); + } + // no file unblock + BlockKind::NvidiaFile | BlockKind::File | BlockKind::NvidiaSetting => (), + BlockKind::Render | BlockKind::Card | BlockKind::Nvidia => { + let mut map: HashMap<_, u32, u8> = HashMap::try_from( + self.ebpf + .map_mut(&kind_string) + .ok_or_else(|| CardwireEbpfError::missing_map(&kind_string))?, + ) + .map_err(CardwireEbpfError::aya)?; - pub fn is_nvidia_blocked(&self, id: u32) -> CardwireEbpfResult { - let map: HashMap<_, u32, u8> = HashMap::try_from( - self.ebpf - .map("BLOCKED_NVIDIAID") - .ok_or_else(|| Self::missing_entity("map", "BLOCKED_NVIDIAID"))?, - ) - .map_err(CardwireEbpfError::aya)?; - match map.get(&id, 0) { - Ok(_) => Ok(true), - Err(MapError::KeyNotFound) => Ok(false), - Err(err) => Err(CardwireEbpfError::aya(err)), + if let Ok(value) = entity.parse::() { + let _ = map.remove(&value); + } + } } - } - /* - This part is for blocking a specific PCI - */ - pub fn block_pci(&mut self, pci: &str) -> CardwireEbpfResult<()> { - let mut map: HashMap<_, [u8; 16], u8> = HashMap::try_from( - self.ebpf - .map_mut("BLOCKED_PCI") - .ok_or_else(|| Self::missing_entity("map", "BLOCKED_PCI"))?, - ) - .map_err(CardwireEbpfError::aya)?; - let key = Self::pci_key(pci); - map.insert(key, 1, 0).map_err(CardwireEbpfError::aya)?; - Ok(()) - } - pub fn unblock_pci(&mut self, pci: &str) -> CardwireEbpfResult<()> { - let mut map: HashMap<_, [u8; 16], u8> = HashMap::try_from( - self.ebpf - .map_mut("BLOCKED_PCI") - .ok_or_else(|| Self::missing_entity("map", "BLOCKED_PCI"))?, - ) - .map_err(CardwireEbpfError::aya)?; - let key = Self::pci_key(pci); - let _ = map.remove(&key); Ok(()) } - pub fn is_pci_blocked(&self, pci: &str) -> CardwireEbpfResult { - let map: HashMap<_, [u8; 16], u8> = HashMap::try_from( - self.ebpf - .map("BLOCKED_PCI") - .ok_or_else(|| Self::missing_entity("map", "BLOCKED_PCI"))?, - ) - .map_err(CardwireEbpfError::aya)?; - let key = Self::pci_key(pci); - match map.get(&key, 0) { - Ok(_) => Ok(true), - Err(MapError::KeyNotFound) => Ok(false), - Err(err) => Err(CardwireEbpfError::aya(err)), + /* + Check a block + */ + pub fn is_kind_blocked(&self, entity: &str, kind: BlockKind) -> CardwireEbpfResult { + // validate input format for the bpf map, else return Err + if !Self::is_format_valid(entity, &kind) { + return Err(CardwireEbpfError::WrongFormat { + kind: kind.to_string(), + input: entity.to_string(), + }); } - } - pub fn set_nvidia_block(&mut self, block: bool) -> CardwireEbpfResult<()> { - let mut map: HashMap<_, u32, u8> = HashMap::try_from( - self.ebpf - .map_mut("SETTINGS") - .ok_or_else(|| Self::missing_entity("map", "SETTINGS"))?, - ) - .map_err(CardwireEbpfError::aya)?; - if block { - map.insert(0, 1, 0).map_err(CardwireEbpfError::aya)?; - } else { - let _ = map.remove(&0); - } - Ok(()) - } + let kind_string = kind.to_string(); - pub fn set_nvidia_file_block(&mut self, file: &str) -> CardwireEbpfResult<()> { - let mut map: HashMap<_, [u8; 30], u8> = HashMap::try_from( - self.ebpf - .map_mut("BLOCKED_NVIDIA_FILES") - .ok_or_else(|| Self::missing_entity("map", "BLOCKED_PCI"))?, - ) - .map_err(CardwireEbpfError::aya)?; - let key = Self::file_key(file); - map.insert(key, 1, 0).map_err(CardwireEbpfError::aya)?; - Ok(()) - } + match kind { + BlockKind::Pci => { + let map: HashMap<_, [u8; 16], u8> = HashMap::try_from( + self.ebpf + .map(&kind_string) + .ok_or_else(|| CardwireEbpfError::missing_map(&kind_string))?, + ) + .map_err(CardwireEbpfError::aya)?; - pub fn set_file_block(&mut self, file: &str) -> CardwireEbpfResult<()> { - let mut map: HashMap<_, [u8; 30], u8> = HashMap::try_from( - self.ebpf - .map_mut("BLOCKED_PCI_FILES") - .ok_or_else(|| Self::missing_entity("map", "BLOCKED_PCI"))?, - ) - .map_err(CardwireEbpfError::aya)?; - let key = Self::file_key(file); - map.insert(key, 1, 0).map_err(CardwireEbpfError::aya)?; - Ok(()) + let value = Self::pci_key(entity); + return match map.get(&value, 0) { + Ok(_) => Ok(true), + Err(MapError::KeyNotFound) => Ok(false), + Err(err) => Err(CardwireEbpfError::aya(err)), + }; + } + // no file unblock + BlockKind::NvidiaFile | BlockKind::File | BlockKind::NvidiaSetting => (), + BlockKind::Render | BlockKind::Card | BlockKind::Nvidia => { + let map: HashMap<_, u32, u8> = HashMap::try_from( + self.ebpf + .map(&kind_string) + .ok_or_else(|| CardwireEbpfError::missing_map(&kind_string))?, + ) + .map_err(CardwireEbpfError::aya)?; + + if let Ok(value) = entity.parse::() { + return match map.get(&value, 0) { + Ok(_) => Ok(true), + Err(MapError::KeyNotFound) => Ok(false), + Err(err) => Err(CardwireEbpfError::aya(err)), + }; + } + } + } + + Ok(false) } }