rusticata_macros/
combinator.rs

1//! General purpose combinators
2
3use nom::bytes::streaming::take;
4use nom::combinator::map_parser;
5use nom::error::{make_error, ErrorKind, ParseError};
6use nom::{IResult, Input, Needed, Parser, ToUsize};
7
8#[deprecated(since = "3.0.1", note = "please use `be_var_u64` instead")]
9/// Read an entire slice as a big-endian value.
10///
11/// Returns the value as `u64`. This function checks for integer overflows, and returns a
12/// `Result::Err` value if the value is too big.
13pub fn bytes_to_u64(s: &[u8]) -> Result<u64, &'static str> {
14    let mut u: u64 = 0;
15
16    if s.is_empty() {
17        return Err("empty");
18    };
19    if s.len() > 8 {
20        return Err("overflow");
21    }
22    for &c in s {
23        let u1 = u << 8;
24        u = u1 | (c as u64);
25    }
26
27    Ok(u)
28}
29
30/// Read the entire slice as a big endian unsigned integer, up to 8 bytes
31#[inline]
32pub fn be_var_u64<'a, E: ParseError<&'a [u8]>>(input: &'a [u8]) -> IResult<&'a [u8], u64, E> {
33    if input.is_empty() {
34        return Err(nom::Err::Incomplete(Needed::new(1)));
35    }
36    if input.len() > 8 {
37        return Err(nom::Err::Error(make_error(input, ErrorKind::TooLarge)));
38    }
39    let mut res = 0u64;
40    for byte in input {
41        res = (res << 8) + *byte as u64;
42    }
43
44    Ok((&b""[..], res))
45}
46
47/// Read the entire slice as a little endian unsigned integer, up to 8 bytes
48#[inline]
49pub fn le_var_u64<'a, E: ParseError<&'a [u8]>>(input: &'a [u8]) -> IResult<&'a [u8], u64, E> {
50    if input.is_empty() {
51        return Err(nom::Err::Incomplete(Needed::new(1)));
52    }
53    if input.len() > 8 {
54        return Err(nom::Err::Error(make_error(input, ErrorKind::TooLarge)));
55    }
56    let mut res = 0u64;
57    for byte in input.iter().rev() {
58        res = (res << 8) + *byte as u64;
59    }
60
61    Ok((&b""[..], res))
62}
63
64/// Read a slice as a big-endian value.
65#[inline]
66pub fn parse_hex_to_u64<S>(i: &[u8], size: S) -> IResult<&[u8], u64>
67where
68    S: ToUsize + Copy,
69{
70    map_parser(take(size.to_usize()), be_var_u64).parse(i)
71}
72
73/// Apply combinator, automatically converts between errors if the underlying type supports it
74pub fn upgrade_error<I, O, E1, E2, F>(mut f: F) -> impl FnMut(I) -> IResult<I, O, E2>
75where
76    E1: ParseError<I>,
77    E2: ParseError<I> + From<E1>,
78    F: FnMut(I) -> IResult<I, O, E1>,
79{
80    move |i| f(i).map_err(nom::Err::convert)
81}
82
83/// Create a combinator that returns the provided value, and input unchanged
84pub fn pure<I, O, E>(val: O) -> impl Fn(I) -> IResult<I, O, E>
85where
86    O: Clone,
87    E: ParseError<I>,
88{
89    move |input: I| Ok((input, val.clone()))
90}
91
92/// Return a closure that takes `len` bytes from input, and applies `parser`.
93pub fn flat_take<I, C, O, E, F>(len: C, mut parser: F) -> impl FnMut(I) -> IResult<I, O, E>
94where
95    I: Input,
96    C: ToUsize + Copy,
97    E: ParseError<I>,
98    F: Parser<I, Output = O, Error = E>,
99{
100    // Note: this is the same as `map_parser(take(len), parser)`
101    move |input: I| {
102        let (input, o1) = take(len.to_usize())(input)?;
103        let (_, o2) = parser.parse(o1)?;
104        Ok((input, o2))
105    }
106}
107
108/// Take `len` bytes from `input`, and apply `parser`.
109pub fn flat_takec<I, O, E, C, F>(input: I, len: C, parser: F) -> IResult<I, O, E>
110where
111    I: Input,
112    E: ParseError<I>,
113    C: ToUsize + Copy,
114    F: Parser<I, Output = O, Error = E>,
115{
116    flat_take(len, parser)(input)
117}
118
119/// Helper macro for nom parsers: run first parser if condition is true, else second parser
120pub fn cond_else<I, O, E, C, F, G>(
121    cond: C,
122    mut first: F,
123    mut second: G,
124) -> impl FnMut(I) -> IResult<I, O, E>
125where
126    E: ParseError<I>,
127    C: Fn() -> bool,
128    F: Parser<I, Output = O, Error = E>,
129    G: Parser<I, Output = O, Error = E>,
130{
131    move |input: I| {
132        if cond() {
133            first.parse(input)
134        } else {
135            second.parse(input)
136        }
137    }
138}
139
140/// Align input value to the next multiple of n bytes
141/// Valid only if n is a power of 2
142pub const fn align_n2(x: usize, n: usize) -> usize {
143    (x + (n - 1)) & !(n - 1)
144}
145
146/// Align input value to the next multiple of 4 bytes
147pub const fn align32(x: usize) -> usize {
148    (x + 3) & !3
149}
150
151#[cfg(test)]
152mod tests {
153    use super::{align32, be_var_u64, cond_else, flat_take, pure};
154    use nom::bytes::streaming::take;
155    use nom::number::streaming::{be_u16, be_u32, be_u8};
156    use nom::{Err, IResult, Needed};
157
158    #[test]
159    fn test_be_var_u64() {
160        let res: IResult<&[u8], u64> = be_var_u64(b"\x12\x34\x56");
161        let (_, v) = res.expect("be_var_u64 failed");
162        assert_eq!(v, 0x123456);
163    }
164
165    #[test]
166    fn test_flat_take() {
167        let input = &[0x00, 0x01, 0xff];
168        // read first 2 bytes and use correct combinator: OK
169        let res: IResult<&[u8], u16> = flat_take(2u8, be_u16)(input);
170        assert_eq!(res, Ok((&input[2..], 0x0001)));
171        // read 3 bytes and use 2: OK (some input is just lost)
172        let res: IResult<&[u8], u16> = flat_take(3u8, be_u16)(input);
173        assert_eq!(res, Ok((&b""[..], 0x0001)));
174        // read 2 bytes and a combinator requiring more bytes
175        let res: IResult<&[u8], u32> = flat_take(2u8, be_u32)(input);
176        assert_eq!(res, Err(Err::Incomplete(Needed::new(2))));
177    }
178
179    #[test]
180    fn test_flat_take_str() {
181        let input = "abcdef";
182        // read first 2 bytes and use correct combinator: OK
183        let res: IResult<&str, &str> = flat_take(2u8, take(2u8))(input);
184        assert_eq!(res, Ok(("cdef", "ab")));
185        // read 3 bytes and use 2: OK (some input is just lost)
186        let res: IResult<&str, &str> = flat_take(3u8, take(2u8))(input);
187        assert_eq!(res, Ok(("def", "ab")));
188        // read 2 bytes and a use combinator requiring more bytes
189        let res: IResult<&str, &str> = flat_take(2u8, take(4u8))(input);
190        assert_eq!(res, Err(Err::Incomplete(Needed::Unknown)));
191    }
192
193    #[test]
194    fn test_cond_else() {
195        let input = &[0x01][..];
196        let empty = &b""[..];
197        let a = 1;
198        fn parse_u8(i: &[u8]) -> IResult<&[u8], u8> {
199            be_u8(i)
200        }
201        assert_eq!(
202            cond_else(|| a == 1, parse_u8, pure(0x02))(input),
203            Ok((empty, 0x01))
204        );
205        assert_eq!(
206            cond_else(|| a == 1, parse_u8, pure(0x02))(input),
207            Ok((empty, 0x01))
208        );
209        assert_eq!(
210            cond_else(|| a == 2, parse_u8, pure(0x02))(input),
211            Ok((input, 0x02))
212        );
213        assert_eq!(
214            cond_else(|| a == 1, pure(0x02), parse_u8)(input),
215            Ok((input, 0x02))
216        );
217        let res: IResult<&[u8], u8> = cond_else(|| a == 1, parse_u8, parse_u8)(input);
218        assert_eq!(res, Ok((empty, 0x01)));
219    }
220
221    #[test]
222    fn test_align32() {
223        assert_eq!(align32(3), 4);
224        assert_eq!(align32(4), 4);
225        assert_eq!(align32(5), 8);
226        assert_eq!(align32(5usize), 8);
227    }
228}
rusticata_macros/combinator.rs

rusticata_macros/
combinator.rs
