tungstenite/handshake/
client.rs

1//! Client handshake machine.
2
3use std::{
4    io::{Read, Write},
5    marker::PhantomData,
6};
7
8use http::{
9    header::HeaderName, HeaderMap, Request as HttpRequest, Response as HttpResponse, StatusCode,
10};
11use httparse::Status;
12use log::*;
13
14use super::{
15    derive_accept_key,
16    headers::{FromHttparse, MAX_HEADERS},
17    machine::{HandshakeMachine, StageResult, TryParse},
18    HandshakeRole, MidHandshake, ProcessingResult,
19};
20use crate::{
21    error::{Error, ProtocolError, Result, SubProtocolError, UrlError},
22    protocol::{Role, WebSocket, WebSocketConfig},
23};
24
25/// Client request type.
26pub type Request = HttpRequest<()>;
27
28/// Client response type.
29pub type Response = HttpResponse<Option<Vec<u8>>>;
30
31/// Client handshake role.
32#[derive(Debug)]
33pub struct ClientHandshake<S> {
34    verify_data: VerifyData,
35    config: Option<WebSocketConfig>,
36    _marker: PhantomData<S>,
37}
38
39impl<S: Read + Write> ClientHandshake<S> {
40    /// Initiate a client handshake.
41    pub fn start(
42        stream: S,
43        request: Request,
44        config: Option<WebSocketConfig>,
45    ) -> Result<MidHandshake<Self>> {
46        if request.method() != http::Method::GET {
47            return Err(Error::Protocol(ProtocolError::WrongHttpMethod));
48        }
49
50        if request.version() < http::Version::HTTP_11 {
51            return Err(Error::Protocol(ProtocolError::WrongHttpVersion));
52        }
53
54        // Check the URI scheme: only ws or wss are supported
55        let _ = crate::client::uri_mode(request.uri())?;
56
57        let subprotocols = extract_subprotocols_from_request(&request)?;
58
59        // Convert and verify the `http::Request` and turn it into the request as per RFC.
60        // Also extract the key from it (it must be present in a correct request).
61        let (request, key) = generate_request(request)?;
62
63        let machine = HandshakeMachine::start_write(stream, request);
64
65        let client = {
66            let accept_key = derive_accept_key(key.as_ref());
67            ClientHandshake {
68                verify_data: VerifyData { accept_key, subprotocols },
69                config,
70                _marker: PhantomData,
71            }
72        };
73
74        trace!("Client handshake initiated.");
75        Ok(MidHandshake { role: client, machine })
76    }
77}
78
79impl<S: Read + Write> HandshakeRole for ClientHandshake<S> {
80    type IncomingData = Response;
81    type InternalStream = S;
82    type FinalResult = (WebSocket<S>, Response);
83    fn stage_finished(
84        &mut self,
85        finish: StageResult<Self::IncomingData, Self::InternalStream>,
86    ) -> Result<ProcessingResult<Self::InternalStream, Self::FinalResult>> {
87        Ok(match finish {
88            StageResult::DoneWriting(stream) => {
89                ProcessingResult::Continue(HandshakeMachine::start_read(stream))
90            }
91            StageResult::DoneReading { stream, result, tail } => {
92                let result = match self.verify_data.verify_response(result) {
93                    Ok(r) => r,
94                    Err(Error::Http(mut e)) => {
95                        *e.body_mut() = Some(tail);
96                        return Err(Error::Http(e));
97                    }
98                    Err(e) => return Err(e),
99                };
100
101                debug!("Client handshake done.");
102                let websocket =
103                    WebSocket::from_partially_read(stream, tail, Role::Client, self.config);
104                ProcessingResult::Done((websocket, result))
105            }
106        })
107    }
108}
109
110/// Verifies and generates a client WebSocket request from the original request and extracts a WebSocket key from it.
111pub fn generate_request(mut request: Request) -> Result<(Vec<u8>, String)> {
112    let mut req = Vec::new();
113    write!(
114        req,
115        "GET {path} {version:?}\r\n",
116        path = request.uri().path_and_query().ok_or(Error::Url(UrlError::NoPathOrQuery))?.as_str(),
117        version = request.version()
118    )
119    .unwrap();
120
121    // Headers that must be present in a correct request.
122    const KEY_HEADERNAME: &str = "Sec-WebSocket-Key";
123    const WEBSOCKET_HEADERS: [&str; 5] =
124        ["Host", "Connection", "Upgrade", "Sec-WebSocket-Version", KEY_HEADERNAME];
125
126    // We must extract a WebSocket key from a properly formed request or fail if it's not present.
127    let key = request
128        .headers()
129        .get(KEY_HEADERNAME)
130        .ok_or_else(|| {
131            Error::Protocol(ProtocolError::InvalidHeader(
132                HeaderName::from_bytes(KEY_HEADERNAME.as_bytes()).unwrap(),
133            ))
134        })?
135        .to_str()?
136        .to_owned();
137
138    // We must check that all necessary headers for a valid request are present. Note that we have to
139    // deal with the fact that some apps seem to have a case-sensitive check for headers which is not
140    // correct and should not considered the correct behavior, but it seems like some apps ignore it.
141    // `http` by default writes all headers in lower-case which is fine (and does not violate the RFC)
142    // but some servers seem to be poorely written and ignore RFC.
143    //
144    // See similar problem in `hyper`: https://github.com/hyperium/hyper/issues/1492
145    let headers = request.headers_mut();
146    for &header in &WEBSOCKET_HEADERS {
147        let value = headers.remove(header).ok_or_else(|| {
148            Error::Protocol(ProtocolError::InvalidHeader(
149                HeaderName::from_bytes(header.as_bytes()).unwrap(),
150            ))
151        })?;
152        write!(
153            req,
154            "{header}: {value}\r\n",
155            header = header,
156            value = value.to_str().map_err(|err| {
157                Error::Utf8(format!("{err} for header name '{header}' with value: {value:?}"))
158            })?
159        )
160        .unwrap();
161    }
162
163    // Now we must ensure that the headers that we've written once are not anymore present in the map.
164    // If they do, then the request is invalid (some headers are duplicated there for some reason).
165    let insensitive: Vec<String> =
166        WEBSOCKET_HEADERS.iter().map(|h| h.to_ascii_lowercase()).collect();
167    for (k, v) in headers {
168        let mut name = k.as_str();
169
170        // We have already written the necessary headers once (above) and removed them from the map.
171        // If we encounter them again, then the request is considered invalid and error is returned.
172        // Note that we can't use `.contains()`, since `&str` does not coerce to `&String` in Rust.
173        if insensitive.iter().any(|x| x == name) {
174            return Err(Error::Protocol(ProtocolError::InvalidHeader(k.clone())));
175        }
176
177        // Relates to the issue of some servers treating headers in a case-sensitive way, please see:
178        // https://github.com/snapview/tungstenite-rs/pull/119 (original fix of the problem)
179        if name == "sec-websocket-protocol" {
180            name = "Sec-WebSocket-Protocol";
181        }
182
183        if name == "origin" {
184            name = "Origin";
185        }
186
187        writeln!(
188            req,
189            "{}: {}\r",
190            name,
191            v.to_str().map_err(|err| {
192                Error::Utf8(format!("{err} for header name '{name}' with value: {v:?}"))
193            })?
194        )
195        .unwrap();
196    }
197
198    writeln!(req, "\r").unwrap();
199    trace!("Request: {:?}", String::from_utf8_lossy(&req));
200    Ok((req, key))
201}
202
203fn extract_subprotocols_from_request(request: &Request) -> Result<Option<Vec<String>>> {
204    if let Some(subprotocols) = request.headers().get("Sec-WebSocket-Protocol") {
205        Ok(Some(subprotocols.to_str()?.split(',').map(|s| s.trim().to_string()).collect()))
206    } else {
207        Ok(None)
208    }
209}
210
211/// Information for handshake verification.
212#[derive(Debug)]
213struct VerifyData {
214    /// Accepted server key.
215    accept_key: String,
216
217    /// Accepted subprotocols
218    subprotocols: Option<Vec<String>>,
219}
220
221impl VerifyData {
222    pub fn verify_response(&self, response: Response) -> Result<Response> {
223        // 1. If the status code received from the server is not 101, the
224        // client handles the response per HTTP [RFC2616] procedures. (RFC 6455)
225        if response.status() != StatusCode::SWITCHING_PROTOCOLS {
226            return Err(Error::Http(response));
227        }
228
229        let headers = response.headers();
230
231        // 2. If the response lacks an |Upgrade| header field or the |Upgrade|
232        // header field contains a value that is not an ASCII case-
233        // insensitive match for the value "websocket", the client MUST
234        // _Fail the WebSocket Connection_. (RFC 6455)
235        if !headers
236            .get("Upgrade")
237            .and_then(|h| h.to_str().ok())
238            .map(|h| h.eq_ignore_ascii_case("websocket"))
239            .unwrap_or(false)
240        {
241            return Err(Error::Protocol(ProtocolError::MissingUpgradeWebSocketHeader));
242        }
243        // 3.  If the response lacks a |Connection| header field or the
244        // |Connection| header field doesn't contain a token that is an
245        // ASCII case-insensitive match for the value "Upgrade", the client
246        // MUST _Fail the WebSocket Connection_. (RFC 6455)
247        if !headers
248            .get("Connection")
249            .and_then(|h| h.to_str().ok())
250            .map(|h| h.eq_ignore_ascii_case("Upgrade"))
251            .unwrap_or(false)
252        {
253            return Err(Error::Protocol(ProtocolError::MissingConnectionUpgradeHeader));
254        }
255        // 4.  If the response lacks a |Sec-WebSocket-Accept| header field or
256        // the |Sec-WebSocket-Accept| contains a value other than the
257        // base64-encoded SHA-1 of ... the client MUST _Fail the WebSocket
258        // Connection_. (RFC 6455)
259        if !headers.get("Sec-WebSocket-Accept").map(|h| h == &self.accept_key).unwrap_or(false) {
260            return Err(Error::Protocol(ProtocolError::SecWebSocketAcceptKeyMismatch));
261        }
262        // 5.  If the response includes a |Sec-WebSocket-Extensions| header
263        // field and this header field indicates the use of an extension
264        // that was not present in the client's handshake (the server has
265        // indicated an extension not requested by the client), the client
266        // MUST _Fail the WebSocket Connection_. (RFC 6455)
267        // TODO
268
269        // 6.  If the response includes a |Sec-WebSocket-Protocol| header field
270        // and this header field indicates the use of a subprotocol that was
271        // not present in the client's handshake (the server has indicated a
272        // subprotocol not requested by the client), the client MUST _Fail
273        // the WebSocket Connection_. (RFC 6455)
274        if headers.get("Sec-WebSocket-Protocol").is_none() && self.subprotocols.is_some() {
275            return Err(Error::Protocol(ProtocolError::SecWebSocketSubProtocolError(
276                SubProtocolError::NoSubProtocol,
277            )));
278        }
279
280        if headers.get("Sec-WebSocket-Protocol").is_some() && self.subprotocols.is_none() {
281            return Err(Error::Protocol(ProtocolError::SecWebSocketSubProtocolError(
282                SubProtocolError::ServerSentSubProtocolNoneRequested,
283            )));
284        }
285
286        if let Some(returned_subprotocol) = headers.get("Sec-WebSocket-Protocol") {
287            if let Some(accepted_subprotocols) = &self.subprotocols {
288                if !accepted_subprotocols.contains(&returned_subprotocol.to_str()?.to_string()) {
289                    return Err(Error::Protocol(ProtocolError::SecWebSocketSubProtocolError(
290                        SubProtocolError::InvalidSubProtocol,
291                    )));
292                }
293            }
294        }
295
296        Ok(response)
297    }
298}
299
300impl TryParse for Response {
301    fn try_parse(buf: &[u8]) -> Result<Option<(usize, Self)>> {
302        let mut hbuffer = [httparse::EMPTY_HEADER; MAX_HEADERS];
303        let mut req = httparse::Response::new(&mut hbuffer);
304        Ok(match req.parse(buf)? {
305            Status::Partial => None,
306            Status::Complete(size) => Some((size, Response::from_httparse(req)?)),
307        })
308    }
309}
310
311impl<'h, 'b: 'h> FromHttparse<httparse::Response<'h, 'b>> for Response {
312    fn from_httparse(raw: httparse::Response<'h, 'b>) -> Result<Self> {
313        if raw.version.expect("Bug: no HTTP version") < /*1.*/1 {
314            return Err(Error::Protocol(ProtocolError::WrongHttpVersion));
315        }
316
317        let headers = HeaderMap::from_httparse(raw.headers)?;
318
319        let mut response = Response::new(None);
320        *response.status_mut() = StatusCode::from_u16(raw.code.expect("Bug: no HTTP status code"))?;
321        *response.headers_mut() = headers;
322        // TODO: httparse only supports HTTP 0.9/1.0/1.1 but not HTTP 2.0
323        // so the only valid value we could get in the response would be 1.1.
324        *response.version_mut() = http::Version::HTTP_11;
325
326        Ok(response)
327    }
328}
329
330/// Generate a random key for the `Sec-WebSocket-Key` header.
331pub fn generate_key() -> String {
332    // a base64-encoded (see Section 4 of [RFC4648]) value that,
333    // when decoded, is 16 bytes in length (RFC 6455)
334    let r: [u8; 16] = rand::random();
335    data_encoding::BASE64.encode(&r)
336}
337
338#[cfg(test)]
339mod tests {
340    use super::{super::machine::TryParse, generate_key, generate_request, Response};
341    use crate::client::IntoClientRequest;
342
343    #[test]
344    fn random_keys() {
345        let k1 = generate_key();
346        println!("Generated random key 1: {k1}");
347        let k2 = generate_key();
348        println!("Generated random key 2: {k2}");
349        assert_ne!(k1, k2);
350        assert_eq!(k1.len(), k2.len());
351        assert_eq!(k1.len(), 24);
352        assert_eq!(k2.len(), 24);
353        assert!(k1.ends_with("=="));
354        assert!(k2.ends_with("=="));
355        assert!(k1[..22].find('=').is_none());
356        assert!(k2[..22].find('=').is_none());
357    }
358
359    fn construct_expected(host: &str, key: &str) -> Vec<u8> {
360        format!(
361            "\
362            GET /getCaseCount HTTP/1.1\r\n\
363            Host: {host}\r\n\
364            Connection: Upgrade\r\n\
365            Upgrade: websocket\r\n\
366            Sec-WebSocket-Version: 13\r\n\
367            Sec-WebSocket-Key: {key}\r\n\
368            \r\n"
369        )
370        .into_bytes()
371    }
372
373    #[test]
374    fn request_formatting() {
375        let request = "ws://localhost/getCaseCount".into_client_request().unwrap();
376        let (request, key) = generate_request(request).unwrap();
377        let correct = construct_expected("localhost", &key);
378        assert_eq!(&request[..], &correct[..]);
379    }
380
381    #[test]
382    fn request_formatting_with_host() {
383        let request = "wss://localhost:9001/getCaseCount".into_client_request().unwrap();
384        let (request, key) = generate_request(request).unwrap();
385        let correct = construct_expected("localhost:9001", &key);
386        assert_eq!(&request[..], &correct[..]);
387    }
388
389    #[test]
390    fn request_formatting_with_at() {
391        let request = "wss://user:pass@localhost:9001/getCaseCount".into_client_request().unwrap();
392        let (request, key) = generate_request(request).unwrap();
393        let correct = construct_expected("localhost:9001", &key);
394        assert_eq!(&request[..], &correct[..]);
395    }
396
397    #[test]
398    fn response_parsing() {
399        const DATA: &[u8] = b"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n";
400        let (_, resp) = Response::try_parse(DATA).unwrap().unwrap();
401        assert_eq!(resp.status(), http::StatusCode::OK);
402        assert_eq!(resp.headers().get("Content-Type").unwrap(), &b"text/html"[..],);
403    }
404
405    #[test]
406    fn invalid_custom_request() {
407        let request = http::Request::builder().method("GET").body(()).unwrap();
408        assert!(generate_request(request).is_err());
409    }
410}
tungstenite/handshake/client.rs

tungstenite/handshake/
client.rs
