mirror of
https://github.com/denoland/deno.git
synced 2024-12-04 17:18:23 -05:00
355 lines
11 KiB
Rust
355 lines
11 KiB
Rust
// Copyright 2018-2024 the Deno authors. All rights reserved. MIT license.
|
|
|
|
use std::marker::PhantomData;
|
|
|
|
use bytes::Bytes;
|
|
use bytes::BytesMut;
|
|
use httparse::Status;
|
|
use hyper::header::HeaderName;
|
|
use hyper::header::HeaderValue;
|
|
use hyper::Response;
|
|
use memmem::Searcher;
|
|
use memmem::TwoWaySearcher;
|
|
use once_cell::sync::OnceCell;
|
|
|
|
#[derive(Debug, thiserror::Error)]
|
|
pub enum WebSocketUpgradeError {
|
|
#[error("invalid headers")]
|
|
InvalidHeaders,
|
|
#[error("{0}")]
|
|
HttpParse(#[from] httparse::Error),
|
|
#[error("{0}")]
|
|
Http(#[from] http::Error),
|
|
#[error("{0}")]
|
|
Utf8(#[from] std::str::Utf8Error),
|
|
#[error("{0}")]
|
|
InvalidHeaderName(#[from] http::header::InvalidHeaderName),
|
|
#[error("{0}")]
|
|
InvalidHeaderValue(#[from] http::header::InvalidHeaderValue),
|
|
#[error("invalid HTTP status line")]
|
|
InvalidHttpStatusLine,
|
|
#[error("attempted to write to completed upgrade buffer")]
|
|
UpgradeBufferAlreadyCompleted,
|
|
}
|
|
|
|
/// Given a buffer that ends in `\n\n` or `\r\n\r\n`, returns a parsed [`Request<Body>`].
|
|
fn parse_response<T: Default>(
|
|
header_bytes: &[u8],
|
|
) -> Result<(usize, Response<T>), WebSocketUpgradeError> {
|
|
let mut headers = [httparse::EMPTY_HEADER; 16];
|
|
let status = httparse::parse_headers(header_bytes, &mut headers)?;
|
|
match status {
|
|
Status::Complete((index, parsed)) => {
|
|
let mut resp = Response::builder().status(101).body(T::default())?;
|
|
for header in parsed.iter() {
|
|
resp.headers_mut().append(
|
|
HeaderName::from_bytes(header.name.as_bytes())?,
|
|
HeaderValue::from_str(std::str::from_utf8(header.value)?)?,
|
|
);
|
|
}
|
|
Ok((index, resp))
|
|
}
|
|
_ => Err(WebSocketUpgradeError::InvalidHeaders),
|
|
}
|
|
}
|
|
|
|
/// Find a newline in a slice.
|
|
fn find_newline(slice: &[u8]) -> Option<usize> {
|
|
for (i, byte) in slice.iter().enumerate() {
|
|
if *byte == b'\n' {
|
|
return Some(i);
|
|
}
|
|
}
|
|
None
|
|
}
|
|
|
|
/// WebSocket upgrade state machine states.
|
|
#[derive(Default)]
|
|
enum WebSocketUpgradeState {
|
|
#[default]
|
|
Initial,
|
|
StatusLine,
|
|
Headers,
|
|
Complete,
|
|
}
|
|
|
|
static HEADER_SEARCHER: OnceCell<TwoWaySearcher> = OnceCell::new();
|
|
static HEADER_SEARCHER2: OnceCell<TwoWaySearcher> = OnceCell::new();
|
|
|
|
#[derive(Default)]
|
|
pub struct WebSocketUpgrade<T: Default> {
|
|
state: WebSocketUpgradeState,
|
|
buf: BytesMut,
|
|
_t: PhantomData<T>,
|
|
}
|
|
|
|
impl<T: Default> WebSocketUpgrade<T> {
|
|
/// Ensures that the status line starts with "HTTP/1.1 101 " which matches all of the node.js
|
|
/// WebSocket libraries that are known. We don't care about the trailing status text.
|
|
fn validate_status(
|
|
&self,
|
|
status: &[u8],
|
|
) -> Result<(), WebSocketUpgradeError> {
|
|
if status.starts_with(b"HTTP/1.1 101 ") {
|
|
Ok(())
|
|
} else {
|
|
Err(WebSocketUpgradeError::InvalidHttpStatusLine)
|
|
}
|
|
}
|
|
|
|
/// Writes bytes to our upgrade buffer, returning [`Ok(None)`] if we need to keep feeding it data,
|
|
/// [`Ok(Some(Response))`] if we got a valid upgrade header, or [`Err`] if something went badly.
|
|
pub fn write(
|
|
&mut self,
|
|
bytes: &[u8],
|
|
) -> Result<Option<(Response<T>, Bytes)>, WebSocketUpgradeError> {
|
|
use WebSocketUpgradeState::*;
|
|
|
|
match self.state {
|
|
Initial => {
|
|
if let Some(index) = find_newline(bytes) {
|
|
let (status, rest) = bytes.split_at(index + 1);
|
|
self.validate_status(status)?;
|
|
|
|
// Fast path for the most common node.js WebSocket libraries that use \r\n as the
|
|
// separator between header lines and send the whole response in one packet.
|
|
if rest.ends_with(b"\r\n\r\n") {
|
|
let (index, response) = parse_response(rest)?;
|
|
if index == rest.len() {
|
|
return Ok(Some((response, Bytes::default())));
|
|
} else {
|
|
let bytes = Bytes::copy_from_slice(&rest[index..]);
|
|
return Ok(Some((response, bytes)));
|
|
}
|
|
}
|
|
|
|
self.state = Headers;
|
|
self.write(rest)
|
|
} else {
|
|
self.state = StatusLine;
|
|
self.buf.extend_from_slice(bytes);
|
|
Ok(None)
|
|
}
|
|
}
|
|
StatusLine => {
|
|
if let Some(index) = find_newline(bytes) {
|
|
let (status, rest) = bytes.split_at(index + 1);
|
|
self.buf.extend_from_slice(status);
|
|
self.validate_status(&self.buf)?;
|
|
self.buf.clear();
|
|
// Recursively process this write
|
|
self.state = Headers;
|
|
self.write(rest)
|
|
} else {
|
|
self.buf.extend_from_slice(bytes);
|
|
Ok(None)
|
|
}
|
|
}
|
|
Headers => {
|
|
self.buf.extend_from_slice(bytes);
|
|
let header_searcher =
|
|
HEADER_SEARCHER.get_or_init(|| TwoWaySearcher::new(b"\r\n\r\n"));
|
|
let header_searcher2 =
|
|
HEADER_SEARCHER2.get_or_init(|| TwoWaySearcher::new(b"\n\n"));
|
|
if header_searcher.search_in(&self.buf).is_some()
|
|
|| header_searcher2.search_in(&self.buf).is_some()
|
|
{
|
|
let (index, response) = parse_response(&self.buf)?;
|
|
let mut buf = std::mem::take(&mut self.buf);
|
|
self.state = Complete;
|
|
Ok(Some((response, buf.split_off(index).freeze())))
|
|
} else {
|
|
Ok(None)
|
|
}
|
|
}
|
|
Complete => Err(WebSocketUpgradeError::UpgradeBufferAlreadyCompleted),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use hyper_v014::Body;
|
|
|
|
type ExpectedResponseAndHead = Option<(Response<Body>, &'static [u8])>;
|
|
|
|
fn assert_response(
|
|
result: Result<Option<(Response<Body>, Bytes)>, WebSocketUpgradeError>,
|
|
expected: Result<ExpectedResponseAndHead, WebSocketUpgradeError>,
|
|
chunk_info: Option<(usize, usize)>,
|
|
) {
|
|
let formatted = format!("{result:?}");
|
|
match expected {
|
|
Ok(Some((resp1, remainder1))) => match result {
|
|
Ok(Some((resp2, remainder2))) => {
|
|
assert_eq!(format!("{resp1:?}"), format!("{resp2:?}"));
|
|
if let Some((byte_len, chunk_size)) = chunk_info {
|
|
// We need to compute how many bytes should be in the trailing data
|
|
|
|
// We know how many bytes of header data we had
|
|
let last_packet_header_size =
|
|
(byte_len - remainder1.len() + chunk_size - 1) % chunk_size + 1;
|
|
|
|
// Which means we can compute how much was in the remainder
|
|
let remaining =
|
|
(chunk_size - last_packet_header_size).min(remainder1.len());
|
|
|
|
assert_eq!(remainder1[..remaining], remainder2);
|
|
} else {
|
|
assert_eq!(remainder1, remainder2);
|
|
}
|
|
}
|
|
_ => panic!("Expected Ok(Some(...)), was {formatted}"),
|
|
},
|
|
Ok(None) => assert!(
|
|
result.ok().unwrap().is_none(),
|
|
"Expected Ok(None), was {formatted}",
|
|
),
|
|
Err(e) => assert_eq!(
|
|
format!("{e:?}"),
|
|
format!("{:?}", result.unwrap_err()),
|
|
"Expected error, was {formatted}",
|
|
),
|
|
}
|
|
}
|
|
|
|
fn validate_upgrade_all_at_once(
|
|
s: &str,
|
|
expected: Result<ExpectedResponseAndHead, WebSocketUpgradeError>,
|
|
) {
|
|
let mut upgrade = WebSocketUpgrade::default();
|
|
let res = upgrade.write(s.as_bytes());
|
|
|
|
assert_response(res, expected, None);
|
|
}
|
|
|
|
fn validate_upgrade_chunks(
|
|
s: &str,
|
|
size: usize,
|
|
expected: Result<ExpectedResponseAndHead, WebSocketUpgradeError>,
|
|
) {
|
|
let chunk_info = Some((s.as_bytes().len(), size));
|
|
let mut upgrade = WebSocketUpgrade::default();
|
|
let mut result = Ok(None);
|
|
for chunk in s.as_bytes().chunks(size) {
|
|
result = upgrade.write(chunk);
|
|
if let Ok(Some(..)) = &result {
|
|
assert_response(result, expected, chunk_info);
|
|
return;
|
|
}
|
|
}
|
|
assert_response(result, expected, chunk_info);
|
|
}
|
|
|
|
fn validate_upgrade(
|
|
s: &str,
|
|
expected: fn() -> Result<ExpectedResponseAndHead, WebSocketUpgradeError>,
|
|
) {
|
|
validate_upgrade_all_at_once(s, expected());
|
|
validate_upgrade_chunks(s, 1, expected());
|
|
validate_upgrade_chunks(s, 2, expected());
|
|
validate_upgrade_chunks(s, 10, expected());
|
|
|
|
// Replace \n with \r\n, but only in headers
|
|
let (headers, trailing) = s.split_once("\n\n").unwrap();
|
|
let s = headers.replace('\n', "\r\n") + "\r\n\r\n" + trailing;
|
|
let s = s.as_ref();
|
|
|
|
validate_upgrade_all_at_once(s, expected());
|
|
validate_upgrade_chunks(s, 1, expected());
|
|
validate_upgrade_chunks(s, 2, expected());
|
|
validate_upgrade_chunks(s, 10, expected());
|
|
}
|
|
|
|
#[test]
|
|
fn upgrade1() {
|
|
validate_upgrade(
|
|
"HTTP/1.1 101 Switching Protocols\nConnection: Upgrade\n\n",
|
|
|| {
|
|
let mut expected =
|
|
Response::builder().status(101).body(Body::empty()).unwrap();
|
|
expected.headers_mut().append(
|
|
HeaderName::from_static("connection"),
|
|
HeaderValue::from_static("Upgrade"),
|
|
);
|
|
Ok(Some((expected, b"")))
|
|
},
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn upgrade_trailing() {
|
|
validate_upgrade(
|
|
"HTTP/1.1 101 Switching Protocols\nConnection: Upgrade\n\ntrailing data",
|
|
|| {
|
|
let mut expected =
|
|
Response::builder().status(101).body(Body::empty()).unwrap();
|
|
expected.headers_mut().append(
|
|
HeaderName::from_static("connection"),
|
|
HeaderValue::from_static("Upgrade"),
|
|
);
|
|
Ok(Some((expected, b"trailing data")))
|
|
},
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn upgrade_trailing_with_newlines() {
|
|
validate_upgrade(
|
|
"HTTP/1.1 101 Switching Protocols\nConnection: Upgrade\n\ntrailing data\r\n\r\n",
|
|
|| {
|
|
let mut expected =
|
|
Response::builder().status(101).body(Body::empty()).unwrap();
|
|
expected.headers_mut().append(
|
|
HeaderName::from_static("connection"),
|
|
HeaderValue::from_static("Upgrade"),
|
|
);
|
|
Ok(Some((expected, b"trailing data\r\n\r\n")))
|
|
},
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn upgrade2() {
|
|
validate_upgrade(
|
|
"HTTP/1.1 101 Switching Protocols\nConnection: Upgrade\nOther: 123\n\n",
|
|
|| {
|
|
let mut expected =
|
|
Response::builder().status(101).body(Body::empty()).unwrap();
|
|
expected.headers_mut().append(
|
|
HeaderName::from_static("connection"),
|
|
HeaderValue::from_static("Upgrade"),
|
|
);
|
|
expected.headers_mut().append(
|
|
HeaderName::from_static("other"),
|
|
HeaderValue::from_static("123"),
|
|
);
|
|
Ok(Some((expected, b"")))
|
|
},
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn upgrade_invalid_status() {
|
|
validate_upgrade("HTTP/1.1 200 OK\nConnection: Upgrade\n\n", || {
|
|
Err(WebSocketUpgradeError::InvalidHttpStatusLine)
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn upgrade_too_many_headers() {
|
|
let headers = (0..20)
|
|
.map(|i| format!("h{i}: {i}"))
|
|
.collect::<Vec<_>>()
|
|
.join("\n");
|
|
validate_upgrade(
|
|
&format!("HTTP/1.1 101 Switching Protocols\n{headers}\n\n"),
|
|
|| {
|
|
Err(WebSocketUpgradeError::HttpParse(
|
|
httparse::Error::TooManyHeaders,
|
|
))
|
|
},
|
|
);
|
|
}
|
|
}
|