Skip to content

Commit

Permalink
Outsource features in seperate file
Browse files Browse the repository at this point in the history 
Signed-off-by: Mr-Kanister <[email protected]>
  • Loading branch information
@Mr-Kanister
Mr-Kanister committed Nov 25, 2024
1 parent bb00a1e commit 12b0f29
Show file tree
Hide file tree
Showing 4 changed files with 149 additions and 207 deletions.
216 changes: 9 additions & 207 deletions rust/backend/daemon/src/ebpf_utils.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,15 +4,10 @@
//
// SPDX-License-Identifier: MIT

use aya::{
programs::{kprobe::KProbeLinkId, KProbe},
Ebpf, EbpfError,
};
use aya::{Ebpf, EbpfError};
use thiserror::Error;
// pub enum ProbeID {
// KProbeID(KProbeLinkId),
// UProbeID(UProbeLinkId),
// }

use crate::features::VfsFeature;

#[derive(Debug, Error)]
pub enum EbpfErrorWrapper {
Expand All @@ -26,203 +21,6 @@ impl From<EbpfErrorWrapper> for tonic::Status {
}
}

pub struct VfsFeature {
vfs_write_id: Option<KProbeLinkId>,
vfs_write_ret_id: Option<KProbeLinkId>,
}

impl VfsFeature {
pub fn new() -> VfsFeature {
VfsFeature {
vfs_write_id: None,
vfs_write_ret_id: None,
}
}

pub fn create(&mut self, ebpf: &mut Ebpf) -> Result<(), EbpfError> {
let vfs_write: &mut KProbe = ebpf
.program_mut("vfs_write")
.ok_or(EbpfError::ProgramError(
aya::programs::ProgramError::InvalidName {
name: "vfs_write".to_string(),
},
))?
.try_into()?;
vfs_write.load()?;

let vfs_write_ret: &mut KProbe = ebpf
.program_mut("vfs_write_ret")
.ok_or(EbpfError::ProgramError(
aya::programs::ProgramError::InvalidName {
name: "vfs_write_ret".to_string(),
},
))?
.try_into()?;
vfs_write_ret.load()?;

Ok(())
}

pub fn attach(&mut self, ebpf: &mut Ebpf) -> Result<(), EbpfError> {
let vfs_write: &mut KProbe = ebpf
.program_mut("vfs_write")
.ok_or(EbpfError::ProgramError(
aya::programs::ProgramError::InvalidName {
name: "vfs_write".to_string(),
},
))?
.try_into()?;
self.vfs_write_id = Some(vfs_write.attach("vfs_write", 0)?);

let vfs_write_ret: &mut KProbe = ebpf
.program_mut("vfs_write_ret")
.ok_or(EbpfError::ProgramError(
aya::programs::ProgramError::InvalidName {
name: "vfs_write_ret".to_string(),
},
))?
.try_into()?;
self.vfs_write_ret_id = Some(vfs_write_ret.attach("vfs_write", 0)?);
Ok(())
}

pub fn _detach(&mut self, ebpf: &mut Ebpf) -> Result<(), EbpfError> {
let vfs_write: &mut KProbe = ebpf
.program_mut("vfs_write")
.ok_or(EbpfError::ProgramError(
aya::programs::ProgramError::InvalidName {
name: "vfs_write".to_string(),
},
))?
.try_into()?;

if let Some(vfs_write_id) = self.vfs_write_id.take() {
vfs_write.detach(vfs_write_id)?;
} else {
return Err(EbpfError::ProgramError(
aya::programs::ProgramError::NotAttached,
));
}

let vfs_write_ret: &mut KProbe = ebpf
.program_mut("vfs_write_ret")
.ok_or(EbpfError::ProgramError(
aya::programs::ProgramError::InvalidName {
name: "vfs_write_ret".to_string(),
},
))?
.try_into()?;

if let Some(vfs_write_ret_id) = self.vfs_write_ret_id.take() {
vfs_write_ret.detach(vfs_write_ret_id)?;
} else {
return Err(EbpfError::ProgramError(
aya::programs::ProgramError::NotAttached,
));
}

Ok(())
}

// fn update(&mut self, ebpf: &mut Ebpf) {
// // update pids that are attached
// !todo!();
// }

pub fn _events(&mut self, _ebpf: &mut Ebpf) {
// return buffered stream of events
// will be discussed by Felix and Beni
}

pub fn _destroy(&mut self, ebpf: &mut Ebpf) -> Result<(), EbpfError> {
// TODO Error handling
let vfs_write: &mut KProbe = ebpf
.program_mut("vfs_write")
.ok_or(EbpfError::ProgramError(
aya::programs::ProgramError::InvalidName {
name: "vfs_write".to_string(),
},
))?
.try_into()?;
vfs_write.unload()?;

let vfs_write_ret: &mut KProbe = ebpf
.program_mut("vfs_write_ret")
.ok_or(EbpfError::ProgramError(
aya::programs::ProgramError::InvalidName {
name: "vfs_write_ret".to_string(),
},
))?
.try_into()?;
vfs_write_ret.unload()?;

Ok(())
}
}

// fn load_function(
// ebpf: &mut Ebpf,
// hash_map: &mut HashMap<String, ProbeID>,
// probe_type: Option<UprobeConfig>,
// func: &str,
// hook: &str,
// ) {
// /* examples:
// * func: "kprobetcp"
// * hook: "tcp_connect"
// */
// // TODO: Error checking

// match probe_type {
// // UPROBE
// Some(UprobeConfig {
// offset,
// target,
// pid,
// }) => {
// // get ebpf program
// let program: &mut UProbe = ebpf.program_mut(func).unwrap().try_into().unwrap();

// // load ebpf program
// program.load().unwrap();

// // attach ebpf program and insert its ProbeID into the hash map
// hash_map.insert(
// func.to_string(),
// ProbeID::UProbeID(program.attach(Some(hook), offset, target, pid).unwrap()),
// );
// }
// // KPROBE/KRETPROBE
// None => {
// let program: &mut KProbe = ebpf.program_mut(func).unwrap().try_into().unwrap();
// program.load().unwrap();
// hash_map.insert(
// func.to_string(),
// ProbeID::KProbeID(program.attach(hook, 0).unwrap()),
// );
// }
// }
// }

// fn unload_function(ebpf: &mut Ebpf, hash_map: &mut HashMap<String, ProbeID>, func: &str) {
// // get ProbeID and remove it from hash map
// let probe = hash_map.remove(func).unwrap();

// match probe {
// ProbeID::UProbeID(_link_id) => {
// // get ebpf program
// let program: &mut UProbe = ebpf.program_mut(func).unwrap().try_into().unwrap();

// // unload ebpf program
// program.unload().unwrap();
// }
// ProbeID::KProbeID(_link_id) => {
// let program: &mut KProbe = ebpf.program_mut(func).unwrap().try_into().unwrap();
// program.unload().unwrap();
// }
// }
// }

pub struct State {
vfs_write_feature: VfsFeature,
}
Expand All @@ -233,14 +31,18 @@ impl State {
vfs_write_feature: VfsFeature::new(),
}
}

pub fn init(&mut self, ebpf: &mut Ebpf) -> Result<(), EbpfError> {
self.vfs_write_feature.create(ebpf)?;

Ok(())
}

pub fn update_from_config(&mut self, ebpf: &mut Ebpf, _config_path: &str) -> Result<(), EbpfError> {
pub fn update_from_config(
&mut self,
ebpf: &mut Ebpf,
_config_path: &str,
) -> Result<(), EbpfError> {
self.vfs_write_feature.attach(ebpf)?;

Ok(())
Expand Down
138 changes: 138 additions & 0 deletions rust/backend/daemon/src/features.rs
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,138 @@
use aya::{
programs::{kprobe::KProbeLinkId, KProbe},
Ebpf, EbpfError,
};

pub struct VfsFeature {
vfs_write_id: Option<KProbeLinkId>,
vfs_write_ret_id: Option<KProbeLinkId>,
}

impl VfsFeature {
pub fn new() -> VfsFeature {
VfsFeature {
vfs_write_id: None,
vfs_write_ret_id: None,
}
}

pub fn create(&mut self, ebpf: &mut Ebpf) -> Result<(), EbpfError> {
let vfs_write: &mut KProbe = ebpf
.program_mut("vfs_write")
.ok_or(EbpfError::ProgramError(
aya::programs::ProgramError::InvalidName {
name: "vfs_write".to_string(),
},
))?
.try_into()?;
vfs_write.load()?;

let vfs_write_ret: &mut KProbe = ebpf
.program_mut("vfs_write_ret")
.ok_or(EbpfError::ProgramError(
aya::programs::ProgramError::InvalidName {
name: "vfs_write_ret".to_string(),
},
))?
.try_into()?;
vfs_write_ret.load()?;

Ok(())
}

pub fn attach(&mut self, ebpf: &mut Ebpf) -> Result<(), EbpfError> {
let vfs_write: &mut KProbe = ebpf
.program_mut("vfs_write")
.ok_or(EbpfError::ProgramError(
aya::programs::ProgramError::InvalidName {
name: "vfs_write".to_string(),
},
))?
.try_into()?;
self.vfs_write_id = Some(vfs_write.attach("vfs_write", 0)?);

let vfs_write_ret: &mut KProbe = ebpf
.program_mut("vfs_write_ret")
.ok_or(EbpfError::ProgramError(
aya::programs::ProgramError::InvalidName {
name: "vfs_write_ret".to_string(),
},
))?
.try_into()?;
self.vfs_write_ret_id = Some(vfs_write_ret.attach("vfs_write", 0)?);
Ok(())
}

pub fn detach(&mut self, ebpf: &mut Ebpf) -> Result<(), EbpfError> {
let vfs_write: &mut KProbe = ebpf
.program_mut("vfs_write")
.ok_or(EbpfError::ProgramError(
aya::programs::ProgramError::InvalidName {
name: "vfs_write".to_string(),
},
))?
.try_into()?;

if let Some(vfs_write_id) = self.vfs_write_id.take() {
vfs_write.detach(vfs_write_id)?;
} else {
return Err(EbpfError::ProgramError(
aya::programs::ProgramError::NotAttached,
));
}

let vfs_write_ret: &mut KProbe = ebpf
.program_mut("vfs_write_ret")
.ok_or(EbpfError::ProgramError(
aya::programs::ProgramError::InvalidName {
name: "vfs_write_ret".to_string(),
},
))?
.try_into()?;

if let Some(vfs_write_ret_id) = self.vfs_write_ret_id.take() {
vfs_write_ret.detach(vfs_write_ret_id)?;
} else {
return Err(EbpfError::ProgramError(
aya::programs::ProgramError::NotAttached,
));
}

Ok(())
}

// fn update(&mut self, ebpf: &mut Ebpf) {
// // update pids that are attached
// !todo!();
// }

// pub fn events(&mut self, ebpf: &mut Ebpf) {
// // return buffered stream of events
// // will be discussed by Felix and Beni
// }

pub fn destroy(&mut self, ebpf: &mut Ebpf) -> Result<(), EbpfError> {
// TODO Error handling
let vfs_write: &mut KProbe = ebpf
.program_mut("vfs_write")
.ok_or(EbpfError::ProgramError(
aya::programs::ProgramError::InvalidName {
name: "vfs_write".to_string(),
},
))?
.try_into()?;
vfs_write.unload()?;

let vfs_write_ret: &mut KProbe = ebpf
.program_mut("vfs_write_ret")
.ok_or(EbpfError::ProgramError(
aya::programs::ProgramError::InvalidName {
name: "vfs_write_ret".to_string(),
},
))?
.try_into()?;
vfs_write_ret.unload()?;

Ok(())
}
}
1 change: 1 addition & 0 deletions rust/backend/daemon/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ mod ebpf_utils;
mod helpers;
mod procfs_utils;
mod server;
mod features;

pub async fn run_server() {
helpers::bump_rlimit();
Expand Down
1 change: 1 addition & 0 deletions rust/backend/daemon/src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ mod ebpf_utils;
mod helpers;
mod procfs_utils;
mod server;
mod features;

#[tokio::main]
async fn main() {
Expand Down

0 comments on commit 12b0f29

Please sign in to comment.