wlan_common/ie/
merger.rs

Help
1// Copyright 2020 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, IeSummaryIter, IeType};
6use anyhow::format_err;
7use std::collections::BTreeMap;
8use std::mem::size_of;
9use std::ops::Range;
10
11const IES_MERGER_BUFFER_LIMIT: usize = 10000;
12
13/// IesMerger is intended to be used to merge beacon and probe response IEs from scan results, as
14/// some IEs are only available on one frame type, or an IE may exist on both but contains more
15/// information on one frame type over the other.
16///
17/// IesMerger therefore has two purposes:
18/// 1. Combine unique IEs from each one into a single IE block.
19/// 2. Pick which IE to keep when the same type is available in multiple IE blocks but has
20///    different information:
21///    a. IesMerger knows about and uses custom logic for specific IE (e.g. for SSID, we
22///       prioritize non-hidden SSID over hidden SSID).
23///    b. In the general case, IesMerger keeps the IE that has more bytes (known use case:
24///       for WPS IE, the one in probe response contains more information).
25///    c. If the number of bytes are the same, IesMerger keeps the later one.
26#[derive(Debug)]
27pub struct IesMerger {
28    ies_updater: IesUpdater,
29    buffer_overflow: bool,
30    merge_ie_failures: usize,
31}
32
33impl IesMerger {
34    pub fn new(ies: Vec<u8>) -> Self {
35        Self { ies_updater: IesUpdater::new(ies), buffer_overflow: false, merge_ie_failures: 0 }
36    }
37
38    pub fn merge(&mut self, ies: &[u8]) {
39        for (ie_type, range) in IeSummaryIter::new(ies) {
40            let add_new = match self.ies_updater.get(&ie_type) {
41                Some(old) => should_add_new(ie_type, old, &ies[range.clone()]),
42                None => true,
43            };
44            let new_addition_len = range.end - range.start;
45            // IesMerger has a buffer limit so an unfriendly AP can't cause us to run out of
46            // memory by repeatedly changing the IEs.
47            if add_new {
48                if self.ies_updater.buf_len() + new_addition_len <= IES_MERGER_BUFFER_LIMIT {
49                    // Setting IE should not fail because we parsed them from an IE chain in the
50                    // first place, so the length of the IE body would not exceed 255 bytes.
51                    if let Err(_e) = self.ies_updater.set(ie_type, &ies[range]) {
52                        self.merge_ie_failures += 1;
53                    }
54                } else {
55                    self.buffer_overflow = true;
56                }
57            }
58        }
59    }
60
61    /// Build and return merged IEs, sorted by order of IeType.
62    pub fn finalize(&mut self) -> Vec<u8> {
63        self.ies_updater.finalize()
64    }
65
66    /// Return a bool indicating whether an IE was not merged because it would have exceeded
67    /// IesMerger's buffer.
68    pub fn buffer_overflow(&mut self) -> bool {
69        self.buffer_overflow
70    }
71
72    /// Return the number of times IesMerger decides to merge an IE but it fails to.
73    /// This should never occur, but we keep a counter in case the assumption is violated.
74    pub fn merge_ie_failures(&self) -> usize {
75        self.merge_ie_failures
76    }
77}
78
79#[derive(Debug)]
80pub struct IesUpdater {
81    // An index of the IEs we are going to keep. Mapping from IeType to the range of the
82    // corresponding bytes in the underlying buffer.
83    ies_summaries: BTreeMap<IeType, Range<usize>>,
84    ies_buf: Vec<u8>,
85}
86
87impl IesUpdater {
88    pub fn new(ies: Vec<u8>) -> Self {
89        let mut ies_summaries = BTreeMap::new();
90        for (ie_type, range) in IeSummaryIter::new(&ies[..]) {
91            ies_summaries.insert(ie_type, range);
92        }
93        Self { ies_summaries, ies_buf: ies }
94    }
95
96    /// Remove any IE with the corresponding `ie_type`.
97    pub fn remove(&mut self, ie_type: &IeType) {
98        self.ies_summaries.remove(ie_type);
99    }
100
101    /// Set an IE with the corresponding `ie_type`, replacing any existing entry with the same type.
102    /// The IE is rejected if it's too large.
103    pub fn set(&mut self, ie_type: IeType, ie_content: &[u8]) -> Result<(), anyhow::Error> {
104        // If length of this IE is too large, ignore it because later on we cannot construct the
105        // IE anyway on `finalize`.
106        if ie_type.extra_len() + ie_content.len() > std::u8::MAX.into() {
107            return Err(format_err!("ie_content too large"));
108        }
109
110        let start_idx = self.ies_buf.len();
111        self.ies_buf.extend_from_slice(ie_content);
112        self.ies_summaries.insert(ie_type, start_idx..self.ies_buf.len());
113        Ok(())
114    }
115
116    /// Set the raw IE (including the IE header).
117    /// The IE is rejected if the IE's length field (denoted by the second byte) is larger than
118    /// the length of the remaining IE bytes, and truncated if it's less than.
119    pub fn set_raw(&mut self, ie: &[u8]) -> Result<(), anyhow::Error> {
120        if ie.is_empty() {
121            return Ok(());
122        }
123        match IeSummaryIter::new(&ie[..]).next() {
124            Some((ie_type, range)) => self.set(ie_type, &ie[range]),
125            None => Err(format_err!("failed parsing `ie`")),
126        }
127    }
128
129    /// Get the IE bytes of the given `ie_type`, if it exists.
130    pub fn get(&self, ie_type: &IeType) -> Option<&[u8]> {
131        self.ies_summaries.get(ie_type).map(|range| &self.ies_buf[range.clone()])
132    }
133
134    fn buf_len(&self) -> usize {
135        self.ies_buf.len()
136    }
137
138    /// Build and return the modified IEs, sorted by order of IeType.
139    pub fn finalize(&mut self) -> Vec<u8> {
140        let total_len = self
141            .ies_summaries
142            .iter()
143            .map(|(ie_type, r)| size_of::<Header>() + ie_type.extra_len() + (r.end - r.start))
144            .sum();
145        let mut ies = Vec::with_capacity(total_len);
146        for (ie_type, range) in self.ies_summaries.iter() {
147            let id = ie_type.basic_id().0;
148            let len = ie_type.extra_len() + (range.end - range.start);
149            // Casting `len` to u8 is safe because in `set`, we reject any IE that is too large
150            ies.extend_from_slice(&[id, len as u8]);
151            ies.extend_from_slice(ie_type.extra_bytes());
152            ies.extend_from_slice(&self.ies_buf[range.clone()]);
153        }
154        ies
155    }
156}
157
158fn should_add_new(ie_type: IeType, old: &[u8], new: &[u8]) -> bool {
159    // If both IEs are the same, no need to add new IE
160    if old == new {
161        return false;
162    }
163
164    // If the new SSID is blank (hidden SSID), prefer the old one
165    if ie_type == IeType::SSID {
166        if new.len() < size_of::<Header>() {
167            return false;
168        }
169        let ssid = &new[2..];
170        if ssid.iter().all(|c| *c == 0u8) {
171            return false;
172        }
173    }
174
175    // In general case, prioritize the one with more information.
176    // For example, for WPS vendor IE, the AP sends a short one in the beacon and
177    // a longer one in the probe response.
178    //
179    // If they are the same length but different, then prioritize the new one.
180    new.len() >= old.len()
181}
182
183#[cfg(test)]
184mod tests {
185    use super::*;
186
187    // The beacon IEs here and the probe response IEs below are taken from the same router, with
188    // the following modification:
189    // 1. SSID is changed
190    // 2. Order of IEs are sorted according to `IeType` ordering to make test code simpler.
191    #[rustfmt::skip]
192    const BEACON_FRAME_IES: &'static [u8] = &[
193        // SSID: "foo-ssid"
194        0x00, 0x08, 0x66, 0x6f, 0x6f, 0x2d, 0x73, 0x73, 0x69, 0x64,
195        // Supported Rates: 6(B), 9, 12(B), 18, 24(B), 36, 48, 54, [Mbit/sec]
196        0x01, 0x08, 0x8c, 0x12, 0x98, 0x24, 0xb0, 0x48, 0x60, 0x6c,
197        // DS Parameter set: Current Channel: 157
198        0x03, 0x01, 0x9d,
199        // Traffic Indication Map (TIM): DTIM 0 of 0 bitmap
200        0x05, 0x04, 0x00, 0x01, 0x00, 0x00,
201        // Power Constraint: 3
202        0x20, 0x01, 0x03,
203        // HT Capabilities (802.11n D1.10)
204        0x2d, 0x1a,
205        0xef, 0x09, // HT Capabilities Info
206        0x1b, // A-MPDU Parameters: 0x1b
207        0xff, 0xff, 0xff, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, // MCS Set
208        0x00, 0x00, // HT Extended Capabilities
209        0x00, 0x00, 0x00, 0x00, // Transmit Beamforming Capabilities
210        0x00, // Antenna Selection Capabilities
211        // RSN Information
212        0x30, 0x14, 0x01, 0x00,
213        0x00, 0x0f, 0xac, 0x04, // Group Cipher: AES (CCM)
214        0x01, 0x00, 0x00, 0x0f, 0xac, 0x04, // Pairwise Cipher: AES (CCM)
215        0x01, 0x00, 0x00, 0x0f, 0xac, 0x02, // AKM: PSK
216        0x00, 0x00, // RSN Capabilities
217        // HT Information (802.11n D1.10)
218        0x3d, 0x16,
219        0x9d, // Primary Channel: 157
220        0x0d, // HT Info Subset - secondary channel above, any channel width, RIFS permitted
221        0x00, 0x00, 0x00, 0x00, // HT Info Subsets
222        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // Basic MCS Set
223        // Overlapping BSS Scan Parameters
224        0x4a, 0x0e, 0x14, 0x00, 0x0a, 0x00, 0x2c, 0x01, 0xc8, 0x00, 0x14, 0x00, 0x05, 0x00, 0x19, 0x00,
225        // Extended Capabilities
226        0x7f, 0x08, 0x01, 0x00, 0x0f, 0x00, 0x00, 0x00, 0x00, 0x40,
227        // VHT Capabilities
228        0xbf, 0x0c,
229        0xb2, 0x01, 0x80, 0x33, // VHT Capabilities Info
230        0xea, 0xff, 0x00, 0x00, 0xea, 0xff, 0x00, 0x00, // VHT Supported MCS Set
231        // VHT Operation
232        0xc0, 0x05, 0x01, 0x9b, 0x00, 0xfc, 0xff,
233        // VHT Tx Power Envelope
234        0xc3, 0x04, 0x02, 0xc4, 0xc4, 0xc4,
235        // Vendor Specific: Atheros Communications, Inc.: Advanced Capability
236        0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f,
237        // Vendor Specific: Microsoft Corp.: WMM/WME: Parameter Element
238        0xdd, 0x18, 0x00, 0x50, 0xf2, 0x02, 0x01, 0x01,
239        0x80, // U-APSD enabled
240        0x00,
241        0x03, 0xa4, 0x00, 0x00, // AC_BE parameters
242        0x27, 0xa4, 0x00, 0x00, // AC_BK parameters
243        0x42, 0x43, 0x5e, 0x00, // AC_VI parameters
244        0x62, 0x32, 0x2f, 0x00, // AC_VO parameters
245        // Vendor Specific: Microsoft Corp.: WPS
246        0xdd, 0x1d, 0x00, 0x50, 0xf2, 0x04, 0x10, 0x4a, 0x00, 0x01, 0x10, 0x10, 0x44, 0x00, 0x01,
247        0x02, 0x10, 0x3c, 0x00, 0x01, 0x03, 0x10, 0x49, 0x00, 0x06, 0x00, 0x37, 0x2a, 0x00, 0x01,
248        0x20,
249    ];
250
251    #[rustfmt::skip]
252    const PROBE_RESP_IES: &'static [u8] = &[
253        // SSID: "foo-ssid"
254        0x00, 0x08, 0x66, 0x6f, 0x6f, 0x2d, 0x73, 0x73, 0x69, 0x64,
255        // Supported Rates: 6(B), 9, 12(B), 18, 24(B), 36, 48, 54, [Mbit/sec]
256        0x01, 0x08, 0x8c, 0x12, 0x98, 0x24, 0xb0, 0x48, 0x60, 0x6c,
257        // DS Parameter set: Current Channel: 157
258        0x03, 0x01, 0x9d,
259        // Power Constraint: 3
260        0x20, 0x01, 0x03,
261        // HT Capabilities (802.11n D1.10)
262        0x2d, 0x1a,
263        0xef, 0x09, // HT Capabilities Info
264        0x1b, // A-MPDU Parameters: 0x1b
265        0xff, 0xff, 0xff, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, // MCS Set
266        0x00, 0x00, // HT Extended Capabilities
267        0x00, 0x00, 0x00, 0x00, // Transmit Beamforming Capabilities
268        0x00, // Antenna Selection Capabilities
269        // RSN Information
270        0x30, 0x14, 0x01, 0x00,
271        0x00, 0x0f, 0xac, 0x04, // Group Cipher: AES (CCM)
272        0x01, 0x00, 0x00, 0x0f, 0xac, 0x04, // Pairwise Cipher: AES (CCM)
273        0x01, 0x00, 0x00, 0x0f, 0xac, 0x02, // AKM: PSK
274        0x00, 0x00, // RSN Capabilities
275        // HT Information (802.11n D1.10)
276        0x3d, 0x16,
277        0x9d, // Primary Channel: 157
278        0x0d, // HT Info Subset - secondary channel above, any channel width, RIFS permitted
279        0x00, 0x00, 0x00, 0x00, // HT Info Subsets
280        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // Basic MCS Set
281        // Overlapping BSS Scan Parameters
282        0x4a, 0x0e, 0x14, 0x00, 0x0a, 0x00, 0x2c, 0x01, 0xc8, 0x00, 0x14, 0x00, 0x05, 0x00, 0x19, 0x00,
283        // Extended Capabilities
284        0x7f, 0x08, 0x01, 0x00, 0x0f, 0x00, 0x00, 0x00, 0x00, 0x40,
285        // VHT Capabilities
286        0xbf, 0x0c,
287        0xb2, 0x01, 0x80, 0x33, // VHT Capabilities Info
288        0xea, 0xff, 0x00, 0x00, 0xea, 0xff, 0x00, 0x00, // VHT Supported MCS Set
289        // VHT Operation
290        0xc0, 0x05, 0x01, 0x9b, 0x00, 0xfc, 0xff,
291        // VHT Tx Power Envelope
292        0xc3, 0x04, 0x02, 0xc4, 0xc4, 0xc4,
293        // Vendor Specific: Atheros Communications, Inc.: Advanced Capability
294        0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f,
295        // Vendor Specific: Microsoft Corp.: WMM/WME: Parameter Element
296        0xdd, 0x18, 0x00, 0x50, 0xf2, 0x02, 0x01, 0x01,
297        0x80, // U-APSD enabled
298        0x00,
299        0x03, 0xa4, 0x00, 0x00, // AC_BE parameters
300        0x27, 0xa4, 0x00, 0x00, // AC_BK parameters
301        0x42, 0x43, 0x5e, 0x00, // AC_VI parameters
302        0x62, 0x32, 0x2f, 0x00, // AC_VO parameters
303        // Vendor Specific: Microsoft Corp.: WPS
304        0xdd, 0x85, 0x00, 0x50, 0xf2, 0x04, 0x10, 0x4a, 0x00, 0x01, 0x10, 0x10, 0x44, 0x00, 0x01,
305        0x02, 0x10, 0x3b, 0x00, 0x01, 0x03, 0x10, 0x47, 0x00, 0x10, 0x87, 0x65, 0x43, 0x21, 0x9a,
306        0xbc, 0xde, 0xf0, 0x12, 0x34, 0x98, 0xda, 0xc4, 0x8e, 0x86, 0xc8, 0x10, 0x21, 0x00, 0x07,
307        0x54, 0x50, 0x2d, 0x4c, 0x69, 0x6e, 0x6b, 0x10, 0x23, 0x00, 0x09, 0x41, 0x72, 0x63, 0x68,
308        0x65, 0x72, 0x20, 0x41, 0x37, 0x10, 0x24, 0x00, 0x03, 0x31, 0x2e, 0x30, 0x10, 0x42, 0x00,
309        0x0c, 0x41, 0x72, 0x63, 0x68, 0x65, 0x72, 0x20, 0x41, 0x37, 0x20, 0x76, 0x35, 0x10, 0x54,
310        0x00, 0x08, 0x00, 0x06, 0x00, 0x50, 0xf2, 0x04, 0x00, 0x01, 0x10, 0x11, 0x00, 0x0a, 0x41,
311        0x72, 0x63, 0x68, 0x65, 0x72, 0x41, 0x37, 0x76, 0x35, 0x10, 0x08, 0x00, 0x02, 0x00, 0x04,
312        0x10, 0x3c, 0x00, 0x01, 0x03, 0x10, 0x49, 0x00, 0x06, 0x00, 0x37, 0x2a, 0x00, 0x01, 0x20,
313    ];
314
315    #[rustfmt::skip]
316    const MERGED_IES: &'static [u8] = &[
317        // SSID: "foo-ssid"
318        0x00, 0x08, 0x66, 0x6f, 0x6f, 0x2d, 0x73, 0x73, 0x69, 0x64,
319        // Supported Rates: 6(B), 9, 12(B), 18, 24(B), 36, 48, 54, [Mbit/sec]
320        0x01, 0x08, 0x8c, 0x12, 0x98, 0x24, 0xb0, 0x48, 0x60, 0x6c,
321        // DS Parameter set: Current Channel: 157
322        0x03, 0x01, 0x9d,
323        // Traffic Indication Map (TIM): DTIM 0 of 0 bitmap
324        0x05, 0x04, 0x00, 0x01, 0x00, 0x00,
325        // Power Constraint: 3
326        0x20, 0x01, 0x03,
327        // HT Capabilities (802.11n D1.10)
328        0x2d, 0x1a,
329        0xef, 0x09, // HT Capabilities Info
330        0x1b, // A-MPDU Parameters: 0x1b
331        0xff, 0xff, 0xff, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, // MCS Set
332        0x00, 0x00, // HT Extended Capabilities
333        0x00, 0x00, 0x00, 0x00, // Transmit Beamforming Capabilities
334        0x00, // Antenna Selection Capabilities
335        // RSN Information
336        0x30, 0x14, 0x01, 0x00,
337        0x00, 0x0f, 0xac, 0x04, // Group Cipher: AES (CCM)
338        0x01, 0x00, 0x00, 0x0f, 0xac, 0x04, // Pairwise Cipher: AES (CCM)
339        0x01, 0x00, 0x00, 0x0f, 0xac, 0x02, // AKM: PSK
340        0x00, 0x00, // RSN Capabilities
341        // HT Information (802.11n D1.10)
342        0x3d, 0x16,
343        0x9d, // Primary Channel: 157
344        0x0d, // HT Info Subset - secondary channel above, any channel width, RIFS permitted
345        0x00, 0x00, 0x00, 0x00, // HT Info Subsets
346        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // Basic MCS Set
347        // Overlapping BSS Scan Parameters
348        0x4a, 0x0e, 0x14, 0x00, 0x0a, 0x00, 0x2c, 0x01, 0xc8, 0x00, 0x14, 0x00, 0x05, 0x00, 0x19, 0x00,
349        // Extended Capabilities
350        0x7f, 0x08, 0x01, 0x00, 0x0f, 0x00, 0x00, 0x00, 0x00, 0x40,
351        // VHT Capabilities
352        0xbf, 0x0c,
353        0xb2, 0x01, 0x80, 0x33, // VHT Capabilities Info
354        0xea, 0xff, 0x00, 0x00, 0xea, 0xff, 0x00, 0x00, // VHT Supported MCS Set
355        // VHT Operation
356        0xc0, 0x05, 0x01, 0x9b, 0x00, 0xfc, 0xff,
357        // VHT Tx Power Envelope
358        0xc3, 0x04, 0x02, 0xc4, 0xc4, 0xc4,
359        // Vendor Specific: Atheros Communications, Inc.: Advanced Capability
360        0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f,
361        // Vendor Specific: Microsoft Corp.: WMM/WME: Parameter Element
362        0xdd, 0x18, 0x00, 0x50, 0xf2, 0x02, 0x01, 0x01,
363        0x80, // U-APSD enabled
364        0x00,
365        0x03, 0xa4, 0x00, 0x00, // AC_BE parameters
366        0x27, 0xa4, 0x00, 0x00, // AC_BK parameters
367        0x42, 0x43, 0x5e, 0x00, // AC_VI parameters
368        0x62, 0x32, 0x2f, 0x00, // AC_VO parameters
369        // Vendor Specific: Microsoft Corp.: WPS
370        0xdd, 0x85, 0x00, 0x50, 0xf2, 0x04, 0x10, 0x4a, 0x00, 0x01, 0x10, 0x10, 0x44, 0x00, 0x01,
371        0x02, 0x10, 0x3b, 0x00, 0x01, 0x03, 0x10, 0x47, 0x00, 0x10, 0x87, 0x65, 0x43, 0x21, 0x9a,
372        0xbc, 0xde, 0xf0, 0x12, 0x34, 0x98, 0xda, 0xc4, 0x8e, 0x86, 0xc8, 0x10, 0x21, 0x00, 0x07,
373        0x54, 0x50, 0x2d, 0x4c, 0x69, 0x6e, 0x6b, 0x10, 0x23, 0x00, 0x09, 0x41, 0x72, 0x63, 0x68,
374        0x65, 0x72, 0x20, 0x41, 0x37, 0x10, 0x24, 0x00, 0x03, 0x31, 0x2e, 0x30, 0x10, 0x42, 0x00,
375        0x0c, 0x41, 0x72, 0x63, 0x68, 0x65, 0x72, 0x20, 0x41, 0x37, 0x20, 0x76, 0x35, 0x10, 0x54,
376        0x00, 0x08, 0x00, 0x06, 0x00, 0x50, 0xf2, 0x04, 0x00, 0x01, 0x10, 0x11, 0x00, 0x0a, 0x41,
377        0x72, 0x63, 0x68, 0x65, 0x72, 0x41, 0x37, 0x76, 0x35, 0x10, 0x08, 0x00, 0x02, 0x00, 0x04,
378        0x10, 0x3c, 0x00, 0x01, 0x03, 0x10, 0x49, 0x00, 0x06, 0x00, 0x37, 0x2a, 0x00, 0x01, 0x20,
379    ];
380
381    #[test]
382    fn test_no_merge() {
383        let ies = IesMerger::new(BEACON_FRAME_IES.to_vec()).finalize();
384        assert_eq!(&ies[..], BEACON_FRAME_IES);
385    }
386
387    #[test]
388    fn test_merge_same_ies() {
389        let mut ies_merger = IesMerger::new(BEACON_FRAME_IES.to_vec());
390        ies_merger.merge(BEACON_FRAME_IES);
391        let ies = ies_merger.finalize();
392        assert_eq!(&ies[..], BEACON_FRAME_IES);
393        assert!(!ies_merger.buffer_overflow());
394    }
395
396    #[test]
397    fn test_merge_different_ies() {
398        let mut ies_merger = IesMerger::new(BEACON_FRAME_IES.to_vec());
399        ies_merger.merge(PROBE_RESP_IES);
400        let ies = ies_merger.finalize();
401        assert_eq!(&ies[..], MERGED_IES);
402        assert!(!ies_merger.buffer_overflow());
403    }
404
405    // Verify that merging IEs in a different order should still produce the same result
406    // (this is not true in the general case, but is true with the data we mock)
407    #[test]
408    fn test_merge_different_ies_commutative() {
409        let mut ies_merger1 = IesMerger::new(BEACON_FRAME_IES.to_vec());
410        ies_merger1.merge(PROBE_RESP_IES);
411        let ies1 = ies_merger1.finalize();
412
413        let mut ies_merger2 = IesMerger::new(PROBE_RESP_IES.to_vec());
414        ies_merger2.merge(BEACON_FRAME_IES);
415        let ies2 = ies_merger2.finalize();
416
417        assert_eq!(ies1, ies2);
418    }
419
420    // In the course of a scan, we may merge the same thing multiple times.
421    // This tests such use case.
422    #[test]
423    fn test_merge_redundant() {
424        let mut ies_merger = IesMerger::new(BEACON_FRAME_IES.to_vec());
425        ies_merger.merge(BEACON_FRAME_IES);
426        ies_merger.merge(PROBE_RESP_IES);
427        ies_merger.merge(PROBE_RESP_IES);
428        ies_merger.merge(BEACON_FRAME_IES);
429        let ies = ies_merger.finalize();
430
431        assert_eq!(&ies[..], MERGED_IES);
432    }
433
434    // Test that we don't overuse memory if an unfriendly AP repeatedly updates the IEs.
435    #[test]
436    fn test_merge_is_resilient() {
437        const LEN: u8 = 255;
438        let mut ies = vec![0xdd, LEN];
439        ies.extend(vec![0xff; LEN as usize].iter());
440
441        let mut ies_merger = IesMerger::new(ies.clone());
442        for i in 0..255 {
443            let last_idx = ies.len() - 1;
444            ies[last_idx] = i; // Tweak one value so the IEs is different, forcing a merge.
445            ies_merger.merge(&ies[..]);
446        }
447        // Verify we don't use too much memory.
448        assert!(ies_merger.ies_updater.buf_len() <= IES_MERGER_BUFFER_LIMIT);
449        // Verify buffer overflow flag is set to true.
450        assert!(ies_merger.buffer_overflow());
451
452        // We should still produce a result.
453        let result_ies = ies_merger.finalize();
454        assert_eq!(ies.len(), result_ies.len());
455        // Verify one IE picked (whichever IE is picked, all but the last byte are the same).
456        assert_eq!(ies[..ies.len() - 1], result_ies[..result_ies.len() - 1]);
457    }
458
459    #[test]
460    fn test_merge_prioritize_non_hidden_ssid_1() {
461        let hidden_ssid_ie = vec![0x00, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];
462        let mut ies_merger = IesMerger::new(BEACON_FRAME_IES.to_vec());
463        ies_merger.merge(&hidden_ssid_ie[..]);
464        let ies = ies_merger.finalize();
465
466        // The non-hidden SSID IE is kept.
467        assert_eq!(&ies[..], BEACON_FRAME_IES);
468
469        // Sanity check by initializing with hidden SSID first.
470        let mut ies_merger = IesMerger::new(hidden_ssid_ie.clone());
471        // This verifies that hidden SSID IE is kept if it's the only one.
472        assert_eq!(ies_merger.finalize(), hidden_ssid_ie);
473
474        let mut ies_merger = IesMerger::new(hidden_ssid_ie);
475        ies_merger.merge(BEACON_FRAME_IES);
476        let ies = ies_merger.finalize();
477
478        // The non-hidden SSID IE is kept.
479        assert_eq!(&ies[..], BEACON_FRAME_IES);
480    }
481
482    #[test]
483    fn test_ie_updater_get() {
484        let ies = vec![
485            0, 2, 10, 20, // IE with no extension ID
486            0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, // Vendor IE
487            255, 2, 5, 1, // IE with extension ID
488        ];
489        let ies_updater = IesUpdater::new(ies);
490
491        assert_eq!(ies_updater.get(&IeType::SSID), Some(&[10, 20][..]));
492        assert_eq!(
493            ies_updater.get(&IeType::new_vendor([0x00, 0x03, 0x7f, 0x01, 0x01, 0x00])),
494            Some(&[0x00, 0xff, 0x7f][..])
495        );
496        assert_eq!(ies_updater.get(&IeType::new_extended(5)), Some(&[1][..]));
497
498        assert_eq!(ies_updater.get(&IeType::SUPPORTED_RATES), None);
499        assert_eq!(
500            ies_updater.get(&IeType::new_vendor([0x00, 0x03, 0x7f, 0x01, 0x01, 0x01])),
501            None
502        );
503        assert_eq!(ies_updater.get(&IeType::new_extended(6)), None);
504    }
505
506    #[test]
507    fn test_ie_updater_set_replace() {
508        let ies = vec![
509            0, 2, 10, 20, // IE with no extension ID
510            0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, // Vendor IE
511            255, 2, 5, 1, // IE with extension ID
512        ];
513        let mut ies_updater = IesUpdater::new(ies);
514
515        ies_updater.set(IeType::SSID, &[30, 40, 50]).expect("set basic succeeds");
516        ies_updater
517            .set(IeType::new_vendor([0x00, 0x03, 0x7f, 0x01, 0x01, 0x00]), &[1, 3, 3, 7])
518            .expect("set vendor succeeds");
519        ies_updater.set(IeType::new_extended(5), &[4, 2]).expect("set extended succeeds");
520
521        assert_eq!(ies_updater.get(&IeType::SSID), Some(&[30, 40, 50][..]));
522        assert_eq!(
523            ies_updater.get(&IeType::new_vendor([0x00, 0x03, 0x7f, 0x01, 0x01, 0x00])),
524            Some(&[1, 3, 3, 7][..])
525        );
526        assert_eq!(ies_updater.get(&IeType::new_extended(5)), Some(&[4, 2][..]));
527
528        assert_eq!(
529            &ies_updater.finalize()[..],
530            &[
531                0, 3, 30, 40, 50, // IE with no extension ID
532                0xdd, 0x0a, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 1, 3, 3, 7, // Vendor IE
533                255, 3, 5, 4, 2, // IE with extension ID
534            ]
535        );
536    }
537
538    #[test]
539    fn test_ie_updater_set_new() {
540        let ies = vec![
541            0, 2, 10, 20, // IE with no extension ID
542            0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, // Vendor IE
543            255, 2, 5, 1, // IE with extension ID
544        ];
545        let mut ies_updater = IesUpdater::new(ies);
546
547        ies_updater.set(IeType::SUPPORTED_RATES, &[30, 40, 50]).expect("set basic succeeds");
548        ies_updater
549            .set(IeType::new_vendor([0x00, 0x03, 0x7f, 0x01, 0x01, 0x01]), &[1, 3, 3, 7])
550            .expect("set vendor succeeds");
551        ies_updater.set(IeType::new_extended(6), &[4, 2]).expect("set extended succeeds");
552
553        assert_eq!(ies_updater.get(&IeType::SUPPORTED_RATES), Some(&[30, 40, 50][..]));
554        assert_eq!(
555            ies_updater.get(&IeType::new_vendor([0x00, 0x03, 0x7f, 0x01, 0x01, 0x01])),
556            Some(&[1, 3, 3, 7][..])
557        );
558        assert_eq!(ies_updater.get(&IeType::new_extended(6)), Some(&[4, 2][..]));
559
560        assert_eq!(
561            &ies_updater.finalize()[..],
562            &[
563                0, 2, 10, 20, // IE with no extension ID
564                1, 3, 30, 40, 50, // New IE with no extension ID
565                0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, // Vendor IE
566                0xdd, 0x0a, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x01, 1, 3, 3, 7, // New vendor IE
567                255, 2, 5, 1, // IE with extension ID
568                255, 3, 6, 4, 2, // New IE with extension ID
569            ]
570        )
571    }
572
573    #[test]
574    fn test_ie_updater_set_ie_too_large() {
575        let ies = vec![
576            0, 2, 10, 20, // IE with no extension ID
577            0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, // Vendor IE
578            255, 2, 5, 1, // IE with extension ID
579        ];
580        let mut ies_updater = IesUpdater::new(ies);
581        ies_updater.set(IeType::SSID, &[11; 256]).expect_err("set basic fails");
582        ies_updater
583            .set(IeType::new_vendor([0x00, 0x03, 0x7f, 0x01, 0x01, 0x00]), &[11; 250])
584            .expect_err("set vendor fails");
585        ies_updater.set(IeType::new_extended(5), &[11; 255]).expect_err("set extended fails");
586
587        // None of the IEs got replaced because all the IEs set in this test are too large
588        assert_eq!(
589            &ies_updater.finalize()[..],
590            &[
591                0, 2, 10, 20, // IE with no extension ID
592                0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, // Vendor IE
593                255, 2, 5, 1, // IE with extension ID
594            ]
595        );
596    }
597
598    #[test]
599    fn test_ie_updater_set_raw_ie() {
600        let ies = vec![];
601
602        let mut ies_updater = IesUpdater::new(ies.clone());
603        ies_updater.set_raw(&[0, 2, 10, 20]).expect("set right length succeeds");
604        ies_updater.set_raw(&[1, 2, 70]).expect_err("set buffer too small fails");
605        ies_updater.set_raw(&[]).expect("set empty doesn't return error");
606        ies_updater.set_raw(&[2, 2, 30, 40, 50, 60]).expect("set truncated succeeds");
607
608        assert_eq!(&ies_updater.finalize()[..], &[0, 2, 10, 20, 2, 2, 30, 40])
609    }
610
611    #[test]
612    fn test_ie_updater_remove() {
613        let ies = vec![
614            0, 2, 10, 20, // IE with no extension ID
615            0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, // Vendor IE
616            255, 2, 5, 1, // IE with extension ID
617        ];
618
619        let mut ies_updater = IesUpdater::new(ies.clone());
620        ies_updater.remove(&IeType::SSID);
621        assert_eq!(ies_updater.get(&IeType::SSID), None);
622        assert_eq!(
623            ies_updater.finalize(),
624            &[
625                0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, // Vendor IE
626                255, 2, 5, 1, // IE with extension ID
627            ]
628        );
629
630        let mut ies_updater = IesUpdater::new(ies.clone());
631        ies_updater.remove(&IeType::new_vendor([0x00, 0x03, 0x7f, 0x01, 0x01, 0x00]));
632        assert_eq!(
633            ies_updater.get(&IeType::new_vendor([0x00, 0x03, 0x7f, 0x01, 0x01, 0x00])),
634            None
635        );
636        assert_eq!(
637            ies_updater.finalize(),
638            &[
639                0, 2, 10, 20, // IE with no extension ID
640                255, 2, 5, 1, // IE with extension ID
641            ],
642        );
643
644        let mut ies_updater = IesUpdater::new(ies);
645        ies_updater.remove(&IeType::new_extended(5));
646        assert_eq!(ies_updater.get(&IeType::new_extended(5)), None);
647        assert_eq!(
648            ies_updater.finalize(),
649            &[
650                0, 2, 10, 20, // IE with no extension ID
651                0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, // Vendor IE
652            ],
653        );
654    }
655}
wlan_common/ie/merger.rs

wlan_common/ie/
merger.rs
