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 => 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            // If IE type is valid, return the IE block. Otherwise, skip to the next one.
84            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, // IE with no extension ID
133            1, 0, // Empty IE
134            0xdd, 0x05, 0x50, 0x6f, 0x9a, 0x1c, 0xaa, // Vendor IE with known header
135            0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, // Vendor IE
136            255, 2, 5, 1, // IE with extension ID
137        ];
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, // IE with no extension ID
153            1, 0, // Empty IE
154            0xdd, 0x05, 0x00, 0x03, 0x7f, 0x01, 0x01, // Not long enough for vendor IE
155            0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, // Vendor IE
156            255, 0, // Not long enough for IE with extension ID
157            255, 2, 5, 1, // IE with extension ID
158            2, 2, 1, // Not enough trailing bytes
159        ];
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}