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

tungstenite/handshake/
client.rs
