Skip to content

Commit

Permalink
Merge pull request #113 from KillingSpark/parse_refactor
Browse files Browse the repository at this point in the history
Refactor parsing. This changes the UnmarshalResult to not by default include a "bytes used" component. This was rarely used, bug-prone, and made implementations more unwieldy.

Also cleans up a lot of related code.
  • Loading branch information
KillingSpark authored Mar 5, 2024
2 parents 7edb138 + faa48d1 commit 4bc3326
Show file tree
Hide file tree
Showing 28 changed files with 867 additions and 859 deletions.
2 changes: 1 addition & 1 deletion rustbus/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ homepage = "https://github.com/KillingSpark/rustbus"

[dependencies]
nix = { version = "0.28", features = ["fs", "poll", "socket", "uio", "user"] }
rustbus_derive = {version = "0.5.0", path = "../rustbus_derive"}
rustbus_derive = {version = "0.6.0", path = "../rustbus_derive"}
thiserror = "1.0"

[dev-dependencies]
Expand Down
11 changes: 7 additions & 4 deletions rustbus/benches/marshal_benchmark.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,16 +6,19 @@ use rustbus::wire::marshal::marshal;
use rustbus::wire::unmarshal::unmarshal_dynamic_header;
use rustbus::wire::unmarshal::unmarshal_header;
use rustbus::wire::unmarshal::unmarshal_next_message;
use rustbus::wire::unmarshal_context::Cursor;

fn marsh(msg: &rustbus::message_builder::MarshalledMessage, buf: &mut Vec<u8>) {
marshal(msg, 0, buf).unwrap();
}

fn unmarshal(buf: &[u8]) {
let (hdrbytes, header) = unmarshal_header(buf, 0).unwrap();
let (dynhdrbytes, dynheader) = unmarshal_dynamic_header(&header, buf, hdrbytes).unwrap();
let (_, _unmarshed_msg) =
unmarshal_next_message(&header, dynheader, buf.to_vec(), hdrbytes + dynhdrbytes).unwrap();
let mut cursor = Cursor::new(buf);
let header = unmarshal_header(&mut cursor).unwrap();
let dynheader = unmarshal_dynamic_header(&header, &mut cursor).unwrap();
let _unmarshed_msg =
unmarshal_next_message(&header, dynheader, buf.to_vec(), cursor.consumed(), vec![])
.unwrap();
}

fn criterion_benchmark(c: &mut Criterion) {
Expand Down
24 changes: 11 additions & 13 deletions rustbus/fuzz/fuzz_targets/fuzz_unmarshal.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,24 +4,22 @@ extern crate libfuzzer_sys;
extern crate rustbus;

fuzz_target!(|data: &[u8]| {
let (hdrbytes, header) = match rustbus::wire::unmarshal::unmarshal_header(data, 0) {
Ok(head) => head,
Err(_) => return,
let mut cursor = rustbus::wire::unmarshal_context::Cursor::new(data);
let Ok(header) = rustbus::wire::unmarshal::unmarshal_header(&mut cursor) else {
return;
};
let Ok(dynheader) = rustbus::wire::unmarshal::unmarshal_dynamic_header(&header, &mut cursor) else {
return;
};
let (dynhdrbytes, dynheader) =
match rustbus::wire::unmarshal::unmarshal_dynamic_header(&header, data, hdrbytes) {
Ok(head) => head,
Err(_) => return,
};

let (_bytes_used, msg) = match rustbus::wire::unmarshal::unmarshal_next_message(
let Ok(msg) = rustbus::wire::unmarshal::unmarshal_next_message(
&header,
dynheader,
data.to_vec(),
hdrbytes + dynhdrbytes,
) {
Ok(msg) => msg,
Err(_) => return,
cursor.consumed(),
vec![],
) else {
return;
};
try_unmarhal_traits(&msg);
msg.unmarshall_all().ok();
Expand Down
27 changes: 14 additions & 13 deletions rustbus/src/bin/fuzz_artifact.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@
use std::io::Read;

use rustbus::wire::unmarshal_context::Cursor;

fn main() {
for path in std::env::args().skip(1) {
println!("Run artifact: {}", path);
Expand All @@ -16,29 +18,28 @@ fn run_artifact(path: &str) {
file.read_to_end(&mut data).unwrap();
let data = &data;

let (hdrbytes, header) = match rustbus::wire::unmarshal::unmarshal_header(data, 0) {
Ok(head) => head,
Err(_) => return,
let mut cursor = Cursor::new(data);
let Ok(header) = rustbus::wire::unmarshal::unmarshal_header(&mut cursor) else {
return;
};

println!("Header: {:?}", header);

let (dynhdrbytes, dynheader) =
match rustbus::wire::unmarshal::unmarshal_dynamic_header(&header, data, hdrbytes) {
Ok(head) => head,
Err(_) => return,
};
let Ok(dynheader) = rustbus::wire::unmarshal::unmarshal_dynamic_header(&header, &mut cursor)
else {
return;
};

println!("Dynheader: {:?}", dynheader);

let (_bytes_used, msg) = match rustbus::wire::unmarshal::unmarshal_next_message(
let Ok(msg) = rustbus::wire::unmarshal::unmarshal_next_message(
&header,
dynheader,
data.clone(),
hdrbytes + dynhdrbytes,
) {
Ok(msg) => msg,
Err(_) => return,
cursor.consumed(),
vec![],
) else {
return;
};

println!("Message: {:?}", msg);
Expand Down
39 changes: 18 additions & 21 deletions rustbus/src/connection/ll_conn.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@ use nix::sys::socket::{
SockaddrStorage, UnixAddr,
};

use crate::wire::unmarshal_context::Cursor;

/// A lowlevel abstraction over the raw unix socket
#[derive(Debug)]
pub struct SendConn {
Expand Down Expand Up @@ -166,8 +168,8 @@ impl RecvConn {
return Ok(16);
}
let msg_buf_in = &self.msg_buf_in.peek();
let (_, header) = unmarshal::unmarshal_header(msg_buf_in, 0)?;
let (_, header_fields_len) =
let header = unmarshal::unmarshal_header(&mut Cursor::new(msg_buf_in))?;
let header_fields_len =
crate::wire::util::parse_u32(&msg_buf_in[unmarshal::HEADER_LEN..], header.byteorder)?;
let complete_header_size = unmarshal::HEADER_LEN + header_fields_len as usize + 4; // +4 because the length of the header fields does not count

Expand Down Expand Up @@ -225,22 +227,22 @@ impl RecvConn {
/// Blocks until a message has been read from the conn or the timeout has been reached
pub fn get_next_message(&mut self, timeout: Timeout) -> Result<MarshalledMessage> {
self.read_whole_message(timeout)?;
let (hdrbytes, header) = unmarshal::unmarshal_header(self.msg_buf_in.peek(), 0)?;
let (dynhdrbytes, dynheader) =
unmarshal::unmarshal_dynamic_header(&header, self.msg_buf_in.peek(), hdrbytes)?;

let buf = self.msg_buf_in.take();
let buf_len = buf.len();
let (bytes_used, mut msg) =
unmarshal::unmarshal_next_message(&header, dynheader, buf, hdrbytes + dynhdrbytes)?;

msg.body.raw_fds = std::mem::take(&mut self.fds_in);
let mut cursor = Cursor::new(self.msg_buf_in.peek());
let header = unmarshal::unmarshal_header(&mut cursor)?;
let dynheader = unmarshal::unmarshal_dynamic_header(&header, &mut cursor)?;
let header_bytes_consumed = cursor.consumed();

if buf_len != bytes_used + hdrbytes + dynhdrbytes {
return Err(Error::UnmarshalError(UnmarshalError::NotAllBytesUsed));
}
let buf = self.msg_buf_in.take();
let raw_fds = std::mem::take(&mut self.fds_in);

Ok(msg)
Ok(unmarshal::unmarshal_next_message(
&header,
dynheader,
buf,
header_bytes_consumed,
raw_fds,
)?)
}
}

Expand Down Expand Up @@ -454,12 +456,7 @@ impl SendMessageContext<'_> {
// if this is not the first write for this message do not send the raw_fds again. This would lead to unexpected
// duplicated FDs on the other end!
let raw_fds = if self.state.bytes_sent == 0 {
self.msg
.body
.raw_fds
.iter()
.filter_map(|fd| fd.get_raw_fd())
.collect::<Vec<RawFd>>()
self.msg.body.get_raw_fds()
} else {
vec![]
};
Expand Down
60 changes: 36 additions & 24 deletions rustbus/src/message_builder.rs
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
//! Build new messages that you want to send over a connection
use std::os::fd::RawFd;

use crate::params::message;
use crate::signature::SignatureIter;
use crate::wire::errors::MarshalError;
use crate::wire::errors::UnmarshalError;
use crate::wire::marshal::traits::{Marshal, SignatureBuffer};
use crate::wire::marshal::MarshalContext;
use crate::wire::unmarshal::UnmarshalContext;
use crate::wire::unmarshal_context::UnmarshalContext;
use crate::wire::validate_raw;
use crate::wire::UnixFd;
use crate::ByteOrder;
Expand Down Expand Up @@ -269,14 +271,13 @@ impl MarshalledMessage {
} else {
let sigs: Vec<_> = crate::signature::Type::parse_description(&self.body.sig)?;

let (_, params) = crate::wire::unmarshal::unmarshal_body(
crate::wire::unmarshal::unmarshal_body(
self.body.byteorder,
&sigs,
self.body.get_buf(),
&self.body.raw_fds,
0,
)?;
params
)?
};
Ok(message::Message {
dynheader: self.dynheader,
Expand All @@ -295,10 +296,10 @@ pub struct MarshalledMessageBody {
buf_offset: usize,

// out of band data
pub(crate) raw_fds: Vec<crate::wire::UnixFd>,
raw_fds: Vec<crate::wire::UnixFd>,

sig: SignatureBuffer,
pub(crate) byteorder: ByteOrder,
byteorder: ByteOrder,
}

impl Default for MarshalledMessageBody {
Expand Down Expand Up @@ -377,11 +378,22 @@ impl MarshalledMessageBody {
&self.buf[self.buf_offset..]
}

pub fn get_raw_fds(&self) -> Vec<RawFd> {
self.raw_fds
.iter()
.filter_map(|fd| fd.get_raw_fd())
.collect::<Vec<RawFd>>()
}

pub fn byteorder(&self) -> ByteOrder {
self.byteorder
}

/// Get a clone of all the `UnixFd`s in the body.
///
/// Some of the `UnixFd`s may already have their `RawFd`s taken.
pub fn get_fds(&self) -> Vec<UnixFd> {
self.raw_fds.clone()
pub fn get_fds(&self) -> &[UnixFd] {
&self.raw_fds
}
/// Clears the buffer and signature but holds on to the memory allocations. You can now start pushing new
/// params as if this were a new message. This allows to reuse the OutMessage for the same dbus-message with different
Expand Down Expand Up @@ -790,15 +802,15 @@ impl<'fds, 'body: 'fds> MessageBodyParser<'body> {
return Err(UnmarshalError::WrongSignature);
}

let mut ctx = UnmarshalContext {
byteorder: self.body.byteorder,
buf: self.body.get_buf(),
offset: self.buf_idx,
fds: &self.body.raw_fds,
};
let mut ctx = UnmarshalContext::new(
&self.body.raw_fds,
self.body.byteorder,
self.body.get_buf(),
self.buf_idx,
);
match T::unmarshal(&mut ctx) {
Ok((bytes, res)) => {
self.buf_idx += bytes;
Ok(res) => {
self.buf_idx = self.body.get_buf().len() - ctx.remainder().len();
self.sig_idx += expected_sig.len();
Ok(res)
}
Expand Down Expand Up @@ -904,18 +916,18 @@ impl<'fds, 'body: 'fds> MessageBodyParser<'body> {
/// This checks if there are params left in the message and if the type you requested fits the signature of the message.
pub fn get_param(&mut self) -> Result<crate::params::Param, UnmarshalError> {
if let Some(sig_str) = self.get_next_sig() {
let mut ctx = UnmarshalContext {
byteorder: self.body.byteorder,
buf: self.body.get_buf(),
offset: self.buf_idx,
fds: &self.body.raw_fds,
};
let mut ctx = UnmarshalContext::new(
&self.body.raw_fds,
self.body.byteorder,
self.body.get_buf(),
self.buf_idx,
);

let sig = &crate::signature::Type::parse_description(sig_str).unwrap()[0];

match crate::wire::unmarshal::container::unmarshal_with_sig(sig, &mut ctx) {
Ok((bytes, res)) => {
self.buf_idx += bytes;
Ok(res) => {
self.buf_idx = self.body.get_buf().len() - ctx.remainder().len();
self.sig_idx += sig_str.len();
Ok(res)
}
Expand Down
2 changes: 1 addition & 1 deletion rustbus/src/params/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -227,7 +227,7 @@ impl Marshal for Variant<'_, '_> {
}
impl<'buf, 'fds> Unmarshal<'buf, 'fds> for Variant<'buf, 'fds> {
fn unmarshal(
ctx: &mut crate::wire::unmarshal::UnmarshalContext<'fds, 'buf>,
ctx: &mut crate::wire::unmarshal_context::UnmarshalContext<'fds, 'buf>,
) -> crate::wire::unmarshal::UnmarshalResult<Self> {
crate::wire::unmarshal::container::unmarshal_variant(ctx)
}
Expand Down
13 changes: 8 additions & 5 deletions rustbus/src/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ use crate::wire::marshal::marshal;
use crate::wire::unmarshal::unmarshal_dynamic_header;
use crate::wire::unmarshal::unmarshal_header;
use crate::wire::unmarshal::unmarshal_next_message;
use crate::wire::unmarshal_context::Cursor;

mod dbus_send;
mod fdpassing;
Expand Down Expand Up @@ -42,14 +43,16 @@ fn test_marshal_unmarshal() {
msg.dynheader.serial = Some(1);
let mut buf = Vec::new();
marshal(&msg, 0, &mut buf).unwrap();
let (hdrbytes, header) = unmarshal_header(&buf, 0).unwrap();
let (dynhdrbytes, dynheader) = unmarshal_dynamic_header(&header, &buf, hdrbytes).unwrap();

let headers_plus_padding = hdrbytes + dynhdrbytes + (8 - ((hdrbytes + dynhdrbytes) % 8));
let mut cursor = Cursor::new(&buf);
let header = unmarshal_header(&mut cursor).unwrap();
let dynheader = unmarshal_dynamic_header(&header, &mut cursor).unwrap();

let headers_plus_padding = cursor.consumed() + (8 - ((cursor.consumed()) % 8));
assert_eq!(headers_plus_padding, buf.len());

let (_, unmarshed_msg) =
unmarshal_next_message(&header, dynheader, msg.get_buf().to_vec(), 0).unwrap();
let unmarshed_msg =
unmarshal_next_message(&header, dynheader, msg.get_buf().to_vec(), 0, vec![]).unwrap();

let msg = unmarshed_msg.unmarshall_all().unwrap();

Expand Down
12 changes: 7 additions & 5 deletions rustbus/src/tests/fdpassing.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,10 @@ const TEST_STRING: &str = "This will be sent over the fd\n";

#[test]
fn test_fd_passing() {
let mut con1 =
connection::rpc_conn::RpcConn::system_conn(connection::Timeout::Infinite).unwrap();
let Ok(mut con1) = connection::rpc_conn::RpcConn::system_conn(connection::Timeout::Infinite)
else {
return;
};
let mut con2 =
connection::rpc_conn::RpcConn::system_conn(connection::Timeout::Infinite).unwrap();
con1.send_message(&mut crate::standard_messages::hello())
Expand Down Expand Up @@ -108,9 +110,9 @@ fn test_fd_marshalling() {

// assert the correct representation, where fds have been put into the fd array and the index is marshalled in the message
assert_eq!(sig.body.get_buf(), &[0, 0, 0, 0, 1, 0, 0, 0, 2, 0, 0, 0]);
assert_ne!(&sig.body.raw_fds[0], &test_fd1);
assert_ne!(&sig.body.raw_fds[1], &test_fd2);
assert_ne!(&sig.body.raw_fds[2], &test_fd3);
assert_ne!(&sig.body.get_fds()[0], &test_fd1);
assert_ne!(&sig.body.get_fds()[1], &test_fd2);
assert_ne!(&sig.body.get_fds()[2], &test_fd3);

// assert that unmarshalling yields the correct fds
let mut parser = sig.body.parser();
Expand Down
1 change: 1 addition & 0 deletions rustbus/src/wire.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
pub mod errors;
pub mod marshal;
pub mod unmarshal;
pub mod unmarshal_context;
pub mod util;
pub mod validate_raw;
pub mod variant_macros;
Expand Down
Loading

0 comments on commit 4bc3326

Please sign in to comment.