1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
// Copyright (c) 2019 Parity Technologies (UK) Ltd.
//
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

//! Websocket [handshake]s.
//!
//! [handshake]: https://tools.ietf.org/html/rfc6455#section-4

pub mod client;
#[cfg(feature = "http")]
pub mod http;
pub mod server;

use crate::extension::{Extension, Param};
use bytes::BytesMut;
use sha1::{Digest, Sha1};
use std::{fmt, io, str};

pub use client::{Client, ServerResponse};
pub use server::{ClientRequest, Server};

// Defined in RFC 6455 and used to generate the `Sec-WebSocket-Accept` header
// in the server handshake response.
const KEY: &[u8] = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

// How many HTTP headers do we support during parsing?
const MAX_NUM_HEADERS: usize = 32;

// Some HTTP headers we need to check during parsing.
const SEC_WEBSOCKET_EXTENSIONS: &str = "Sec-WebSocket-Extensions";
const SEC_WEBSOCKET_PROTOCOL: &str = "Sec-WebSocket-Protocol";

/// Check a set of headers contains a specific one.
fn expect_ascii_header(headers: &[httparse::Header], name: &str, ours: &str) -> Result<(), Error> {
	enum State {
		Init,  // Start state
		Name,  // Header name found
		Match, // Header value matches
	}

	headers
		.iter()
		.filter(|h| h.name.eq_ignore_ascii_case(name))
		.fold(Ok(State::Init), |result, header| {
			if let Ok(State::Match) = result {
				return result;
			}
			if str::from_utf8(header.value)?.split(',').any(|v| v.trim().eq_ignore_ascii_case(ours)) {
				return Ok(State::Match);
			}
			Ok(State::Name)
		})
		.and_then(|state| match state {
			State::Init => Err(Error::HeaderNotFound(name.into())),
			State::Name => Err(Error::UnexpectedHeader(name.into())),
			State::Match => Ok(()),
		})
}

/// Pick the first header with the given name and apply the given closure to it.
fn with_first_header<'a, F, R>(headers: &[httparse::Header<'a>], name: &str, f: F) -> Result<R, Error>
where
	F: Fn(&'a [u8]) -> Result<R, Error>,
{
	if let Some(h) = headers.iter().find(|h| h.name.eq_ignore_ascii_case(name)) {
		f(h.value)
	} else {
		Err(Error::HeaderNotFound(name.into()))
	}
}

// Configure all extensions with parsed parameters.
fn configure_extensions(extensions: &mut [Box<dyn Extension + Send>], line: &str) -> Result<(), Error> {
	for e in line.split(',') {
		let mut ext_parts = e.split(';');
		if let Some(name) = ext_parts.next() {
			let name = name.trim();
			if let Some(ext) = extensions.iter_mut().find(|x| x.name().eq_ignore_ascii_case(name)) {
				let mut params = Vec::new();
				for p in ext_parts {
					let mut key_value = p.split('=');
					if let Some(key) = key_value.next().map(str::trim) {
						let val = key_value.next().map(|v| v.trim().trim_matches('"'));
						let mut p = Param::new(key);
						p.set_value(val);
						params.push(p)
					}
				}
				ext.configure(&params).map_err(Error::Extension)?
			}
		}
	}
	Ok(())
}

// Write all extensions to the given buffer.
fn append_extensions<'a, I>(extensions: I, bytes: &mut BytesMut)
where
	I: IntoIterator<Item = &'a Box<dyn Extension + Send>>,
{
	let mut iter = extensions.into_iter().peekable();

	if iter.peek().is_some() {
		bytes.extend_from_slice(b"\r\nSec-WebSocket-Extensions: ")
	}

	append_extension_header_value(iter, bytes)
}

// Write the extension header value to the given buffer.
fn append_extension_header_value<'a, I>(mut extensions_iter: std::iter::Peekable<I>, bytes: &mut BytesMut)
where
	I: Iterator<Item = &'a Box<dyn Extension + Send>>,
{
	while let Some(e) = extensions_iter.next() {
		bytes.extend_from_slice(e.name().as_bytes());
		for p in e.params() {
			bytes.extend_from_slice(b"; ");
			bytes.extend_from_slice(p.name().as_bytes());
			if let Some(v) = p.value() {
				bytes.extend_from_slice(b"=");
				bytes.extend_from_slice(v.as_bytes())
			}
		}
		if extensions_iter.peek().is_some() {
			bytes.extend_from_slice(b", ")
		}
	}
}

// This function takes a 16 byte key (base64 encoded, and so 24 bytes of input) that is expected via
// the `Sec-WebSocket-Key` header during a websocket handshake, and writes the response that's expected
// to be handed back in the response header `Sec-WebSocket-Accept`.
//
// The response is a base64 encoding of a 160bit hash. base64 encoding uses 1 ascii character per 6 bits,
// and 160 / 6 = 26.66 characters. The output is padded with '=' to the nearest 4 characters, so we need 28
// bytes in total for all of the characters.
//
// See https://datatracker.ietf.org/doc/html/rfc6455#section-1.3 for more information on this.
fn generate_accept_key<'k>(key_base64: &WebSocketKey) -> [u8; 28] {
	let mut digest = Sha1::new();
	digest.update(key_base64);
	digest.update(KEY);
	let d = digest.finalize();

	let mut output_buf = [0; 28];
	let n = base64::encode_config_slice(&d, base64::STANDARD, &mut output_buf);
	debug_assert_eq!(n, 28, "encoding to base64 should be exactly 28 bytes");
	output_buf
}

/// Enumeration of possible handshake errors.
#[non_exhaustive]
#[derive(Debug)]
pub enum Error {
	/// An I/O error has been encountered.
	Io(io::Error),
	/// An HTTP version =/= 1.1 was encountered.
	UnsupportedHttpVersion,
	/// An incomplete HTTP request.
	IncompleteHttpRequest,
	/// The value of the `Sec-WebSocket-Key` header is of unexpected length.
	SecWebSocketKeyInvalidLength(usize),
	/// The handshake request was not a GET request.
	InvalidRequestMethod,
	/// An HTTP header has not been present.
	HeaderNotFound(String),
	/// An HTTP header value was not expected.
	UnexpectedHeader(String),
	/// The Sec-WebSocket-Accept header value did not match.
	InvalidSecWebSocketAccept,
	/// The server returned an extension we did not ask for.
	UnsolicitedExtension,
	/// The server returned a protocol we did not ask for.
	UnsolicitedProtocol,
	/// An extension produced an error while encoding or decoding.
	Extension(crate::BoxedError),
	/// The HTTP entity could not be parsed successfully.
	Http(crate::BoxedError),
	/// UTF-8 decoding failed.
	Utf8(str::Utf8Error),
}

impl fmt::Display for Error {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		match self {
			Error::Io(e) => write!(f, "i/o error: {}", e),
			Error::UnsupportedHttpVersion => f.write_str("http version was not 1.1"),
			Error::IncompleteHttpRequest => f.write_str("http request was incomplete"),
			Error::SecWebSocketKeyInvalidLength(len) => {
				write!(f, "Sec-WebSocket-Key header was {} bytes long, expected 24", len)
			}
			Error::InvalidRequestMethod => f.write_str("handshake was not a GET request"),
			Error::HeaderNotFound(name) => write!(f, "header {} not found", name),
			Error::UnexpectedHeader(name) => write!(f, "header {} had an unexpected value", name),
			Error::InvalidSecWebSocketAccept => f.write_str("websocket key mismatch"),
			Error::UnsolicitedExtension => f.write_str("unsolicited extension returned"),
			Error::UnsolicitedProtocol => f.write_str("unsolicited protocol returned"),
			Error::Extension(e) => write!(f, "extension error: {}", e),
			Error::Http(e) => write!(f, "http parser error: {}", e),
			Error::Utf8(e) => write!(f, "utf-8 decoding error: {}", e),
		}
	}
}

impl std::error::Error for Error {
	fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
		match self {
			Error::Io(e) => Some(e),
			Error::Extension(e) => Some(&**e),
			Error::Http(e) => Some(&**e),
			Error::Utf8(e) => Some(e),
			Error::UnsupportedHttpVersion
			| Error::IncompleteHttpRequest
			| Error::SecWebSocketKeyInvalidLength(_)
			| Error::InvalidRequestMethod
			| Error::HeaderNotFound(_)
			| Error::UnexpectedHeader(_)
			| Error::InvalidSecWebSocketAccept
			| Error::UnsolicitedExtension
			| Error::UnsolicitedProtocol => None,
		}
	}
}

impl From<io::Error> for Error {
	fn from(e: io::Error) -> Self {
		Error::Io(e)
	}
}

impl From<str::Utf8Error> for Error {
	fn from(e: str::Utf8Error) -> Self {
		Error::Utf8(e)
	}
}

/// Owned value of the `Sec-WebSocket-Key` header.
///
/// Per [RFC 6455](https://datatracker.ietf.org/doc/html/rfc6455#section-4.1):
///
/// ```text
/// (...) The value of this header field MUST be a
/// nonce consisting of a randomly selected 16-byte value that has
/// been base64-encoded (see Section 4 of [RFC4648]). (...)
/// ```
///
/// Base64 encoding of the nonce produces 24 ASCII bytes, padding included.
pub type WebSocketKey = [u8; 24];

#[cfg(test)]
mod tests {
	use super::expect_ascii_header;

	#[test]
	fn header_match() {
		let headers = &[
			httparse::Header { name: "foo", value: b"a,b,c,d" },
			httparse::Header { name: "foo", value: b"x" },
			httparse::Header { name: "foo", value: b"y, z, a" },
			httparse::Header { name: "bar", value: b"xxx" },
			httparse::Header { name: "bar", value: b"sdfsdf 423 42 424" },
			httparse::Header { name: "baz", value: b"123" },
		];

		assert!(expect_ascii_header(headers, "foo", "a").is_ok());
		assert!(expect_ascii_header(headers, "foo", "b").is_ok());
		assert!(expect_ascii_header(headers, "foo", "c").is_ok());
		assert!(expect_ascii_header(headers, "foo", "d").is_ok());
		assert!(expect_ascii_header(headers, "foo", "x").is_ok());
		assert!(expect_ascii_header(headers, "foo", "y").is_ok());
		assert!(expect_ascii_header(headers, "foo", "z").is_ok());
		assert!(expect_ascii_header(headers, "foo", "a").is_ok());
		assert!(expect_ascii_header(headers, "bar", "xxx").is_ok());
		assert!(expect_ascii_header(headers, "bar", "sdfsdf 423 42 424").is_ok());
		assert!(expect_ascii_header(headers, "baz", "123").is_ok());
		assert!(expect_ascii_header(headers, "baz", "???").is_err());
		assert!(expect_ascii_header(headers, "???", "x").is_err());
	}
}