wlan_common/ie/
merger.rs

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, 0x05, 0x50, 0x6f, 0x9a, 0x1c, 0xaa, // Vendor IE with known header
487            0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, // Vendor IE
488            255, 2, 5, 1, // IE with extension ID
489        ];
490        let ies_updater = IesUpdater::new(ies);
491
492        assert_eq!(ies_updater.get(&IeType::SSID), Some(&[10, 20][..]));
493        assert_eq!(
494            ies_updater.get(&IeType::new_vendor4([0x50, 0x6f, 0x9a, 0x1c])),
495            Some(&[0xaa][..])
496        );
497        assert_eq!(
498            ies_updater.get(&IeType::new_vendor6([0x00, 0x03, 0x7f, 0x01, 0x01, 0x00])),
499            Some(&[0x00, 0xff, 0x7f][..])
500        );
501        assert_eq!(ies_updater.get(&IeType::new_extended(5)), Some(&[1][..]));
502
503        assert_eq!(ies_updater.get(&IeType::SUPPORTED_RATES), None);
504        assert_eq!(ies_updater.get(&IeType::new_vendor4([0x50, 0x6f, 0x9a, 0x1d])), None);
505        assert_eq!(
506            ies_updater.get(&IeType::new_vendor6([0x00, 0x03, 0x7f, 0x01, 0x01, 0x01])),
507            None
508        );
509        assert_eq!(ies_updater.get(&IeType::new_extended(6)), None);
510    }
511
512    #[test]
513    fn test_ie_updater_set_replace() {
514        let ies = vec![
515            0, 2, 10, 20, // IE with no extension ID
516            0xdd, 0x05, 0x50, 0x6f, 0x9a, 0x1c, 0xaa, // Vendor IE with known header
517            0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, // Vendor IE
518            255, 2, 5, 1, // IE with extension ID
519        ];
520        let mut ies_updater = IesUpdater::new(ies);
521
522        ies_updater.set(IeType::SSID, &[30, 40, 50]).expect("set basic succeeds");
523        ies_updater
524            .set(IeType::new_vendor4([0x50, 0x6f, 0x9a, 0x1c]), &[0xbb])
525            .expect("set vendor succeeds");
526        ies_updater
527            .set(IeType::new_vendor6([0x00, 0x03, 0x7f, 0x01, 0x01, 0x00]), &[1, 3, 3, 7])
528            .expect("set vendor succeeds");
529        ies_updater.set(IeType::new_extended(5), &[4, 2]).expect("set extended succeeds");
530
531        assert_eq!(ies_updater.get(&IeType::SSID), Some(&[30, 40, 50][..]));
532        assert_eq!(
533            ies_updater.get(&IeType::new_vendor4([0x50, 0x6f, 0x9a, 0x1c])),
534            Some(&[0xbb][..])
535        );
536        assert_eq!(
537            ies_updater.get(&IeType::new_vendor6([0x00, 0x03, 0x7f, 0x01, 0x01, 0x00])),
538            Some(&[1, 3, 3, 7][..])
539        );
540        assert_eq!(ies_updater.get(&IeType::new_extended(5)), Some(&[4, 2][..]));
541
542        assert_eq!(
543            &ies_updater.finalize()[..],
544            &[
545                0, 3, 30, 40, 50, // IE with no extension ID
546                0xdd, 0x0a, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 1, 3, 3, 7, // Vendor IE
547                0xdd, 0x05, 0x50, 0x6f, 0x9a, 0x1c, 0xbb, // Vendor IE with known header
548                255, 3, 5, 4, 2, // IE with extension ID
549            ]
550        );
551    }
552
553    #[test]
554    fn test_ie_updater_set_new() {
555        let ies = vec![
556            0, 2, 10, 20, // IE with no extension ID
557            0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, // Vendor IE
558            255, 2, 5, 1, // IE with extension ID
559        ];
560        let mut ies_updater = IesUpdater::new(ies);
561
562        ies_updater.set(IeType::SUPPORTED_RATES, &[30, 40, 50]).expect("set basic succeeds");
563        ies_updater
564            .set(IeType::new_vendor6([0x00, 0x03, 0x7f, 0x01, 0x01, 0x01]), &[1, 3, 3, 7])
565            .expect("set vendor succeeds");
566        ies_updater.set(IeType::new_extended(6), &[4, 2]).expect("set extended succeeds");
567
568        assert_eq!(ies_updater.get(&IeType::SUPPORTED_RATES), Some(&[30, 40, 50][..]));
569        assert_eq!(
570            ies_updater.get(&IeType::new_vendor6([0x00, 0x03, 0x7f, 0x01, 0x01, 0x01])),
571            Some(&[1, 3, 3, 7][..])
572        );
573        assert_eq!(ies_updater.get(&IeType::new_extended(6)), Some(&[4, 2][..]));
574
575        assert_eq!(
576            &ies_updater.finalize()[..],
577            &[
578                0, 2, 10, 20, // IE with no extension ID
579                1, 3, 30, 40, 50, // New IE with no extension ID
580                0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, // Vendor IE
581                0xdd, 0x0a, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x01, 1, 3, 3, 7, // New vendor IE
582                255, 2, 5, 1, // IE with extension ID
583                255, 3, 6, 4, 2, // New IE with extension ID
584            ]
585        )
586    }
587
588    #[test]
589    fn test_ie_updater_set_ie_too_large() {
590        let ies = vec![
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        let mut ies_updater = IesUpdater::new(ies);
596        ies_updater.set(IeType::SSID, &[11; 256]).expect_err("set basic fails");
597        ies_updater
598            .set(IeType::new_vendor6([0x00, 0x03, 0x7f, 0x01, 0x01, 0x00]), &[11; 250])
599            .expect_err("set vendor fails");
600        ies_updater.set(IeType::new_extended(5), &[11; 255]).expect_err("set extended fails");
601
602        // None of the IEs got replaced because all the IEs set in this test are too large
603        assert_eq!(
604            &ies_updater.finalize()[..],
605            &[
606                0, 2, 10, 20, // IE with no extension ID
607                0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, // Vendor IE
608                255, 2, 5, 1, // IE with extension ID
609            ]
610        );
611    }
612
613    #[test]
614    fn test_ie_updater_set_raw_ie() {
615        let ies = vec![];
616
617        let mut ies_updater = IesUpdater::new(ies.clone());
618        ies_updater.set_raw(&[0, 2, 10, 20]).expect("set right length succeeds");
619        ies_updater.set_raw(&[1, 2, 70]).expect_err("set buffer too small fails");
620        ies_updater.set_raw(&[]).expect("set empty doesn't return error");
621        ies_updater.set_raw(&[2, 2, 30, 40, 50, 60]).expect("set truncated succeeds");
622
623        assert_eq!(&ies_updater.finalize()[..], &[0, 2, 10, 20, 2, 2, 30, 40])
624    }
625
626    #[test]
627    fn test_ie_updater_remove() {
628        let ies = vec![
629            0, 2, 10, 20, // IE with no extension ID
630            0xdd, 0x05, 0x50, 0x6f, 0x9a, 0x1c, 0xaa, // Vendor IE with known header
631            0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, // Vendor IE
632            255, 2, 5, 1, // IE with extension ID
633        ];
634
635        let mut ies_updater = IesUpdater::new(ies.clone());
636        ies_updater.remove(&IeType::SSID);
637        assert_eq!(ies_updater.get(&IeType::SSID), None);
638        assert_eq!(
639            ies_updater.finalize(),
640            &[
641                0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, // Vendor IE
642                0xdd, 0x05, 0x50, 0x6f, 0x9a, 0x1c, 0xaa, // Vendor IE with known header
643                255, 2, 5, 1, // IE with extension ID
644            ]
645        );
646
647        let mut ies_updater = IesUpdater::new(ies.clone());
648        ies_updater.remove(&IeType::new_vendor4([0x50, 0x6f, 0x9a, 0x1c]));
649        assert_eq!(ies_updater.get(&IeType::new_vendor4([0x50, 0x6f, 0x9a, 0x1c])), None);
650        assert_eq!(
651            ies_updater.finalize(),
652            &[
653                0, 2, 10, 20, // IE with no extension ID
654                0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, // Vendor IE
655                255, 2, 5, 1, // IE with extension ID
656            ],
657        );
658
659        let mut ies_updater = IesUpdater::new(ies.clone());
660        ies_updater.remove(&IeType::new_vendor6([0x00, 0x03, 0x7f, 0x01, 0x01, 0x00]));
661        assert_eq!(
662            ies_updater.get(&IeType::new_vendor6([0x00, 0x03, 0x7f, 0x01, 0x01, 0x00])),
663            None
664        );
665        assert_eq!(
666            ies_updater.finalize(),
667            &[
668                0, 2, 10, 20, // IE with no extension ID
669                0xdd, 0x05, 0x50, 0x6f, 0x9a, 0x1c, 0xaa, // Vendor IE with known header
670                255, 2, 5, 1, // IE with extension ID
671            ],
672        );
673
674        let mut ies_updater = IesUpdater::new(ies);
675        ies_updater.remove(&IeType::new_extended(5));
676        assert_eq!(ies_updater.get(&IeType::new_extended(5)), None);
677        assert_eq!(
678            ies_updater.finalize(),
679            &[
680                0, 2, 10, 20, // IE with no extension ID
681                0xdd, 0x09, 0x00, 0x03, 0x7f, 0x01, 0x01, 0x00, 0x00, 0xff, 0x7f, // Vendor IE
682                0xdd, 0x05, 0x50, 0x6f, 0x9a, 0x1c, 0xaa // Vendor IE with known header
683            ],
684        );
685    }
686}
wlan_common/ie/merger.rs

wlan_common/ie/
merger.rs
