async_tftp/
parse.rs

1use nom::branch::alt;
2use nom::bytes::complete::{tag, tag_no_case, take_till};
3use nom::combinator::{map, map_opt, map_res, rest};
4use nom::multi::many0;
5use nom::number::complete::be_u16;
6use nom::sequence::tuple;
7use nom::IResult;
8use num_traits::FromPrimitive;
9use std::str::{self, FromStr};
10
11use crate::error::Result;
12use crate::packet::{self, *};
13
14#[derive(Debug)]
15enum Opt<'a> {
16    BlkSize(u16),
17    Timeout(u8),
18    Tsize(u64),
19    Invalid(&'a str, &'a str),
20}
21
22pub(crate) fn parse_packet(input: &[u8]) -> Result<Packet> {
23    let (rest, packet) = match parse_packet_type(input)? {
24        (data, PacketType::Rrq) => parse_rrq(data)?,
25        (data, PacketType::Wrq) => parse_wrq(data)?,
26        (data, PacketType::Data) => parse_data(data)?,
27        (data, PacketType::Ack) => parse_ack(data)?,
28        (data, PacketType::Error) => parse_error(data)?,
29        (data, PacketType::OAck) => parse_oack(data)?,
30    };
31
32    if rest.is_empty() {
33        Ok(packet)
34    } else {
35        Err(crate::Error::InvalidPacket)
36    }
37}
38
39fn nul_str(input: &[u8]) -> IResult<&[u8], &str> {
40    map_res(
41        tuple((take_till(|c| c == b'\0'), tag(b"\0"))),
42        |(s, _): (&[u8], _)| str::from_utf8(s),
43    )(input)
44}
45
46fn parse_packet_type(input: &[u8]) -> IResult<&[u8], PacketType> {
47    map_opt(be_u16, PacketType::from_u16)(input)
48}
49
50fn parse_mode(input: &[u8]) -> IResult<&[u8], Mode> {
51    alt((
52        map(tag_no_case(b"netascii\0"), |_| Mode::Netascii),
53        map(tag_no_case(b"octet\0"), |_| Mode::Octet),
54        map(tag_no_case(b"mail\0"), |_| Mode::Mail),
55    ))(input)
56}
57
58fn parse_opt_blksize(input: &[u8]) -> IResult<&[u8], Opt> {
59    map_opt(tuple((tag_no_case(b"blksize\0"), nul_str)), |(_, n): (_, &str)| {
60        u16::from_str(n)
61            .ok()
62            .filter(|n| *n >= 8 && *n <= 65464)
63            .map(Opt::BlkSize)
64    })(input)
65}
66
67fn parse_opt_timeout(input: &[u8]) -> IResult<&[u8], Opt> {
68    map_opt(tuple((tag_no_case(b"timeout\0"), nul_str)), |(_, n): (_, &str)| {
69        u8::from_str(n).ok().filter(|n| *n >= 1).map(Opt::Timeout)
70    })(input)
71}
72
73fn parse_opt_tsize(input: &[u8]) -> IResult<&[u8], Opt> {
74    map_opt(tuple((tag_no_case(b"tsize\0"), nul_str)), |(_, n): (_, &str)| {
75        u64::from_str(n).ok().map(Opt::Tsize)
76    })(input)
77}
78
79pub(crate) fn parse_opts(input: &[u8]) -> IResult<&[u8], Opts> {
80    many0(alt((
81        parse_opt_blksize,
82        parse_opt_timeout,
83        parse_opt_tsize,
84        map(tuple((nul_str, nul_str)), |(k, v)| Opt::Invalid(k, v)),
85    )))(input)
86    .map(|(i, opt_vec)| (i, to_opts(opt_vec)))
87}
88
89fn to_opts(opt_vec: Vec<Opt>) -> Opts {
90    let mut opts = Opts::default();
91
92    for opt in opt_vec {
93        match opt {
94            Opt::BlkSize(size) => {
95                if opts.block_size.is_none() {
96                    opts.block_size.replace(size);
97                }
98            }
99            Opt::Timeout(timeout) => {
100                if opts.timeout.is_none() {
101                    opts.timeout.replace(timeout);
102                }
103            }
104            Opt::Tsize(size) => {
105                if opts.transfer_size.is_none() {
106                    opts.transfer_size.replace(size);
107                }
108            }
109            Opt::Invalid(..) => {}
110        }
111    }
112
113    opts
114}
115
116fn parse_rrq(input: &[u8]) -> IResult<&[u8], Packet> {
117    let (input, (filename, mode, opts)) =
118        tuple((nul_str, parse_mode, parse_opts))(input)?;
119
120    Ok((
121        input,
122        Packet::Rrq(RwReq {
123            filename: filename.to_owned(),
124            mode,
125            opts,
126        }),
127    ))
128}
129
130fn parse_wrq(input: &[u8]) -> IResult<&[u8], Packet> {
131    let (input, (filename, mode, opts)) =
132        tuple((nul_str, parse_mode, parse_opts))(input)?;
133
134    Ok((
135        input,
136        Packet::Wrq(RwReq {
137            filename: filename.to_owned(),
138            mode,
139            opts,
140        }),
141    ))
142}
143
144fn parse_data(input: &[u8]) -> IResult<&[u8], Packet> {
145    tuple((be_u16, rest))(input)
146        .map(|(i, (block_nr, data))| (i, Packet::Data(block_nr, data)))
147}
148
149fn parse_ack(input: &[u8]) -> IResult<&[u8], Packet> {
150    be_u16(input).map(|(i, block_nr)| (i, Packet::Ack(block_nr)))
151}
152
153fn parse_error(input: &[u8]) -> IResult<&[u8], Packet> {
154    tuple((be_u16, nul_str))(input).map(|(i, (code, msg))| {
155        (i, packet::Error::from_code(code, Some(msg)).into())
156    })
157}
158
159fn parse_oack(input: &[u8]) -> IResult<&[u8], Packet> {
160    parse_opts(input).map(|(i, opts)| (i, Packet::OAck(opts)))
161}
async_tftp/parse.rs

async_tftp/
parse.rs
