1use super::{Header, Id, IeType};
6use crate::buffer_reader::BufferReader;
7use std::mem::size_of;
8use std::ops::Range;
9use zerocopy::SplitByteSlice;
10
11pub struct Reader<B>(BufferReader<B>);
14
15impl<B: SplitByteSlice> Reader<B> {
16 pub fn new(bytes: B) -> Self {
17 Reader(BufferReader::new(bytes))
18 }
19}
20
21impl<B: SplitByteSlice> Iterator for Reader<B> {
22 type Item = (Id, B);
23
24 fn next(&mut self) -> Option<Self::Item> {
25 let header = self.0.peek::<Header>()?;
26 let body_len = header.body_len as usize;
27 if self.0.bytes_remaining() < size_of::<Header>() + body_len {
28 None
29 } else {
30 let header = self.0.read::<Header>().unwrap();
32 let body = self.0.read_bytes(body_len).unwrap();
33 Some((header.id, body))
34 }
35 }
36}
37
38pub struct IeSummaryIter<B>(BufferReader<B>);
48
49impl<B: SplitByteSlice> IeSummaryIter<B> {
50 pub fn new(bytes: B) -> Self {
51 Self(BufferReader::new(bytes))
52 }
53}
54
55impl<B: SplitByteSlice> Iterator for IeSummaryIter<B> {
56 type Item = (IeType, Range<usize>);
57
58 fn next(&mut self) -> Option<Self::Item> {
59 loop {
60 let header = self.0.peek::<Header>()?;
61 let body_len = header.body_len as usize;
62
63 if self.0.bytes_remaining() < size_of::<Header>() + body_len {
65 return None;
66 }
67
68 let header = self.0.read::<Header>().unwrap();
70 let start_idx = self.0.bytes_read();
71 let body = self.0.read_bytes(body_len).unwrap();
72 let ie_type = match header.id {
73 Id::VENDOR_SPECIFIC => {
74 if body.len() >= 6 {
75 Some(IeType::new_vendor(body[0..6].try_into().unwrap()))
76 } else {
77 None
78 }
79 }
80 Id::EXTENSION => {
81 if body.len() >= 1 {
82 Some(IeType::new_extended(body[0]))
83 } else {
84 None
85 }
86 }
87 _ => Some(IeType::new_basic(header.id)),
88 };
89 match ie_type {
91 Some(ie_type) => {
92 return Some((ie_type, start_idx + ie_type.extra_len()..start_idx + body_len))
93 }
94 None => (),
95 }
96 }
97 }
98}
99
100#[cfg(test)]
101mod tests {
102 use super::*;
103
104 #[test]
105 pub fn empty() {
106 assert_eq!(None, Reader::new(&[][..]).next());
107 }
108
109 #[test]
110 pub fn less_than_header() {
111 assert_eq!(None, Reader::new(&[0][..]).next());
112 }
113
114 #[test]
115 pub fn body_too_short() {
116 assert_eq!(None, Reader::new(&[0, 2, 10][..]).next());
117 }
118
119 #[test]
120 pub fn empty_body() {
121 let elems: Vec<_> = Reader::new(&[0, 0][..]).collect();
122 assert_eq!(&[(Id::SSID, &[][..])], &elems[..]);
123 }
124
125 #[test]
126 pub fn two_elements() {
127 let bytes = vec![0, 2, 10, 20, 1, 3, 11, 22, 33];
128 let elems: Vec<_> = Reader::new(&bytes[..]).collect();
129 assert_eq!(
130 &[(Id::SSID, &[10, 20][..]), (Id::SUPPORTED_RATES, &[11, 22, 33][..])],
131 &elems[..]
132 );
133 }
134
135 #[test]
136 pub fn ie_summary_iter() {
137 let bytes = vec![
138 0, 2, 10, 20, 1, 0, 0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, 255, 2, 5, 1, ];
143 let elems: Vec<_> = IeSummaryIter::new(&bytes[..]).collect();
144 let expected = &[
145 (IeType::new_basic(Id::SSID), 2..4),
146 (IeType::new_basic(Id::SUPPORTED_RATES), 6..6),
147 (IeType::new_vendor([0x00, 0x03, 0x7f, 0x01, 0x01, 0x00]), 14..17),
148 (IeType::new_extended(5), 20..21),
149 ];
150 assert_eq!(&elems[..], expected);
151 }
152
153 #[test]
154 pub fn ie_summary_iter_skip_invalid_ies() {
155 let bytes = vec![
156 0, 2, 10, 20, 1, 0, 0xdd, 0x05, 0x00, 0x03, 0x7f, 0x01, 0x01, 0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, 255, 0, 255, 2, 5, 1, 2, 2, 1, ];
164 let elems: Vec<_> = IeSummaryIter::new(&bytes[..]).collect();
165 let expected = &[
166 (IeType::new_basic(Id::SSID), 2..4),
167 (IeType::new_basic(Id::SUPPORTED_RATES), 6..6),
168 (IeType::new_vendor([0x00, 0x03, 0x7f, 0x01, 0x01, 0x00]), 21..24),
169 (IeType::new_extended(5), 29..30),
170 ];
171 assert_eq!(&elems[..], expected);
172 }
173}