blitz_ws/handshake/
client.rs

1//! Client handshake machine
2
3use std::{
4    io::{Read, Write},
5    marker::PhantomData,
6};
7
8use base64::Engine;
9use http::{
10    HeaderMap, HeaderName, Method, Request as HttpRequest, Response as HttpResponse, StatusCode,
11    Version,
12};
13use httparse::{Status, EMPTY_HEADER};
14
15use crate::{
16    error::{Error, ProtocolError, Result, SubProtocolError, UrlError},
17    handshake::{
18        core::{derive_accept_key, HandshakeRole, MidHandshake, ProcessingResult},
19        headers::{FromHttparse, MAX_HEADERS},
20        machine::{HandshakeMachine, StageResult, TryParse},
21    },
22    protocol::{
23        config::WebSocketConfig,
24        websocket::{OperationMode, WebSocket},
25    },
26};
27
28/// Client Request type
29pub type Request = HttpRequest<()>;
30/// Client Response Type
31pub type Response = HttpResponse<Option<Vec<u8>>>;
32
33/// Client handshake
34#[derive(Debug)]
35pub struct ClientHandshake<S> {
36    verify_data: VerifyData,
37    config: Option<WebSocketConfig>,
38    _marker: PhantomData<S>,
39}
40
41impl<S: Read + Write> ClientHandshake<S> {
42    /// Initiate a client handshake
43    pub fn start(
44        stream: S,
45        req: Request,
46        config: Option<WebSocketConfig>,
47    ) -> Result<MidHandshake<Self>> {
48        if req.method() != Method::GET {
49            return Err(Error::Protocol(ProtocolError::InvalidHttpMethod));
50        }
51        if req.version() < Version::HTTP_11 {
52            return Err(Error::Protocol(ProtocolError::InvalidHttpVersion));
53        }
54
55        let subprotocols = extract_subprotocols(&req)?;
56
57        let (request, key) = generate_request(req)?;
58
59        let machine = HandshakeMachine::start_write(stream, request);
60
61        let client = {
62            let accept_key = derive_accept_key(key.as_ref());
63            ClientHandshake {
64                verify_data: VerifyData { accept_key, subprotocols },
65                config,
66                _marker: PhantomData,
67            }
68        };
69
70        Ok(MidHandshake { role: client, machine })
71    }
72}
73
74impl<S: Read + Write> HandshakeRole for ClientHandshake<S> {
75    type IncomingData = Response;
76    type InternalStream = S;
77    type FinalResult = (WebSocket<S>, Response);
78
79    fn stage_finished(
80        &mut self,
81        finish: StageResult<Self::IncomingData, Self::InternalStream>,
82    ) -> Result<ProcessingResult<Self::InternalStream, Self::FinalResult>> {
83        Ok(match finish {
84            StageResult::DoneWriting(stream) => {
85                ProcessingResult::Continue(HandshakeMachine::start_read(stream))
86            }
87            StageResult::DoneReading { result, stream, tail } => {
88                let res = match self.verify_data.verify_response(result) {
89                    Ok(r) => r,
90                    Err(Error::Http(mut e)) => {
91                        *e.body_mut() = Some(tail);
92                        return Err(Error::Http(e));
93                    }
94                    Err(e) => return Err(e),
95                };
96
97                let websocket = WebSocket::from_partially_read(
98                    stream,
99                    tail,
100                    OperationMode::Client,
101                    self.config,
102                );
103                ProcessingResult::Done((websocket, res))
104            }
105        })
106    }
107}
108
109/// Verifies and generates a client WebSocket request from a raw request and extracts a WebSocket key from it
110pub fn generate_request(mut request: Request) -> Result<(Vec<u8>, String)> {
111    let mut req = Vec::new();
112    write!(
113        req,
114        "GET {path} {version:?}\r\n",
115        path = request.uri().path_and_query().ok_or(Error::Url(UrlError::NoPathOrQuery))?.as_str(),
116        version = request.version()
117    )
118    .unwrap();
119
120    const KEY_HEADERNAME: &str = "Sec-WebSocket-Key";
121    const WEBSOCKET_HEADERS: [&str; 5] =
122        ["Host", "Connection", "Upgrade", "Sec-WebSocket-Version", KEY_HEADERNAME];
123
124    let key = request
125        .headers()
126        .get(KEY_HEADERNAME)
127        .ok_or_else(|| {
128            Error::Protocol(ProtocolError::InvalidHeader(
129                HeaderName::from_bytes(KEY_HEADERNAME.as_bytes()).unwrap(),
130            ))
131        })?
132        .to_str()?
133        .to_owned();
134
135    let headers = request.headers_mut();
136    for &header in &WEBSOCKET_HEADERS {
137        let val = headers.remove(header).ok_or_else(|| {
138            Error::Protocol(ProtocolError::InvalidHeader(
139                HeaderName::from_bytes(header.as_bytes()).unwrap(),
140            ))
141        })?;
142
143        write!(
144            req,
145            "{header}: {value}\r\n",
146            header = header,
147            value = val.to_str().map_err(|e| {
148                Error::Utf8(format!("{e} for header name '{header}' with value: {val:?}"))
149            })?
150        )
151        .unwrap();
152    }
153
154    let insensitive: Vec<String> =
155        WEBSOCKET_HEADERS.iter().map(|h| h.to_ascii_lowercase()).collect();
156    for (k, v) in headers {
157        let mut name = k.as_str();
158
159        if insensitive.iter().any(|h| h == name) {
160            return Err(Error::Protocol(ProtocolError::InvalidHeader(k.clone())));
161        }
162
163        if name == "sec-websocket-protocol" {
164            name = "Sec-WebSocket-Protocol";
165        }
166        if name == "origin" {
167            name = "Origin";
168        }
169
170        writeln!(
171            req,
172            "{}: {}\r",
173            name,
174            v.to_str().map_err(|e| Error::Utf8(format!(
175                "{e} for header name '{name}' with value: {v:?}"
176            )))?
177        )
178        .unwrap();
179    }
180
181    writeln!(req, "\r").unwrap();
182    Ok((req, key))
183}
184
185fn extract_subprotocols(req: &Request) -> Result<Option<Vec<String>>> {
186    if let Some(subprotocols) = req.headers().get("Sec-WebSocket-Protocol") {
187        Ok(Some(subprotocols.to_str()?.split(',').map(|s| s.trim().to_string()).collect()))
188    } else {
189        Ok(None)
190    }
191}
192
193#[derive(Debug)]
194struct VerifyData {
195    accept_key: String,
196    subprotocols: Option<Vec<String>>,
197}
198
199impl VerifyData {
200    pub fn verify_response(&self, res: Response) -> Result<Response> {
201        if res.status() != StatusCode::SWITCHING_PROTOCOLS {
202            return Err(Error::Http(res));
203        }
204
205        let headers = res.headers();
206
207        if !headers
208            .get("Connection")
209            .and_then(|h| h.to_str().ok())
210            .map(|v| v.split([',', ' ']).any(|s| s.eq_ignore_ascii_case("Upgrade")))
211            .unwrap_or(false)
212        {
213            return Err(Error::Protocol(ProtocolError::MissingConnectionUpgradeHeader));
214        }
215
216        if !headers
217            .get("Upgrade")
218            .and_then(|h| h.to_str().ok())
219            .map(|v| v.eq_ignore_ascii_case("websocket"))
220            .unwrap_or(false)
221        {
222            return Err(Error::Protocol(ProtocolError::MissingUpgradeHeader));
223        }
224
225        if !headers.get("Sec-WebSocket-Accept").map(|h| h == &self.accept_key).unwrap_or(false) {
226            return Err(Error::Protocol(ProtocolError::AcceptKeyMismatch));
227        }
228
229        if headers.get("Sec-WebSocket-Protocol").is_none() && self.subprotocols.is_some() {
230            return Err(Error::Protocol(ProtocolError::SecWebSocketSubProtocolError(
231                SubProtocolError::NoSubProtocol,
232            )));
233        }
234        if headers.get("Sec-Websocket-Protocol").is_some() && self.subprotocols.is_none() {
235            return Err(Error::Protocol(ProtocolError::SecWebSocketSubProtocolError(
236                SubProtocolError::ServerSentSubProtocolNoneRequested,
237            )));
238        }
239        if let Some(returned_subprotocol) = headers.get("Sec-WebSocket-Protocol") {
240            if let Some(accepted_subprotocols) = &self.subprotocols {
241                if !accepted_subprotocols.contains(&returned_subprotocol.to_str()?.to_string()) {
242                    return Err(Error::Protocol(ProtocolError::SecWebSocketSubProtocolError(
243                        SubProtocolError::InvalidSubProtocol,
244                    )));
245                }
246            }
247        }
248
249        Ok(res)
250    }
251}
252
253impl TryParse for Response {
254    fn try_parse(data: &[u8]) -> crate::error::Result<Option<(usize, Self)>> {
255        let mut hbuffer = [EMPTY_HEADER; MAX_HEADERS];
256        let mut req = httparse::Response::new(&mut hbuffer);
257
258        Ok(match req.parse(data)? {
259            Status::Partial => None,
260            Status::Complete(n) => Some((n, Response::from_httparse(req)?)),
261        })
262    }
263}
264
265impl<'b: 'h, 'h> FromHttparse<httparse::Response<'h, 'b>> for Response {
266    fn from_httparse(raw: httparse::Response<'h, 'b>) -> crate::error::Result<Self> {
267        if raw.version != Some(1) {
268            return Err(Error::Protocol(ProtocolError::InvalidHttpVersion));
269        }
270
271        let headers = HeaderMap::from_httparse(raw.headers)?;
272
273        let mut res = Response::new(None);
274        *res.status_mut() = StatusCode::from_u16(raw.code.expect("Bug: no HTTP status code"))?;
275        *res.headers_mut() = headers;
276        *res.version_mut() = Version::HTTP_11;
277
278        Ok(res)
279    }
280}
281
282/// Generates a random accept key for the `Sec-WebSocket-Key` header
283pub fn generate_key() -> String {
284    let r: [u8; 16] = rand::random();
285    base64::engine::general_purpose::STANDARD.encode(r)
286}
blitz_ws/handshake/client.rs

blitz_ws/handshake/
client.rs
