diff --git a/proxy_agent/src/common/error.rs b/proxy_agent/src/common/error.rs index 2e7346f7..da4daedb 100644 --- a/proxy_agent/src/common/error.rs +++ b/proxy_agent/src/common/error.rs @@ -134,6 +134,9 @@ pub enum BpfErrorType { #[error("Failed to attach program '{0}' with error: {1}")] AttachBpfProgram(String, String), + #[error("Failed to close BPF object: {0}")] + CloseBpfObject(String), + #[error("Failed to convert program to '{0}' with error: {1}")] ConvertBpfProgram(String, String), diff --git a/proxy_agent/src/redirector.rs b/proxy_agent/src/redirector.rs index 56e9d4ba..892bdf88 100644 --- a/proxy_agent/src/redirector.rs +++ b/proxy_agent/src/redirector.rs @@ -36,7 +36,7 @@ //! let status = agent_status_shared_state.get_status(AgentStatusModule::Redirector).await; //! //! // Close the redirector to offload the eBPF program -//! redirector::close(redirector_shared_state.clone(), agent_status_shared_state.clone()).await; +//! redirector::close(redirector_shared_state.clone(), agent_status_shared_state.clone()).await.unwrap(); //! ``` #[cfg(windows)] @@ -68,10 +68,72 @@ use std::net::Ipv4Addr; use std::path::PathBuf; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; +use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; const ALERT_ONLY_EVENT_LOG_INTERVAL: Duration = Duration::from_secs(60 * 15); // 15 minutes +/// Runtime structure for managing event handling tasks associated with a BPF object. +#[derive(Default)] +pub(crate) struct EventTaskRuntime { + event_cancellation_token: Option, + /// Event handling tasks associated with this BPF object. + /// 1. The event-processing task created during start_internal. + /// 2. The platform ring-buffer reader task created during subscribe_alert_only. + event_tasks: Vec>, +} + +impl EventTaskRuntime { + fn set_cancellation_token(&mut self, cancellation_token: CancellationToken) { + self.event_cancellation_token = Some(cancellation_token); + } + + fn add_task(&mut self, task: JoinHandle<()>) { + self.event_tasks.push(task); + } + + /// Shuts down all event handling tasks associated with this BPF object. + /// Cancels the event cancellation token and awaits the completion of all tasks. + /// Returns the first error encountered, if any. + pub(crate) async fn shutdown(mut self) -> Result<()> { + if let Some(cancellation_token) = self.event_cancellation_token.take() { + cancellation_token.cancel(); + } + + let mut first_error = None; + for task in self.event_tasks.drain(..) { + if let Err(e) = task.await { + let error = Error::Bpf(BpfErrorType::CloseBpfObject(format!( + "eBPF event task failed while stopping: {e}" + ))); + logger::write_error(error.to_string()); + if first_error.is_none() { + first_error = Some(error); + } + } + } + + if let Some(error) = first_error { + return Err(error); + } + + Ok(()) + } +} + +/// Implements the `Drop` trait for `EventTaskRuntime`. +/// Ensures that all event handling tasks are aborted and the cancellation token is triggered when the runtime is dropped. +impl Drop for EventTaskRuntime { + fn drop(&mut self) { + if let Some(cancellation_token) = self.event_cancellation_token.take() { + cancellation_token.cancel(); + } + for task in self.event_tasks.drain(..) { + task.abort(); + } + } +} + #[cfg(not(windows))] pub use linux::BpfObject; #[cfg(windows)] @@ -190,13 +252,19 @@ impl AlertOnlyEventRateLimiter { pub struct Redirector { local_port: u16, shared_state: SharedState, + cancellation_token: CancellationToken, } impl Redirector { - pub fn new(local_port: u16, shared_state: &SharedState) -> Self { + pub fn new( + local_port: u16, + shared_state: &SharedState, + cancellation_token: CancellationToken, + ) -> Self { Redirector { local_port, shared_state: shared_state.clone(), + cancellation_token, } } @@ -204,6 +272,10 @@ impl Redirector { const RETRY_INTERVAL_MS: u64 = 10; pub async fn start(&self) { + if self.cancellation_token.is_cancelled() { + return; + } + let message = "eBPF redirector is starting"; if let Err(e) = self .shared_state @@ -217,8 +289,30 @@ impl Redirector { } let level = match self.start_impl().await { - Ok(_) => LoggerLevel::Info, - Err(_) => LoggerLevel::Error, + Ok(_) if !self.cancellation_token.is_cancelled() => LoggerLevel::Info, + Ok(_) => { + if let Err(e) = close( + self.shared_state.get_redirector_shared_state(), + self.shared_state.get_agent_status_shared_state(), + ) + .await + { + logger::write_error(format!( + "Failed to close redirector after startup cancellation: {e}" + )); + LoggerLevel::Error + } else { + LoggerLevel::Info + } + } + Err(_) => { + let _ = self + .shared_state + .get_agent_status_shared_state() + .set_module_state(ModuleState::STOPPED, AgentStatusModule::Redirector) + .await; + LoggerLevel::Error + } }; event_logger::write_event( level, @@ -231,6 +325,9 @@ impl Redirector { async fn start_impl(&self) -> Result<()> { for _ in 0..Self::MAX_RETRIES { + if self.cancellation_token.is_cancelled() { + return Err(Error::Bpf(BpfErrorType::FailedToStartRedirector)); + } match self.start_internal().await { Ok(_) => return Ok(()), Err(e) => { @@ -268,14 +365,21 @@ impl Redirector { self.attach_bpf_prog(&mut bpf_object)?; logger::write_information("Success attached bpf prog.".to_string()); - match bpf_object.subscribe_alert_only(self.shared_state.get_cancellation_token().clone()) { + let event_cancellation_token = self.cancellation_token.child_token(); + bpf_object + .event_runtime_mut() + .set_cancellation_token(event_cancellation_token.clone()); + match bpf_object.subscribe_alert_only(event_cancellation_token.clone()) { Ok(receiver) => { // Handle the received alert-only events here, spawn a task to process them - tokio::spawn(process_alert_only_events( + let event_processor_task = tokio::spawn(process_alert_only_events( receiver, self.shared_state.get_proxy_server_shared_state(), - self.shared_state.get_cancellation_token(), + event_cancellation_token, )); + bpf_object + .event_runtime_mut() + .add_task(event_processor_task); } Err(e) => { logger::write_error(format!("Failed to subscribe to alert-only events: {e}")); @@ -539,17 +643,21 @@ pub async fn update_local_ip_bind_monitor_only( pub async fn close( redirector_shared_state: RedirectorSharedState, agent_status_shared_state: AgentStatusSharedState, -) { - let _ = agent_status_shared_state - .set_module_state(ModuleState::STOPPED, AgentStatusModule::Redirector) - .await; - - // reset ebpf object +) -> Result<()> { #[cfg(windows)] { - windows::close_bpf_object(redirector_shared_state.clone()).await; + windows::close_bpf_object(redirector_shared_state).await?; } - let _ = redirector_shared_state.clear_bpf_object().await; + + #[cfg(not(windows))] + { + linux::close_bpf_object(redirector_shared_state).await?; + } + + agent_status_shared_state + .set_module_state(ModuleState::STOPPED, AgentStatusModule::Redirector) + .await?; + Ok(()) } #[cfg(not(windows))] diff --git a/proxy_agent/src/redirector/linux.rs b/proxy_agent/src/redirector/linux.rs index b3ce088e..f0cf8745 100644 --- a/proxy_agent/src/redirector/linux.rs +++ b/proxy_agent/src/redirector/linux.rs @@ -24,16 +24,38 @@ use proxy_agent_shared::{logger::LoggerLevel, misc_helpers}; use std::convert::TryFrom; use std::path::PathBuf; -pub struct BpfObject(Ebpf); +pub struct BpfObject { + ebpf: Option, + event_runtime: super::EventTaskRuntime, +} // BpfObject is a wrapper around Bpf object to interact with Linux eBPF programs and maps impl BpfObject { pub fn new(ebpf: Ebpf) -> Self { - BpfObject(ebpf) + BpfObject { + ebpf: Some(ebpf), + event_runtime: super::EventTaskRuntime::default(), + } } - pub fn get_bpf(&self) -> &Ebpf { - &self.0 + fn ebpf(&self) -> Result<&Ebpf> { + self.ebpf + .as_ref() + .ok_or(Error::Bpf(BpfErrorType::NullBpfObject)) + } + + fn ebpf_mut(&mut self) -> Result<&mut Ebpf> { + self.ebpf + .as_mut() + .ok_or(Error::Bpf(BpfErrorType::NullBpfObject)) + } + + pub(crate) fn event_runtime_mut(&mut self) -> &mut super::EventTaskRuntime { + &mut self.event_runtime + } + + fn take_resources(&mut self) -> (Option, super::EventTaskRuntime) { + (self.ebpf.take(), std::mem::take(&mut self.event_runtime)) } pub fn from_ebpf_file(bpf_file_path: &PathBuf) -> Result { @@ -60,7 +82,7 @@ impl BpfObject { pub fn update_skip_process_map(&mut self, pid: u32) -> Result<()> { let skip_process_map_name = SKIP_PROCESS_MAP_NAME; - match self.0.map_mut(skip_process_map_name) { + match self.ebpf_mut()?.map_mut(skip_process_map_name) { Some(map) => match HashMap::<&mut MapData, [u32; 1], [u32; 1]>::try_from(map) { Ok(mut skip_process_map) => { let key = sock_addr_skip_process_entry::from_pid(pid); @@ -95,7 +117,7 @@ impl BpfObject { pub fn update_local_ip_bind_monitor_only(&mut self, enabled: bool) -> Result<()> { let config_map_name = CONFIG_MAP_NAME; - match self.0.map_mut(config_map_name) { + match self.ebpf_mut()?.map_mut(config_map_name) { Some(map) => match HashMap::<&mut MapData, u32, [u32; 1]>::try_from(map) { Ok(mut config_map) => config_map .insert( @@ -135,7 +157,7 @@ impl BpfObject { dest_port: u16, ) -> Result<()> { let policy_map_name = POLICY_MAP_NAME; - match self.0.map_mut(policy_map_name) { + match self.ebpf_mut()?.map_mut(policy_map_name) { Some(map) => match HashMap::<&mut MapData, [u32; 6], [u32; 6]>::try_from(map) { Ok(mut policy_map) => { let local_ip = super::string_to_ip(constants::PROXY_AGENT_IP); @@ -182,7 +204,7 @@ impl BpfObject { program_name: &str, ) -> Result<()> { match std::fs::File::open(cgroup2_root_path.clone()) { - Ok(cgroup) => match self.0.program_mut(program_name) { + Ok(cgroup) => match self.ebpf_mut()?.program_mut(program_name) { Some(program) => match program.try_into() { Ok(p) => { let program: &mut CgroupSockAddr = p; @@ -236,7 +258,7 @@ impl BpfObject { pub fn attach_kprobe_program(&mut self) -> Result<()> { let program_name = "tcp_connect_probe"; - match self.0.program_mut(program_name) { + match self.ebpf_mut()?.program_mut(program_name) { Some(program) => match program.try_into() { Ok(p) => { let program: &mut KProbe = p; @@ -282,7 +304,7 @@ impl BpfObject { pub fn lookup_audit(&self, source_port: u16) -> Result { let audit_map_name = AUDIT_MAP_NAME; - match self.0.map(audit_map_name) { + match self.ebpf()?.map(audit_map_name) { Some(map) => match HashMap::<&MapData, AuditMapKey, AuditMapValue>::try_from(map) { Ok(audit_map) => { let key = sock_addr_audit_key::from_source_port(source_port); @@ -317,7 +339,10 @@ impl BpfObject { redirect: bool, ) -> bool { let policy_map_name = POLICY_MAP_NAME; - match self.0.map_mut(policy_map_name) { + let Ok(ebpf) = self.ebpf_mut() else { + return false; + }; + match ebpf.map_mut(policy_map_name) { Some(map) => match HashMap::<&mut MapData, [u32; 6], [u32; 6]>::try_from(map) { Ok(mut policy_map) => { let key = destination_entry::from_ipv4(dest_ipv4, dest_port); @@ -392,7 +417,7 @@ impl BpfObject { pub fn remove_audit_map_entry(&mut self, source_port: u16) -> Result<()> { let audit_map_name = AUDIT_MAP_NAME; - match self.0.map_mut(audit_map_name) { + match self.ebpf_mut()?.map_mut(audit_map_name) { Some(map) => match HashMap::<&mut MapData, AuditMapKey, AuditMapValue>::try_from(map) { Ok(mut audit_map) => { let key = sock_addr_audit_key::from_source_port(source_port); @@ -427,7 +452,7 @@ impl BpfObject { cancellation_token: tokio_util::sync::CancellationToken, ) -> Result> { let audit_map_name = ALERT_ONLY_MAP_NAME; - let map = self.0.take_map(audit_map_name).ok_or_else(|| { + let map = self.ebpf_mut()?.take_map(audit_map_name).ok_or_else(|| { Error::Bpf(BpfErrorType::GetBpfMap( audit_map_name.to_string(), "Map does not exist".to_string(), @@ -446,7 +471,7 @@ impl BpfObject { )) })?; let (sender, receiver) = tokio::sync::mpsc::unbounded_channel(); - tokio::spawn(async move { + let event_reader_task = tokio::spawn(async move { loop { tokio::select! { _ = cancellation_token.cancelled() => return, @@ -475,6 +500,7 @@ impl BpfObject { } } }); + self.event_runtime.add_task(event_reader_task); Ok(receiver) } } @@ -522,6 +548,28 @@ impl super::Redirector { } } +pub async fn close_bpf_object(redirector_shared_state: RedirectorSharedState) -> Result<()> { + let Some(bpf_object) = redirector_shared_state.take_bpf_object().await? else { + return Ok(()); + }; + + let (ebpf, event_runtime) = { + let mut bpf_object = bpf_object.lock().map_err(|e| { + Error::Bpf(BpfErrorType::CloseBpfObject(format!( + "BPF object lock is poisoned: {e}" + ))) + })?; + bpf_object.take_resources() + }; + + // Aya detaches managed links and closes program/map file descriptors when Ebpf is dropped. + drop(ebpf); + event_runtime.shutdown().await?; + + logger::write("Successfully closed Linux BPF object.".to_string()); + Ok(()) +} + pub async fn update_wire_server_redirect_policy( redirect: bool, redirector_shared_state: RedirectorSharedState, @@ -677,7 +725,7 @@ mod tests { // drop map_mut("audit_map") within this scope let mut audit_map: HashMap<&mut aya::maps::MapData, AuditMapKey, AuditMapValue> = HashMap::<&mut aya::maps::MapData, AuditMapKey, AuditMapValue>::try_from( - bpf.0.map_mut("audit_map").unwrap(), + bpf.ebpf_mut().unwrap().map_mut("audit_map").unwrap(), ) .unwrap(); audit_map diff --git a/proxy_agent/src/redirector/windows.rs b/proxy_agent/src/redirector/windows.rs index ad167ae5..b19299c3 100644 --- a/proxy_agent/src/redirector/windows.rs +++ b/proxy_agent/src/redirector/windows.rs @@ -20,6 +20,7 @@ use windows_sys::Win32::Networking::WinSock; pub struct BpfObject( pub *mut bpf_obj::bpf_object, pub Vec<*mut bpf_obj::ebpf_link_t>, + super::EventTaskRuntime, ); // Safety: bpf_object, which is a reference to an eBPF object, has no dependencies on thread-local storage and can // safely be sent to another thread. This is not explicitly documented in the Windows eBPF library, but the library does @@ -30,6 +31,16 @@ pub struct BpfObject( // [1] https://libbpf.readthedocs.io/en/v1.4.5/api.html#error-handling unsafe impl Send for BpfObject {} +impl BpfObject { + pub(crate) fn event_runtime_mut(&mut self) -> &mut super::EventTaskRuntime { + &mut self.2 + } + + fn take_event_runtime(&mut self) -> super::EventTaskRuntime { + std::mem::take(&mut self.2) + } +} + impl Default for BpfObject { fn default() -> Self { Self::new() @@ -100,11 +111,33 @@ impl super::Redirector { } } -pub async fn close_bpf_object(redirector_shared_state: RedirectorSharedState) { - if let Ok(Some(bpf_object)) = redirector_shared_state.get_bpf_object().await { - bpf_object.lock().unwrap().close_bpf_object(); - logger::write("Success closed bpf object.".to_string()); - } +pub async fn close_bpf_object(redirector_shared_state: RedirectorSharedState) -> Result<()> { + let Some(bpf_object) = redirector_shared_state.take_bpf_object().await? else { + return Ok(()); + }; + + let event_runtime = { + let mut bpf_object = bpf_object.lock().map_err(|e| { + Error::Bpf(BpfErrorType::CloseBpfObject(format!( + "BPF object lock is poisoned: {e}" + ))) + })?; + bpf_object.take_event_runtime() + }; + let shutdown_result = event_runtime.shutdown().await; + + bpf_object + .lock() + .map_err(|e| { + Error::Bpf(BpfErrorType::CloseBpfObject(format!( + "BPF object lock is poisoned: {e}" + ))) + })? + .close_bpf_object(); + + shutdown_result?; + logger::write("Success closed bpf object.".to_string()); + Ok(()) } pub fn get_audit_from_redirect_context(raw_socket_id: usize) -> Result { diff --git a/proxy_agent/src/redirector/windows/bpf_api.rs b/proxy_agent/src/redirector/windows/bpf_api.rs index 08d292e9..00652807 100644 --- a/proxy_agent/src/redirector/windows/bpf_api.rs +++ b/proxy_agent/src/redirector/windows/bpf_api.rs @@ -234,7 +234,7 @@ pub fn bpf_object__find_program_by_name( Ok(unsafe { find_program_by_name(obj, c_string.as_ptr()) }) } -pub fn ebpf_prog_attach( +pub fn ebpf_program_attach( prog: *mut ebpf_program_t, attach_type: *const ebpf_attach_type_t, attach_parameters: *const c_void, diff --git a/proxy_agent/src/redirector/windows/bpf_prog.rs b/proxy_agent/src/redirector/windows/bpf_prog.rs index 07541713..5b17252d 100644 --- a/proxy_agent/src/redirector/windows/bpf_prog.rs +++ b/proxy_agent/src/redirector/windows/bpf_prog.rs @@ -61,7 +61,11 @@ impl BpfObject { } pub fn new() -> Self { - Self(std::ptr::null::().cast_mut(), Vec::new()) + Self( + std::ptr::null::().cast_mut(), + Vec::new(), + super::super::EventTaskRuntime::default(), + ) } /** @@ -175,7 +179,7 @@ impl BpfObject { let compartment_id = 1; let mut link: ebpf_link_t = ebpf_link_t::empty(); let mut link: *mut ebpf_link_t = &mut link as *mut ebpf_link_t; - match ebpf_prog_attach( + match ebpf_program_attach( program, std::ptr::null(), &compartment_id as *const i32 as *const c_void, @@ -186,7 +190,7 @@ impl BpfObject { if r != 0 { return Err(Error::Bpf(BpfErrorType::AttachBpfProgram( program_name.to_string(), - format!("ebpf_prog_attach return with error code '{r}'"), + format!("ebpf_program_attach return with error code '{r}'"), ))); } logger::write_information(format!("Successfully attached {program_name} program.")); @@ -273,10 +277,6 @@ impl BpfObject { if self.0.is_null() { return; } - if let Err(e) = bpf_object__close(self.0) { - logger::write_error(format!("bpf_object__close with error: {e}")); - } - self.0 = std::ptr::null::().cast_mut(); for link in self.1.drain(..) { if link.is_null() { @@ -289,6 +289,11 @@ impl BpfObject { logger::write_error(format!("bpf_link_destroy with error: {e}")); } } + + if let Err(e) = bpf_object__close(self.0) { + logger::write_error(format!("bpf_object__close with error: {e}")); + } + self.0 = std::ptr::null::().cast_mut(); } /** @@ -466,7 +471,7 @@ impl BpfObject { } pub fn subscribe_alert_only( - &self, + &mut self, cancellation_token: tokio_util::sync::CancellationToken, ) -> Result> { let map_name = ALERT_ONLY_MAP_NAME; @@ -489,7 +494,7 @@ impl BpfObject { } let ring_address = ring as usize; let context_address = context as usize; - tokio::task::spawn_blocking(move || { + let event_reader_task = tokio::task::spawn_blocking(move || { let ring = ring_address as *mut ring_buffer; while !cancellation_token.is_cancelled() { match ring_buffer__poll(ring, 250) { @@ -498,7 +503,9 @@ impl BpfObject { logger::write_warning(format!( "ring_buffer__poll failed with result {result}" )); - break; + // TODO: Handle negative result appropriately, + // For now, just log and continue to poll the ring buffer + continue; } Err(err) => { logger::write_warning(format!("ring_buffer__poll failed: {err}")); @@ -513,6 +520,7 @@ impl BpfObject { )); } }); + self.2.add_task(event_reader_task); Ok(receiver) } diff --git a/proxy_agent/src/service.rs b/proxy_agent/src/service.rs index 75cfd31e..07eefaf4 100644 --- a/proxy_agent/src/service.rs +++ b/proxy_agent/src/service.rs @@ -85,7 +85,12 @@ pub async fn start_service(shared_state: SharedState) { }); tokio::spawn({ - let redirector: Redirector = Redirector::new(constants::PROXY_AGENT_PORT, &shared_state); + let redirector_cancellation_token = shared_state.get_cancellation_token().child_token(); + let redirector: Redirector = Redirector::new( + constants::PROXY_AGENT_PORT, + &shared_state, + redirector_cancellation_token, + ); async move { redirector.start().await; } @@ -174,7 +179,8 @@ pub fn stop_service(shared_state: SharedState) { tokio::spawn({ let shared_state = shared_state.clone(); async move { - redirector::close( + // Best effort to close the redirector gracefully. The result is ignored as we are stopping the service. + _ = redirector::close( shared_state.get_redirector_shared_state(), shared_state.get_agent_status_shared_state(), ) diff --git a/proxy_agent/src/shared_state/redirector_wrapper.rs b/proxy_agent/src/shared_state/redirector_wrapper.rs index a1e9d2dd..ea193577 100644 --- a/proxy_agent/src/shared_state/redirector_wrapper.rs +++ b/proxy_agent/src/shared_state/redirector_wrapper.rs @@ -41,6 +41,9 @@ enum RedirectorAction { GetBpfObject { response: oneshot::Sender>>>, }, + TakeBpfObject { + response: oneshot::Sender>>>, + }, } #[derive(Clone, Debug)] @@ -92,6 +95,14 @@ impl RedirectorSharedState { ); } } + RedirectorAction::TakeBpfObject { response } => { + if response.send(bpf_object.take()).is_err() { + logger::write_warning( + "Failed to send response to RedirectorAction::TakeBpfObject" + .to_string(), + ); + } + } } } }); @@ -159,6 +170,8 @@ impl RedirectorSharedState { self.set_bpf_object(None).await } + /// Retrieves the current BPF object without modifying the shared state. + /// Returns the BPF object if it exists, otherwise returns None. pub async fn get_bpf_object(&self) -> Result>>> { let (response_tx, response_rx) = oneshot::channel(); self.0 @@ -173,4 +186,21 @@ impl RedirectorSharedState { .await .map_err(|e| Error::RecvError("RedirectorAction::GetBpfObject".to_string(), e)) } + + /// Moves the stored Arc to the caller and replaces the shared-state value with None + /// Returns the BPF object if it exists, otherwise returns None. + pub async fn take_bpf_object(&self) -> Result>>> { + let (response_tx, response_rx) = oneshot::channel(); + self.0 + .send(RedirectorAction::TakeBpfObject { + response: response_tx, + }) + .await + .map_err(|e| { + Error::SendError("RedirectorAction::TakeBpfObject".to_string(), e.to_string()) + })?; + response_rx + .await + .map_err(|e| Error::RecvError("RedirectorAction::TakeBpfObject".to_string(), e)) + } }