1use super::Element;
6use crate::ProtectionInfo;
7use bitfield::bitfield;
8use nom::bytes::streaming::take;
9use nom::combinator::{eof, map};
10use nom::error::{Error, ErrorKind};
11use nom::number::streaming::{le_u8, le_u16};
12use nom::sequence::terminated;
13use nom::{IResult, Parser};
14use wlan_common::append::{BufferTooSmall, TrackedAppend, VecCursor};
15use wlan_common::ie::{wpa, write_wpa1_ie};
16use wlan_common::organization::Oui;
17
18pub const TYPE: u8 = 0xDD;
19const PADDING_DATA_LEN: u8 = 0;
20const HDR_LEN: usize = 6;
21const HDR_OUI_TYPE_LEN: usize = 4;
23
24const GTK_DATA_TYPE: u8 = 1;
26const IGTK_DATA_TYPE: u8 = 9;
27
28const GTK_FIXED_LEN: usize = 2;
31
32const IGTK_IPN_LEN: usize = 6;
33const IGTK_FIXED_LEN: usize = 2 + IGTK_IPN_LEN;
34
35#[derive(Default, Debug, PartialEq)]
37pub struct Header {
38 pub type_: u8,
39 pub len: u8,
40 pub oui: Oui,
41 pub data_type: u8,
42}
43
44impl Header {
45 pub fn new(type_: u8, len: u8, oui: &[u8], data_type: u8) -> Header {
46 let mut oui_buf = [0u8; 3];
47 oui_buf.copy_from_slice(oui);
48 Header { type_, len, data_type, oui: Oui::new(oui_buf), ..Default::default() }
49 }
50
51 pub fn new_dot11(data_type: u8, data_len: usize) -> Header {
52 Header { type_: TYPE, len: (HDR_OUI_TYPE_LEN + data_len) as u8, data_type, oui: Oui::DOT11 }
53 }
54
55 fn data_len(&self) -> usize {
56 if self.len < 4 { 0 } else { (self.len as usize) - 4 }
57 }
58}
59
60pub enum GtkInfoTx {
62 _OnlyRx = 0,
63 BothRxTx = 1,
64}
65
66bitfield! {
68 pub struct GtkInfo(u8);
69 impl Debug;
70 pub key_id, set_key_id: 1, 0;
71 pub tx, set_tx: 2, 2;
72 pub value, _: 7,0;
74}
75
76impl PartialEq for GtkInfo {
77 fn eq(&self, other: &Self) -> bool {
78 self.0 == other.0
79 }
80}
81
82#[derive(Debug, PartialEq)]
85pub struct Gtk {
86 pub info: GtkInfo,
87 pub gtk: Box<[u8]>,
89}
90
91impl Gtk {
92 pub fn new(key_id: u8, tx: GtkInfoTx, gtk: &[u8]) -> Self {
93 let mut gtk_info = GtkInfo(0);
94 gtk_info.set_key_id(key_id);
95 gtk_info.set_tx(tx as u8);
96 Self { info: gtk_info, gtk: gtk.into() }
97 }
98
99 pub fn len(&self) -> usize {
101 GTK_FIXED_LEN + self.gtk.len()
102 }
103}
104
105#[derive(Debug, PartialEq)]
108pub struct Igtk {
109 pub id: u16,
110 pub ipn: [u8; IGTK_IPN_LEN],
112 pub igtk: Vec<u8>,
113}
114
115impl Igtk {
116 pub fn new(id: u16, ipn: &[u8], igtk: &[u8]) -> Self {
117 let mut ipn_buf = [0u8; IGTK_IPN_LEN];
118 ipn_buf.copy_from_slice(ipn);
119 Self { id, ipn: ipn_buf, igtk: igtk.to_vec() }
120 }
121
122 pub fn len(&self) -> usize {
123 IGTK_FIXED_LEN + self.igtk.len()
124 }
125}
126
127pub fn parse(i0: &[u8]) -> IResult<&[u8], Element> {
128 if i0.len() <= 1 {
130 return Ok((&i0[i0.len()..], Element::Padding));
131 }
132
133 let data_len = i0[1];
135 if data_len == PADDING_DATA_LEN {
136 return parse_padding(&i0[1..]);
137 }
138
139 let (i1, hdr) = parse_header(i0)?;
141 let (i2, bytes) = take(hdr.data_len()).parse(i1)?;
142 match hdr.oui {
143 Oui::DOT11 => match hdr.data_type {
144 GTK_DATA_TYPE => {
145 let (_, gtk) = parse_gtk(hdr.data_len()).parse(bytes)?;
146 Ok((i2, Element::Gtk(hdr, gtk)))
147 }
148 IGTK_DATA_TYPE => {
149 let (_, gtk) = parse_igtk(hdr.data_len()).parse(bytes)?;
150 Ok((i2, Element::Igtk(hdr, gtk)))
151 }
152 _ => Ok((i2, Element::UnsupportedKde(hdr))),
153 },
154 Oui::MSFT if hdr.data_type == wpa::VENDOR_SPECIFIC_TYPE => {
157 let (_, wpa) = wpa::from_bytes(&bytes[..])?;
158 Ok((i2, Element::LegacyWpa1(wpa)))
159 }
160 _ => Ok((i2, Element::UnsupportedKde(hdr))),
161 }
162}
163
164fn parse_header(input: &[u8]) -> IResult<&[u8], Header> {
165 map((le_u8, le_u8, take(3usize), le_u8), |(header_type, length, oui, data_type)| {
166 Header::new(header_type, length, oui, data_type)
167 })
168 .parse(input)
169}
170
171fn parse_padding(input: &[u8]) -> IResult<&[u8], Element> {
172 if input.iter().all(|&x| x == 0) {
173 Ok((&[], Element::Padding))
174 } else {
175 Err(nom::Err::Error(Error::new(input, ErrorKind::Eof)))
178 }
179}
180
181fn parse_gtk<'a>(data_len: usize) -> impl Parser<&'a [u8], Output = Gtk, Error = Error<&'a [u8]>> {
182 map(
183 terminated(
184 (
185 map(le_u8::<&'a [u8], Error<_>>, GtkInfo),
186 take(1usize),
187 take(data_len - 2),
188 ),
189 eof,
190 ),
191 |(info, _reserved, gtk)| Gtk { info, gtk: Vec::from(gtk).into_boxed_slice() },
192 )
193}
194
195fn parse_igtk<'a>(
196 data_len: usize,
197) -> impl Parser<&'a [u8], Output = Igtk, Error = Error<&'a [u8]>> {
198 map(
199 terminated(
200 (le_u16::<&'a [u8], Error<_>>, take(IGTK_IPN_LEN), take(data_len - IGTK_FIXED_LEN)),
201 eof,
202 ),
203 |(id, ipn, igtk)| Igtk::new(id, ipn, igtk),
204 )
205}
206
207pub struct Writer<A> {
209 buf: A,
210}
211
212impl Writer<VecCursor> {
213 pub fn new() -> Self {
214 Self { buf: VecCursor::new() }
215 }
216}
217
218impl<A: TrackedAppend> Writer<A> {
219 #[cfg(test)]
220 pub fn new_with(buf: A) -> Self {
221 Self { buf }
222 }
223
224 pub fn write_protection(&mut self, protection: &ProtectionInfo) -> Result<(), BufferTooSmall> {
225 match protection {
226 ProtectionInfo::Rsne(rsne) => rsne.write_into(&mut self.buf),
227 ProtectionInfo::LegacyWpa(wpa) => write_wpa1_ie(&mut self.buf, wpa),
228 }
229 }
230
231 fn write_kde_hdr(&mut self, hdr: Header) -> Result<(), BufferTooSmall> {
232 if !self.buf.can_append(HDR_LEN) {
233 return Err(BufferTooSmall);
234 }
235 self.buf.append_byte(hdr.type_)?;
236 self.buf.append_byte(hdr.len)?;
237 self.buf.append_bytes(&hdr.oui[..])?;
238 self.buf.append_byte(hdr.data_type)?;
239 Ok(())
240 }
241
242 pub fn write_gtk(&mut self, gtk_kde: &Gtk) -> Result<(), BufferTooSmall> {
243 if !self.buf.can_append(HDR_LEN + gtk_kde.len()) {
244 return Err(BufferTooSmall);
245 }
246 let hdr = Header::new_dot11(GTK_DATA_TYPE, gtk_kde.len());
248 self.write_kde_hdr(hdr)?;
249
250 self.buf.append_byte(gtk_kde.info.value())?;
252 self.buf.append_byte(0)?;
253 self.buf.append_bytes(>k_kde.gtk[..])?;
254 Ok(())
255 }
256
257 pub fn write_igtk(&mut self, igtk_kde: &Igtk) -> Result<(), BufferTooSmall> {
258 if !self.buf.can_append(HDR_LEN + igtk_kde.len()) {
259 return Err(BufferTooSmall);
260 }
261 let hdr = Header::new_dot11(IGTK_DATA_TYPE, igtk_kde.len());
263 self.write_kde_hdr(hdr)?;
264
265 self.buf.append_bytes(&igtk_kde.id.to_le_bytes())?;
267 self.buf.append_bytes(&igtk_kde.ipn[..])?;
268 self.buf.append_bytes(&igtk_kde.igtk[..])?;
269 Ok(())
270 }
271
272 pub fn finalize_for_encryption(mut self) -> Result<A, BufferTooSmall> {
273 let written = self.buf.bytes_appended();
278 let padding_len =
279 if written < 16 { 16 - written } else { ((written + 7) / 8) * 8 - written };
280 if !self.buf.can_append(padding_len) {
281 return Err(BufferTooSmall);
282 }
283
284 if padding_len != 0 {
285 self.buf.append_byte(TYPE)?;
286 self.buf.append_bytes_zeroed(padding_len - 1)?;
287 }
288 Ok(self.buf)
289 }
290
291 pub fn finalize_for_plaintext(self) -> Result<A, BufferTooSmall> {
292 Ok(self.buf)
293 }
294}
295
296#[cfg(test)]
297mod tests {
298 use super::*;
299 use crate::key_data::extract_elements;
300 use assert_matches::assert_matches;
301 use wlan_common::test_utils::FixedSizedTestBuffer;
302
303 fn write_and_extract_padding(gtk_len: usize) -> Vec<u8> {
304 let mut w = Writer::new();
305 w.write_gtk(&Gtk::new(2, GtkInfoTx::BothRxTx, &vec![2; gtk_len]))
306 .expect("failure writing GTK KDE");
307 w.finalize_for_encryption()
308 .expect("failure finializing key data")
309 .into_vec()
310 .split_off(HDR_LEN + GTK_FIXED_LEN + gtk_len)
311 }
312
313 #[test]
314 fn test_padding_min_length() {
315 let buf = write_and_extract_padding(0);
316 assert_eq!(buf, vec![TYPE, 0, 0, 0, 0, 0, 0, 0]);
317
318 let buf = write_and_extract_padding(5);
319 assert_eq!(buf, vec![TYPE, 0, 0]);
320
321 let buf = write_and_extract_padding(7);
322 assert_eq!(buf, vec![TYPE]);
323 }
324
325 #[test]
326 fn test_no_padding_expected_for_plaintext() {
327 let mut w = Writer::new();
328 w.write_gtk(&Gtk::new(2, GtkInfoTx::BothRxTx, &[2; 2])).expect("failure writing GTK KDE");
329 let buf = w.finalize_for_plaintext().expect("failure finalizing key data");
330 assert_eq!(
331 &buf[..],
332 &[
333 TYPE,
335 8,
336 0x00,
337 0x0F,
338 0xAC,
339 GTK_DATA_TYPE,
340 0b0000_0110,
341 0,
342 2,
343 2
344 ][..]
345 );
346 }
347
348 #[test]
349 fn test_no_padding() {
350 let buf = write_and_extract_padding(8);
351 assert!(buf.is_empty());
352 }
353
354 #[test]
355 fn test_padding_multiple_8() {
356 let buf = write_and_extract_padding(9);
357 assert_eq!(buf, vec![TYPE, 0, 0, 0, 0, 0, 0]);
358
359 let buf = write_and_extract_padding(16);
360 assert!(buf.is_empty());
361 }
362
363 #[test]
364 fn test_write_read_gtk_with_padding() {
365 let mut w = Writer::new();
367 w.write_gtk(&Gtk::new(2, GtkInfoTx::BothRxTx, &[24; 5])).expect("failure writing GTK KDE");
368 let buf = w.finalize_for_encryption().expect("failure finializing key data");
369 #[rustfmt::skip]
370 assert_eq!(&buf[..], &[
371 TYPE, 11, 0x00, 0x0F, 0xAC, GTK_DATA_TYPE, 0b0000_0110, 0, 24, 24, 24, 24, 24,
373 TYPE, 0, 0,
375 ][..]);
376
377 let result = extract_elements(&buf[..]);
379 assert!(result.is_ok(), "Error: {:?}", result);
380 let mut elements = result.unwrap();
381 assert_eq!(elements.len(), 2);
382
383 assert_matches!(elements.remove(0), Element::Gtk(hdr, kde) => {
384 assert_eq!(hdr, Header { type_: 0xDD, len: 11, oui: Oui::DOT11, data_type: 1 });
385 assert_eq!(kde, Gtk { info: GtkInfo(6), gtk: Box::new([24; 5]) });
386 });
387 assert_matches!(elements.remove(0), Element::Padding);
388 }
389
390 #[test]
391 fn test_write_read_igtk() {
392 let igtk = Igtk::new(10, &[11; 6], &[22; 2]);
393 let mut w = Writer::new();
395 w.write_igtk(&igtk).expect("failure writing IGTK KDE");
396 let buf = w.finalize_for_encryption().expect("failure finializing key data");
397 #[rustfmt::skip]
398 assert_eq!(&buf[..], &[
399 TYPE, 14, 0x00, 0x0F, 0xAC, IGTK_DATA_TYPE, 10, 0, 11,11,11,11,11,11,22,22
400 ][..]);
401
402 let result = extract_elements(&buf[..]);
404 assert!(result.is_ok(), "Error: {:?}", result);
405 let mut elements = result.unwrap();
406 assert_eq!(elements.len(), 1);
407
408 assert_matches!(elements.remove(0), Element::Igtk(hdr, kde) => {
409 assert_eq!(hdr, Header {
410 type_: 0xDD, len: 14, oui: Oui::DOT11, data_type: IGTK_DATA_TYPE
411 });
412 assert_eq!(kde, igtk);
413 });
414 }
415
416 #[test]
417 fn test_write_gtk_too_small_buffer() {
418 Writer::new_with(FixedSizedTestBuffer::new(10))
419 .write_gtk(&Gtk::new(2, GtkInfoTx::BothRxTx, &[24; 5]))
420 .expect_err("expected failure writing GTK KDE");
421 }
422
423 #[test]
424 fn test_write_gtk_sufficient_fixed_buffer() {
425 Writer::new_with(FixedSizedTestBuffer::new(13))
426 .write_gtk(&Gtk::new(2, GtkInfoTx::BothRxTx, &[24; 5]))
427 .expect("expected success writing GTK KDE");
428 }
429
430 #[test]
431 fn test_create_gtk_element() {
432 let gtk = Gtk::new(1, GtkInfoTx::_OnlyRx, &[24; 16][..]);
433 assert_eq!(gtk.info.key_id(), 1);
434 assert_eq!(gtk.info.tx(), 0);
435 assert_eq!(>k.gtk[..], &[24; 16][..]);
436 }
437
438 #[test]
439 fn test_gtk_len() {
440 let gtk_kde = Gtk { info: GtkInfo(0), gtk: Box::new([]) };
441 assert_eq!(gtk_kde.len(), 2);
442
443 let gtk_kde = Gtk { info: GtkInfo(0), gtk: Box::new([0; 16]) };
444 assert_eq!(gtk_kde.len(), 18);
445
446 let gtk_kde = Gtk { info: GtkInfo(0), gtk: Box::new([0; 8]) };
447 assert_eq!(gtk_kde.len(), 10);
448
449 let gtk_kde = Gtk { info: GtkInfo(0), gtk: Box::new([0; 4]) };
450 assert_eq!(gtk_kde.len(), 6);
451
452 let gtk_kde = Gtk { info: GtkInfo(0), gtk: Box::new([0; 32]) };
453 assert_eq!(gtk_kde.len(), 34);
454 }
455}