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