1 use crate::result::{WebSocketError, WebSocketResult}; 2 use hyper; 3 use hyper::header::parsing::from_one_raw_str; 4 use hyper::header::{Header, HeaderFormat}; 5 use std::fmt::{self, Debug}; 6 use std::str::FromStr; 7 8 use websocket_base::header::WebSocketKey as WebSocketKeyLL; 9 10 /// Represents a Sec-WebSocket-Key header. 11 #[derive(PartialEq, Clone, Copy, Default)] 12 pub struct WebSocketKey(pub WebSocketKeyLL); 13 14 impl Debug for WebSocketKey { fmt(&self, f: &mut fmt::Formatter) -> fmt::Result15 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 16 self.0.fmt(f) 17 } 18 } 19 20 impl FromStr for WebSocketKey { 21 type Err = WebSocketError; 22 from_str(key: &str) -> WebSocketResult<WebSocketKey>23 fn from_str(key: &str) -> WebSocketResult<WebSocketKey> { 24 Ok(WebSocketKey(WebSocketKeyLL::from_str(key)?)) 25 } 26 } 27 28 impl WebSocketKey { 29 /// Generate a new, random WebSocketKey new() -> WebSocketKey30 pub fn new() -> WebSocketKey { 31 WebSocketKey(WebSocketKeyLL::new()) 32 } 33 /// Return the Base64 encoding of this WebSocketKey serialize(&self) -> String34 pub fn serialize(&self) -> String { 35 self.0.serialize() 36 } 37 38 /// Create WebSocketKey by explicitly specifying the key from_array(a: [u8; 16]) -> WebSocketKey39 pub fn from_array(a: [u8; 16]) -> WebSocketKey { 40 WebSocketKey(WebSocketKeyLL(a)) 41 } 42 } 43 44 impl Header for WebSocketKey { header_name() -> &'static str45 fn header_name() -> &'static str { 46 "Sec-WebSocket-Key" 47 } 48 parse_header(raw: &[Vec<u8>]) -> hyper::Result<WebSocketKey>49 fn parse_header(raw: &[Vec<u8>]) -> hyper::Result<WebSocketKey> { 50 from_one_raw_str(raw) 51 } 52 } 53 54 impl HeaderFormat for WebSocketKey { fmt_header(&self, fmt: &mut fmt::Formatter) -> fmt::Result55 fn fmt_header(&self, fmt: &mut fmt::Formatter) -> fmt::Result { 56 write!(fmt, "{}", self.serialize()) 57 } 58 } 59 60 #[cfg(all(feature = "nightly", test))] 61 mod tests { 62 use super::*; 63 use hyper::header::Header; 64 use test; 65 66 #[test] test_header_key()67 fn test_header_key() { 68 use crate::header::Headers; 69 70 let extensions = WebSocketKey::from_array([65; 16]); 71 let mut headers = Headers::new(); 72 headers.set(extensions); 73 74 assert_eq!( 75 &headers.to_string()[..], 76 "Sec-WebSocket-Key: QUFBQUFBQUFBQUFBQUFBQQ==\r\n" 77 ); 78 } 79 80 #[test] test_header_from_str()81 fn test_header_from_str() { 82 let key = WebSocketKey::from_str("YSByZWFsbCBnb29kIGtleQ=="); 83 assert!(key.is_ok()); // 16 bytes 84 85 let key = WebSocketKey::from_str("YSBzaG9ydCBrZXk="); 86 assert!(key.is_err()); // < 16 bytes 87 88 let key = WebSocketKey::from_str("YSB2ZXJ5IHZlcnkgbG9uZyBrZXk="); 89 assert!(key.is_err()); // > 16 bytes 90 } 91 92 #[bench] bench_header_key_new(b: &mut test::Bencher)93 fn bench_header_key_new(b: &mut test::Bencher) { 94 b.iter(|| { 95 let mut key = WebSocketKey::new(); 96 test::black_box(&mut key); 97 }); 98 } 99 100 #[bench] bench_header_key_parse(b: &mut test::Bencher)101 fn bench_header_key_parse(b: &mut test::Bencher) { 102 let value = vec![b"QUFBQUFBQUFBQUFBQUFBQQ==".to_vec()]; 103 b.iter(|| { 104 let mut key: WebSocketKey = Header::parse_header(&value[..]).unwrap(); 105 test::black_box(&mut key); 106 }); 107 } 108 109 #[bench] bench_header_key_format(b: &mut test::Bencher)110 fn bench_header_key_format(b: &mut test::Bencher) { 111 let value = vec![b"QUFBQUFBQUFBQUFBQUFBQQ==".to_vec()]; 112 let val: WebSocketKey = Header::parse_header(&value[..]).unwrap(); 113 b.iter(|| { 114 format!("{}", val.serialize()); 115 }); 116 } 117 } 118