Skip to content

Commit

Permalink
Add field magic support
Browse files Browse the repository at this point in the history
  • Loading branch information
wcampbell0x2a committed Nov 15, 2024
1 parent cbad152 commit b1ad7a7
Show file tree
Hide file tree
Showing 5 changed files with 84 additions and 14 deletions.
8 changes: 8 additions & 0 deletions deku-derive/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -502,6 +502,9 @@ struct FieldData {

/// seek from start position
seek_from_start: Option<TokenStream>,

/// magic value that needs to appear before field
magic: Option<syn::LitByteStr>,
}

impl FieldData {
Expand Down Expand Up @@ -547,6 +550,7 @@ impl FieldData {
seek_from_current: receiver.seek_from_current?,
seek_from_end: receiver.seek_from_end?,
seek_from_start: receiver.seek_from_start?,
magic: receiver.magic,
};

FieldData::validate(&data)?;
Expand Down Expand Up @@ -980,6 +984,10 @@ struct DekuFieldReceiver {
/// seek from start position
#[darling(default = "default_res_opt", map = "map_litstr_as_tokenstream")]
seek_from_start: Result<Option<TokenStream>, ReplacementError>,

/// magic value that needs to appear before field
#[darling(default)]
magic: Option<syn::LitByteStr>,
}

/// Receiver for the variant-level attributes inside a enum
Expand Down
34 changes: 23 additions & 11 deletions deku-derive/src/macros/deku_read.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ use darling::ast::{Data, Fields};
use darling::ToTokens;
use proc_macro2::TokenStream;
use quote::quote;
use syn::{Ident, LitByteStr};

use crate::macros::{
gen_ctx_types_and_arg, gen_field_args, gen_internal_field_ident, gen_internal_field_idents,
Expand Down Expand Up @@ -495,20 +496,24 @@ fn emit_enum(input: &DekuData) -> Result<TokenStream, syn::Error> {
fn emit_magic_read(input: &DekuData) -> TokenStream {
let crate_ = super::get_crate_name();
if let Some(magic) = &input.magic {
quote! {
let __deku_magic = #magic;
emit_magic_read_lit(&crate_, magic)
} else {
quote! {}
}
}

for __deku_byte in __deku_magic {
let __deku_read_byte = u8::from_reader_with_ctx(__deku_reader, ())?;
if *__deku_byte != __deku_read_byte {
extern crate alloc;
use alloc::borrow::Cow;
return Err(::#crate_::DekuError::Parse(Cow::from(format!("Missing magic value {:?}", #magic))));
}
fn emit_magic_read_lit(crate_: &Ident, magic: &LitByteStr) -> TokenStream {
quote! {
let __deku_magic = #magic;

for __deku_byte in __deku_magic {
let __deku_read_byte = u8::from_reader_with_ctx(__deku_reader, ())?;
if *__deku_byte != __deku_read_byte {
extern crate alloc;
use alloc::borrow::Cow;
return Err(::#crate_::DekuError::Parse(Cow::from(format!("Missing magic value {:?}", #magic))));
}
}
} else {
quote! {}
}
}

Expand Down Expand Up @@ -738,6 +743,12 @@ fn emit_field_read(
quote! {}
};

let magic_read = if let Some(magic) = &f.magic {
emit_magic_read_lit(&crate_, magic)
} else {
quote! {}
};

let field_read_func = if field_reader.is_some() {
quote! { #field_reader? }
} else {
Expand Down Expand Up @@ -947,6 +958,7 @@ fn emit_field_read(

let field_read = quote! {
#seek
#magic_read
#pad_bits_before

#bit_offset
Expand Down
8 changes: 8 additions & 0 deletions deku-derive/src/macros/deku_write.rs
Original file line number Diff line number Diff line change
Expand Up @@ -528,6 +528,13 @@ fn emit_field_write(
) -> Result<TokenStream, syn::Error> {
let crate_ = super::get_crate_name();
let field_endian = f.endian.as_ref().or(input.endian.as_ref());
let magic_write = if let Some(magic) = &f.magic {
quote! {
::#crate_::DekuWriter::to_writer(#magic, __deku_writer, ())?;
}
} else {
quote! {}
};

let seek = if let Some(num) = &f.seek_from_current {
quote! {
Expand Down Expand Up @@ -705,6 +712,7 @@ fn emit_field_write(

let field_write = quote! {
#seek
#magic_write
#pad_bits_before

#bit_offset
Expand Down
29 changes: 26 additions & 3 deletions src/attributes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ enum DekuEnum {
| Attribute | Scope | Description
|-----------|------------------|------------
| [endian](#endian) | top-level, field | Set the endianness
| [magic](#magic) | top-level | A magic value that must be present at the start of this struct/enum
| [magic](#magic) | top-level, field | A magic value that must be present at the start of this struct/enum/field
| [seek_from_current](#seek_from_current) | top-level, field | Sets the offset of reader and writer to the current position plus the specified number of bytes
| [seek_from_end](#seek_from_end) | top-level, field | Sets the offset to the size of reader and writer plus the specified number of bytes
| [seek_from_start](#seek_from_start) | top-level, field | Sets the offset of reader and writer to provided number of bytes
Expand Down Expand Up @@ -151,10 +151,10 @@ assert_eq!(&*data, value);
# magic
Sets a "magic" value that must be present in the data at the start of
a struct/enum when reading, and that is written out of the start of
a struct/enum or field when reading, and that is written out of the start of
that type's data when writing.
Example:
Example (top-level):
```rust
# use deku::prelude::*;
# use std::convert::{TryInto, TryFrom};
Expand All @@ -177,6 +177,29 @@ let value: Vec<u8> = value.try_into().unwrap();
assert_eq!(data, value);
```
Example (field):
```rust
# use deku::prelude::*;
# use std::convert::{TryInto, TryFrom};
# #[derive(Debug, PartialEq, DekuRead, DekuWrite)]
struct DekuTest {
#[deku(magic = b"deku")]
data: u8
}
let data: &[u8] = &[b'd', b'e', b'k', b'u', 50];
let value = DekuTest::try_from(data).unwrap();
assert_eq!(
DekuTest { data: 50 },
value
);
let value: Vec<u8> = value.try_into().unwrap();
assert_eq!(data, value);
```
# seek_from_current
Using the internal reader, seek to current position plus offset before reading field.
Expand Down
19 changes: 19 additions & 0 deletions tests/test_magic.rs
Original file line number Diff line number Diff line change
Expand Up @@ -72,3 +72,22 @@ fn test_magic_enum(input: &[u8]) {
let ret_write: Vec<u8> = ret_read.try_into().unwrap();
assert_eq!(ret_write, input)
}

#[rstest(input,
case(&hex!("64656b7500")),
)]
fn test_struct_magic_field(input: &[u8]) {
#[derive(PartialEq, Debug, DekuRead, DekuWrite)]
struct TestStruct {
#[deku(magic = b"deku")]
magic: u8,
}
let input = input.to_vec();

let ret_read = TestStruct::try_from(input.as_slice()).unwrap();

assert_eq!(TestStruct { magic: 0 }, ret_read);

let ret_write: Vec<u8> = ret_read.try_into().unwrap();
assert_eq!(ret_write, input)
}

0 comments on commit b1ad7a7

Please sign in to comment.