wlan_common/ie/
reader.rs

1// Copyright 2019 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use super::{Header, Id, IeType};
6use crate::buffer_reader::BufferReader;
7use std::mem::size_of;
8use std::ops::Range;
9use zerocopy::SplitByteSlice;
10
11// TODO(https://fxbug.dev/42164332): Should probably remove Reader in favor of
12// IeSummaryIter everywhere.
13pub 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            // Unwraps are OK because we checked the length above
31            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
38/// An iterator that takes in a chain of IEs and produces summary for each IE.
39/// The summary is a tuple consisting of:
40/// - The IeType
41/// - The range of the rest of the IE:
42///   - If the IeType is basic, this range is the IE body
43///   - If the IeType is vendor, this range is the IE body without the first six bytes that
44///     identify the particular vendor IE
45///   - If the IeType is extended, this range is the IE body without the first byte that identifies
46///     the extension ID
47pub 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            // There are not enough bytes left, return None.
64            if self.0.bytes_remaining() < size_of::<Header>() + body_len {
65                return None;
66            }
67
68            // Unwraps are OK because we checked the length above.
69            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            // If IE type is valid, return the IE block. Otherwise, skip to the next one.
90            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, // IE with no extension ID
139            1, 0, // Empty IE
140            0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, // Vendor IE
141            255, 2, 5, 1, // IE with extension ID
142        ];
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, // IE with no extension ID
157            1, 0, // Empty IE
158            0xdd, 0x05, 0x00, 0x03, 0x7f, 0x01, 0x01, // Not long enough for vendor IE
159            0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, // Vendor IE
160            255, 0, // Not long enough for IE with extension ID
161            255, 2, 5, 1, // IE with extension ID
162            2, 2, 1, // Not enough trailing bytes
163        ];
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}