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
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
//! Client handshake machine.

use std::{
    io::{Read, Write},
    marker::PhantomData,
};

use http::{
    header::HeaderName, HeaderMap, Request as HttpRequest, Response as HttpResponse, StatusCode,
};
use httparse::Status;
use log::*;

use super::{
    derive_accept_key,
    headers::{FromHttparse, MAX_HEADERS},
    machine::{HandshakeMachine, StageResult, TryParse},
    HandshakeRole, MidHandshake, ProcessingResult,
};
use crate::{
    error::{Error, ProtocolError, Result, SubProtocolError, UrlError},
    protocol::{Role, WebSocket, WebSocketConfig},
};

/// Client request type.
pub type Request = HttpRequest<()>;

/// Client response type.
pub type Response = HttpResponse<Option<Vec<u8>>>;

/// Client handshake role.
#[derive(Debug)]
pub struct ClientHandshake<S> {
    verify_data: VerifyData,
    config: Option<WebSocketConfig>,
    _marker: PhantomData<S>,
}

impl<S: Read + Write> ClientHandshake<S> {
    /// Initiate a client handshake.
    pub fn start(
        stream: S,
        request: Request,
        config: Option<WebSocketConfig>,
    ) -> Result<MidHandshake<Self>> {
        if request.method() != http::Method::GET {
            return Err(Error::Protocol(ProtocolError::WrongHttpMethod));
        }

        if request.version() < http::Version::HTTP_11 {
            return Err(Error::Protocol(ProtocolError::WrongHttpVersion));
        }

        // Check the URI scheme: only ws or wss are supported
        let _ = crate::client::uri_mode(request.uri())?;

        let subprotocols = extract_subprotocols_from_request(&request)?;

        // Convert and verify the `http::Request` and turn it into the request as per RFC.
        // Also extract the key from it (it must be present in a correct request).
        let (request, key) = generate_request(request)?;

        let machine = HandshakeMachine::start_write(stream, request);

        let client = {
            let accept_key = derive_accept_key(key.as_ref());
            ClientHandshake {
                verify_data: VerifyData { accept_key, subprotocols },
                config,
                _marker: PhantomData,
            }
        };

        trace!("Client handshake initiated.");
        Ok(MidHandshake { role: client, machine })
    }
}

impl<S: Read + Write> HandshakeRole for ClientHandshake<S> {
    type IncomingData = Response;
    type InternalStream = S;
    type FinalResult = (WebSocket<S>, Response);
    fn stage_finished(
        &mut self,
        finish: StageResult<Self::IncomingData, Self::InternalStream>,
    ) -> Result<ProcessingResult<Self::InternalStream, Self::FinalResult>> {
        Ok(match finish {
            StageResult::DoneWriting(stream) => {
                ProcessingResult::Continue(HandshakeMachine::start_read(stream))
            }
            StageResult::DoneReading { stream, result, tail } => {
                let result = match self.verify_data.verify_response(result) {
                    Ok(r) => r,
                    Err(Error::Http(mut e)) => {
                        *e.body_mut() = Some(tail);
                        return Err(Error::Http(e));
                    }
                    Err(e) => return Err(e),
                };

                debug!("Client handshake done.");
                let websocket =
                    WebSocket::from_partially_read(stream, tail, Role::Client, self.config);
                ProcessingResult::Done((websocket, result))
            }
        })
    }
}

/// Verifies and generates a client WebSocket request from the original request and extracts a WebSocket key from it.
pub fn generate_request(mut request: Request) -> Result<(Vec<u8>, String)> {
    let mut req = Vec::new();
    write!(
        req,
        "GET {path} {version:?}\r\n",
        path = request.uri().path_and_query().ok_or(Error::Url(UrlError::NoPathOrQuery))?.as_str(),
        version = request.version()
    )
    .unwrap();

    // Headers that must be present in a correct request.
    const KEY_HEADERNAME: &str = "Sec-WebSocket-Key";
    const WEBSOCKET_HEADERS: [&str; 5] =
        ["Host", "Connection", "Upgrade", "Sec-WebSocket-Version", KEY_HEADERNAME];

    // We must extract a WebSocket key from a properly formed request or fail if it's not present.
    let key = request
        .headers()
        .get(KEY_HEADERNAME)
        .ok_or_else(|| {
            Error::Protocol(ProtocolError::InvalidHeader(
                HeaderName::from_bytes(KEY_HEADERNAME.as_bytes()).unwrap(),
            ))
        })?
        .to_str()?
        .to_owned();

    // We must check that all necessary headers for a valid request are present. Note that we have to
    // deal with the fact that some apps seem to have a case-sensitive check for headers which is not
    // correct and should not considered the correct behavior, but it seems like some apps ignore it.
    // `http` by default writes all headers in lower-case which is fine (and does not violate the RFC)
    // but some servers seem to be poorely written and ignore RFC.
    //
    // See similar problem in `hyper`: https://github.com/hyperium/hyper/issues/1492
    let headers = request.headers_mut();
    for &header in &WEBSOCKET_HEADERS {
        let value = headers.remove(header).ok_or_else(|| {
            Error::Protocol(ProtocolError::InvalidHeader(
                HeaderName::from_bytes(header.as_bytes()).unwrap(),
            ))
        })?;
        write!(req, "{header}: {value}\r\n", header = header, value = value.to_str()?).unwrap();
    }

    // Now we must ensure that the headers that we've written once are not anymore present in the map.
    // If they do, then the request is invalid (some headers are duplicated there for some reason).
    let insensitive: Vec<String> =
        WEBSOCKET_HEADERS.iter().map(|h| h.to_ascii_lowercase()).collect();
    for (k, v) in headers {
        let mut name = k.as_str();

        // We have already written the necessary headers once (above) and removed them from the map.
        // If we encounter them again, then the request is considered invalid and error is returned.
        // Note that we can't use `.contains()`, since `&str` does not coerce to `&String` in Rust.
        if insensitive.iter().any(|x| x == name) {
            return Err(Error::Protocol(ProtocolError::InvalidHeader(k.clone())));
        }

        // Relates to the issue of some servers treating headers in a case-sensitive way, please see:
        // https://github.com/snapview/tungstenite-rs/pull/119 (original fix of the problem)
        if name == "sec-websocket-protocol" {
            name = "Sec-WebSocket-Protocol";
        }

        if name == "origin" {
            name = "Origin";
        }

        writeln!(req, "{}: {}\r", name, v.to_str()?).unwrap();
    }

    writeln!(req, "\r").unwrap();
    trace!("Request: {:?}", String::from_utf8_lossy(&req));
    Ok((req, key))
}

fn extract_subprotocols_from_request(request: &Request) -> Result<Option<Vec<String>>> {
    if let Some(subprotocols) = request.headers().get("Sec-WebSocket-Protocol") {
        Ok(Some(subprotocols.to_str()?.split(",").map(|s| s.to_string()).collect()))
    } else {
        Ok(None)
    }
}

/// Information for handshake verification.
#[derive(Debug)]
struct VerifyData {
    /// Accepted server key.
    accept_key: String,

    /// Accepted subprotocols
    subprotocols: Option<Vec<String>>,
}

impl VerifyData {
    pub fn verify_response(&self, response: Response) -> Result<Response> {
        // 1. If the status code received from the server is not 101, the
        // client handles the response per HTTP [RFC2616] procedures. (RFC 6455)
        if response.status() != StatusCode::SWITCHING_PROTOCOLS {
            return Err(Error::Http(response));
        }

        let headers = response.headers();

        // 2. If the response lacks an |Upgrade| header field or the |Upgrade|
        // header field contains a value that is not an ASCII case-
        // insensitive match for the value "websocket", the client MUST
        // _Fail the WebSocket Connection_. (RFC 6455)
        if !headers
            .get("Upgrade")
            .and_then(|h| h.to_str().ok())
            .map(|h| h.eq_ignore_ascii_case("websocket"))
            .unwrap_or(false)
        {
            return Err(Error::Protocol(ProtocolError::MissingUpgradeWebSocketHeader));
        }
        // 3.  If the response lacks a |Connection| header field or the
        // |Connection| header field doesn't contain a token that is an
        // ASCII case-insensitive match for the value "Upgrade", the client
        // MUST _Fail the WebSocket Connection_. (RFC 6455)
        if !headers
            .get("Connection")
            .and_then(|h| h.to_str().ok())
            .map(|h| h.eq_ignore_ascii_case("Upgrade"))
            .unwrap_or(false)
        {
            return Err(Error::Protocol(ProtocolError::MissingConnectionUpgradeHeader));
        }
        // 4.  If the response lacks a |Sec-WebSocket-Accept| header field or
        // the |Sec-WebSocket-Accept| contains a value other than the
        // base64-encoded SHA-1 of ... the client MUST _Fail the WebSocket
        // Connection_. (RFC 6455)
        if !headers.get("Sec-WebSocket-Accept").map(|h| h == &self.accept_key).unwrap_or(false) {
            return Err(Error::Protocol(ProtocolError::SecWebSocketAcceptKeyMismatch));
        }
        // 5.  If the response includes a |Sec-WebSocket-Extensions| header
        // field and this header field indicates the use of an extension
        // that was not present in the client's handshake (the server has
        // indicated an extension not requested by the client), the client
        // MUST _Fail the WebSocket Connection_. (RFC 6455)
        // TODO

        // 6.  If the response includes a |Sec-WebSocket-Protocol| header field
        // and this header field indicates the use of a subprotocol that was
        // not present in the client's handshake (the server has indicated a
        // subprotocol not requested by the client), the client MUST _Fail
        // the WebSocket Connection_. (RFC 6455)
        if headers.get("Sec-WebSocket-Protocol").is_none() && self.subprotocols.is_some() {
            return Err(Error::Protocol(ProtocolError::SecWebSocketSubProtocolError(
                SubProtocolError::NoSubProtocol,
            )));
        }

        if headers.get("Sec-WebSocket-Protocol").is_some() && self.subprotocols.is_none() {
            return Err(Error::Protocol(ProtocolError::SecWebSocketSubProtocolError(
                SubProtocolError::ServerSentSubProtocolNoneRequested,
            )));
        }

        if let Some(returned_subprotocol) = headers.get("Sec-WebSocket-Protocol") {
            if let Some(accepted_subprotocols) = &self.subprotocols {
                if !accepted_subprotocols.contains(&returned_subprotocol.to_str()?.to_string()) {
                    return Err(Error::Protocol(ProtocolError::SecWebSocketSubProtocolError(
                        SubProtocolError::InvalidSubProtocol,
                    )));
                }
            }
        }

        Ok(response)
    }
}

impl TryParse for Response {
    fn try_parse(buf: &[u8]) -> Result<Option<(usize, Self)>> {
        let mut hbuffer = [httparse::EMPTY_HEADER; MAX_HEADERS];
        let mut req = httparse::Response::new(&mut hbuffer);
        Ok(match req.parse(buf)? {
            Status::Partial => None,
            Status::Complete(size) => Some((size, Response::from_httparse(req)?)),
        })
    }
}

impl<'h, 'b: 'h> FromHttparse<httparse::Response<'h, 'b>> for Response {
    fn from_httparse(raw: httparse::Response<'h, 'b>) -> Result<Self> {
        if raw.version.expect("Bug: no HTTP version") < /*1.*/1 {
            return Err(Error::Protocol(ProtocolError::WrongHttpVersion));
        }

        let headers = HeaderMap::from_httparse(raw.headers)?;

        let mut response = Response::new(None);
        *response.status_mut() = StatusCode::from_u16(raw.code.expect("Bug: no HTTP status code"))?;
        *response.headers_mut() = headers;
        // TODO: httparse only supports HTTP 0.9/1.0/1.1 but not HTTP 2.0
        // so the only valid value we could get in the response would be 1.1.
        *response.version_mut() = http::Version::HTTP_11;

        Ok(response)
    }
}

/// Generate a random key for the `Sec-WebSocket-Key` header.
pub fn generate_key() -> String {
    // a base64-encoded (see Section 4 of [RFC4648]) value that,
    // when decoded, is 16 bytes in length (RFC 6455)
    let r: [u8; 16] = rand::random();
    data_encoding::BASE64.encode(&r)
}

#[cfg(test)]
mod tests {
    use super::{super::machine::TryParse, generate_key, generate_request, Response};
    use crate::client::IntoClientRequest;

    #[test]
    fn random_keys() {
        let k1 = generate_key();
        println!("Generated random key 1: {}", k1);
        let k2 = generate_key();
        println!("Generated random key 2: {}", k2);
        assert_ne!(k1, k2);
        assert_eq!(k1.len(), k2.len());
        assert_eq!(k1.len(), 24);
        assert_eq!(k2.len(), 24);
        assert!(k1.ends_with("=="));
        assert!(k2.ends_with("=="));
        assert!(k1[..22].find('=').is_none());
        assert!(k2[..22].find('=').is_none());
    }

    fn construct_expected(host: &str, key: &str) -> Vec<u8> {
        format!(
            "\
            GET /getCaseCount HTTP/1.1\r\n\
            Host: {host}\r\n\
            Connection: Upgrade\r\n\
            Upgrade: websocket\r\n\
            Sec-WebSocket-Version: 13\r\n\
            Sec-WebSocket-Key: {key}\r\n\
            \r\n",
            host = host,
            key = key
        )
        .into_bytes()
    }

    #[test]
    fn request_formatting() {
        let request = "ws://localhost/getCaseCount".into_client_request().unwrap();
        let (request, key) = generate_request(request).unwrap();
        let correct = construct_expected("localhost", &key);
        assert_eq!(&request[..], &correct[..]);
    }

    #[test]
    fn request_formatting_with_host() {
        let request = "wss://localhost:9001/getCaseCount".into_client_request().unwrap();
        let (request, key) = generate_request(request).unwrap();
        let correct = construct_expected("localhost:9001", &key);
        assert_eq!(&request[..], &correct[..]);
    }

    #[test]
    fn request_formatting_with_at() {
        let request = "wss://user:pass@localhost:9001/getCaseCount".into_client_request().unwrap();
        let (request, key) = generate_request(request).unwrap();
        let correct = construct_expected("localhost:9001", &key);
        assert_eq!(&request[..], &correct[..]);
    }

    #[test]
    fn response_parsing() {
        const DATA: &[u8] = b"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n";
        let (_, resp) = Response::try_parse(DATA).unwrap().unwrap();
        assert_eq!(resp.status(), http::StatusCode::OK);
        assert_eq!(resp.headers().get("Content-Type").unwrap(), &b"text/html"[..],);
    }

    #[test]
    fn invalid_custom_request() {
        let request = http::Request::builder().method("GET").body(()).unwrap();
        assert!(generate_request(request).is_err());
    }
}