use zerocopy::{Immutable, IntoBytes};
use super::{
attr::{NLA_F_NESTED, NlAttr, nla_align},
message::{NLMSG_HDRLEN, NlMsgHdr, nlmsg_align},
};
#[derive(Debug, Clone, Copy)]
pub struct NestToken {
offset: usize,
}
#[derive(Debug, Clone)]
pub struct MessageBuilder {
buf: Vec<u8>,
}
impl MessageBuilder {
pub fn new(msg_type: u16, flags: u16) -> Self {
let header = NlMsgHdr::new(msg_type, flags);
let mut buf = vec![0u8; NLMSG_HDRLEN];
buf[..std::mem::size_of::<NlMsgHdr>()].copy_from_slice(header.as_bytes());
Self { buf }
}
pub fn with_header(header: NlMsgHdr) -> Self {
let mut buf = vec![0u8; NLMSG_HDRLEN];
buf[..std::mem::size_of::<NlMsgHdr>()].copy_from_slice(header.as_bytes());
Self { buf }
}
pub fn len(&self) -> usize {
self.buf.len()
}
pub fn is_empty(&self) -> bool {
self.buf.len() == NLMSG_HDRLEN
}
pub fn append_bytes(&mut self, data: &[u8]) {
self.buf.extend_from_slice(data);
let aligned = nlmsg_align(self.buf.len());
self.buf.resize(aligned, 0);
}
pub fn append<T: IntoBytes + Immutable>(&mut self, data: &T) {
self.append_bytes(data.as_bytes());
}
pub fn append_attr(&mut self, attr_type: u16, data: &[u8]) {
let attr = NlAttr::new(attr_type, data.len());
self.buf.extend_from_slice(attr.as_bytes());
self.buf.extend_from_slice(data);
let aligned = nla_align(self.buf.len());
self.buf.resize(aligned, 0);
}
pub fn append_attr_u8(&mut self, attr_type: u16, value: u8) {
self.append_attr(attr_type, &[value]);
}
pub fn append_attr_empty(&mut self, attr_type: u16) {
self.append_attr(attr_type, &[]);
}
pub fn append_attr_u16(&mut self, attr_type: u16, value: u16) {
self.append_attr(attr_type, &value.to_ne_bytes());
}
pub fn append_attr_u32(&mut self, attr_type: u16, value: u32) {
self.append_attr(attr_type, &value.to_ne_bytes());
}
pub fn append_attr_u64(&mut self, attr_type: u16, value: u64) {
self.append_attr(attr_type, &value.to_ne_bytes());
}
pub fn append_attr_u64_be(&mut self, attr_type: u16, value: u64) {
self.append_attr(attr_type, &value.to_be_bytes());
}
pub fn append_attr_u16_be(&mut self, attr_type: u16, value: u16) {
self.append_attr(attr_type, &value.to_be_bytes());
}
pub fn append_attr_u32_be(&mut self, attr_type: u16, value: u32) {
self.append_attr(attr_type, &value.to_be_bytes());
}
pub fn append_attr_str(&mut self, attr_type: u16, value: &str) {
let mut data = value.as_bytes().to_vec();
data.push(0); self.append_attr(attr_type, &data);
}
pub fn append_attr_string(&mut self, attr_type: u16, value: &str) {
self.append_attr(attr_type, value.as_bytes());
}
pub fn nest_start(&mut self, attr_type: u16) -> NestToken {
let offset = self.buf.len();
let attr = NlAttr::new(attr_type | NLA_F_NESTED, 0);
self.buf.extend_from_slice(attr.as_bytes());
NestToken { offset }
}
pub fn nest_end(&mut self, token: NestToken) {
let len = self.buf.len() - token.offset;
let len_bytes = (len as u16).to_ne_bytes();
self.buf[token.offset] = len_bytes[0];
self.buf[token.offset + 1] = len_bytes[1];
let aligned = nla_align(self.buf.len());
self.buf.resize(aligned, 0);
}
pub fn set_seq(&mut self, seq: u32) {
let bytes = seq.to_ne_bytes();
self.buf[8..12].copy_from_slice(&bytes);
}
pub fn set_pid(&mut self, pid: u32) {
let bytes = pid.to_ne_bytes();
self.buf[12..16].copy_from_slice(&bytes);
}
pub fn finish(mut self) -> Vec<u8> {
let len = self.buf.len() as u32;
let len_bytes = len.to_ne_bytes();
self.buf[0..4].copy_from_slice(&len_bytes);
self.buf
}
pub fn as_bytes(&self) -> &[u8] {
&self.buf
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::netlink::{attr::NLA_HDRLEN, message::NLM_F_REQUEST};
#[test]
fn test_simple_message() {
let msg = MessageBuilder::new(16, NLM_F_REQUEST).finish();
assert_eq!(msg.len(), NLMSG_HDRLEN);
let header = NlMsgHdr::from_bytes(&msg).unwrap();
assert_eq!(header.nlmsg_len as usize, NLMSG_HDRLEN);
assert_eq!(header.nlmsg_type, 16);
assert_eq!(header.nlmsg_flags, NLM_F_REQUEST);
}
#[test]
fn test_attribute() {
let mut builder = MessageBuilder::new(16, NLM_F_REQUEST);
builder.append_attr_u32(1, 0x12345678);
let msg = builder.finish();
assert!(msg.len() >= NLMSG_HDRLEN + NLA_HDRLEN + 4);
}
#[test]
fn test_nested_attribute() {
let mut builder = MessageBuilder::new(16, NLM_F_REQUEST);
let nest = builder.nest_start(1);
builder.append_attr_u32(2, 100);
builder.nest_end(nest);
let msg = builder.finish();
assert!(msg.len() > NLMSG_HDRLEN);
}
}