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 => IeType::new_vendor(&body[..]),
74 Id::EXTENSION => {
75 if body.len() >= 1 {
76 Some(IeType::new_extended(body[0]))
77 } else {
78 None
79 }
80 }
81 _ => Some(IeType::new_basic(header.id)),
82 };
83 match ie_type {
85 Some(ie_type) => {
86 return Some((ie_type, start_idx + ie_type.extra_len()..start_idx + body_len));
87 }
88 None => (),
89 }
90 }
91 }
92}
93
94#[cfg(test)]
95mod tests {
96 use super::*;
97
98 #[test]
99 pub fn empty() {
100 assert_eq!(None, Reader::new(&[][..]).next());
101 }
102
103 #[test]
104 pub fn less_than_header() {
105 assert_eq!(None, Reader::new(&[0][..]).next());
106 }
107
108 #[test]
109 pub fn body_too_short() {
110 assert_eq!(None, Reader::new(&[0, 2, 10][..]).next());
111 }
112
113 #[test]
114 pub fn empty_body() {
115 let elems: Vec<_> = Reader::new(&[0, 0][..]).collect();
116 assert_eq!(&[(Id::SSID, &[][..])], &elems[..]);
117 }
118
119 #[test]
120 pub fn two_elements() {
121 let bytes = vec![0, 2, 10, 20, 1, 3, 11, 22, 33];
122 let elems: Vec<_> = Reader::new(&bytes[..]).collect();
123 assert_eq!(
124 &[(Id::SSID, &[10, 20][..]), (Id::SUPPORTED_RATES, &[11, 22, 33][..])],
125 &elems[..]
126 );
127 }
128
129 #[test]
130 pub fn ie_summary_iter() {
131 let bytes = vec![
132 0, 2, 10, 20, 1, 0, 0xdd, 0x05, 0x50, 0x6f, 0x9a, 0x1c, 0xaa, 0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, 255, 2, 5, 1, ];
138 let elems: Vec<_> = IeSummaryIter::new(&bytes[..]).collect();
139 let expected = &[
140 (IeType::new_basic(Id::SSID), 2..4),
141 (IeType::new_basic(Id::SUPPORTED_RATES), 6..6),
142 (IeType::new_vendor4([0x50, 0x6f, 0x9a, 0x1c]), 12..13),
143 (IeType::new_vendor6([0x00, 0x03, 0x7f, 0x01, 0x01, 0x00]), 21..24),
144 (IeType::new_extended(5), 27..28),
145 ];
146 assert_eq!(&elems[..], expected);
147 }
148
149 #[test]
150 pub fn ie_summary_iter_skip_invalid_ies() {
151 let bytes = vec![
152 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, ];
160 let elems: Vec<_> = IeSummaryIter::new(&bytes[..]).collect();
161 let expected = &[
162 (IeType::new_basic(Id::SSID), 2..4),
163 (IeType::new_basic(Id::SUPPORTED_RATES), 6..6),
164 (IeType::new_vendor6([0x00, 0x03, 0x7f, 0x01, 0x01, 0x00]), 21..24),
165 (IeType::new_extended(5), 29..30),
166 ];
167 assert_eq!(&elems[..], expected);
168 }
169}