tungstenite/handshake/
server.rs

1//! Server handshake machine.
2
3use std::{
4    io::{self, Read, Write},
5    marker::PhantomData,
6    result::Result as StdResult,
7};
8
9use http::{
10    header::HeaderValue, response::Builder, HeaderMap, Request as HttpRequest,
11    Response as HttpResponse, StatusCode,
12};
13use httparse::Status;
14use log::*;
15
16use super::{
17    derive_accept_key,
18    headers::{FromHttparse, MAX_HEADERS},
19    machine::{HandshakeMachine, StageResult, TryParse},
20    HandshakeRole, MidHandshake, ProcessingResult,
21};
22use crate::{
23    error::{Error, ProtocolError, Result},
24    handshake::version_as_str,
25    protocol::{Role, WebSocket, WebSocketConfig},
26};
27
28/// Server request type.
29pub type Request = HttpRequest<()>;
30
31/// Server response type.
32pub type Response = HttpResponse<()>;
33
34/// Server error response type.
35pub type ErrorResponse = HttpResponse<Option<String>>;
36
37fn create_parts<T>(request: &HttpRequest<T>) -> Result<Builder> {
38    if request.method() != http::Method::GET {
39        return Err(Error::Protocol(ProtocolError::WrongHttpMethod));
40    }
41
42    if request.version() < http::Version::HTTP_11 {
43        return Err(Error::Protocol(ProtocolError::WrongHttpVersion));
44    }
45
46    if !request
47        .headers()
48        .get("Connection")
49        .and_then(|h| h.to_str().ok())
50        .map(|h| h.split([' ', ',']).any(|p| p.eq_ignore_ascii_case("Upgrade")))
51        .unwrap_or(false)
52    {
53        return Err(Error::Protocol(ProtocolError::MissingConnectionUpgradeHeader));
54    }
55
56    if !request
57        .headers()
58        .get("Upgrade")
59        .and_then(|h| h.to_str().ok())
60        .map(|h| h.eq_ignore_ascii_case("websocket"))
61        .unwrap_or(false)
62    {
63        return Err(Error::Protocol(ProtocolError::MissingUpgradeWebSocketHeader));
64    }
65
66    if !request.headers().get("Sec-WebSocket-Version").map(|h| h == "13").unwrap_or(false) {
67        return Err(Error::Protocol(ProtocolError::MissingSecWebSocketVersionHeader));
68    }
69
70    let key = request
71        .headers()
72        .get("Sec-WebSocket-Key")
73        .ok_or(Error::Protocol(ProtocolError::MissingSecWebSocketKey))?;
74
75    if !is_valid_sec_websocket_key(key) {
76        return Err(Error::Protocol(ProtocolError::InvalidSecWebSocketKey));
77    }
78
79    let builder = Response::builder()
80        .status(StatusCode::SWITCHING_PROTOCOLS)
81        .version(request.version())
82        .header("Connection", "Upgrade")
83        .header("Upgrade", "websocket")
84        .header("Sec-WebSocket-Accept", derive_accept_key(key.as_bytes()));
85
86    Ok(builder)
87}
88
89fn is_valid_sec_websocket_key(key: &HeaderValue) -> bool {
90    if key.len() != 24 {
91        return false;
92    }
93
94    let Ok(decoded) = data_encoding::BASE64.decode(key.as_bytes()) else {
95        return false;
96    };
97
98    decoded.len() == 16
99}
100
101/// Create a response for the request.
102pub fn create_response(request: &Request) -> Result<Response> {
103    Ok(create_parts(request)?.body(())?)
104}
105
106/// Create a response for the request with a custom body.
107pub fn create_response_with_body<T1, T2>(
108    request: &HttpRequest<T1>,
109    generate_body: impl FnOnce() -> T2,
110) -> Result<HttpResponse<T2>> {
111    Ok(create_parts(request)?.body(generate_body())?)
112}
113
114/// Write `response` to the stream `w`.
115pub fn write_response<T>(mut w: impl io::Write, response: &HttpResponse<T>) -> Result<()> {
116    writeln!(
117        w,
118        "{version} {status}\r",
119        version = version_as_str(response.version())?,
120        status = response.status()
121    )?;
122
123    for (k, v) in response.headers() {
124        writeln!(w, "{}: {}\r", k, v.to_str()?)?;
125    }
126
127    writeln!(w, "\r")?;
128
129    Ok(())
130}
131
132impl TryParse for Request {
133    fn try_parse(buf: &[u8]) -> Result<Option<(usize, Self)>> {
134        let mut hbuffer = [httparse::EMPTY_HEADER; MAX_HEADERS];
135        let mut req = httparse::Request::new(&mut hbuffer);
136        Ok(match req.parse(buf)? {
137            Status::Partial => None,
138            Status::Complete(size) => Some((size, Request::from_httparse(req)?)),
139        })
140    }
141}
142
143impl<'h, 'b: 'h> FromHttparse<httparse::Request<'h, 'b>> for Request {
144    fn from_httparse(raw: httparse::Request<'h, 'b>) -> Result<Self> {
145        if raw.method.expect("Bug: no method in header") != "GET" {
146            return Err(Error::Protocol(ProtocolError::WrongHttpMethod));
147        }
148
149        if raw.version.expect("Bug: no HTTP version") < /*1.*/1 {
150            return Err(Error::Protocol(ProtocolError::WrongHttpVersion));
151        }
152
153        let headers = HeaderMap::from_httparse(raw.headers)?;
154
155        let mut request = Request::new(());
156        *request.method_mut() = http::Method::GET;
157        *request.headers_mut() = headers;
158        *request.uri_mut() = raw.path.expect("Bug: no path in header").parse()?;
159        // TODO: httparse only supports HTTP 0.9/1.0/1.1 but not HTTP 2.0
160        // so the only valid value we could get in the response would be 1.1.
161        *request.version_mut() = http::Version::HTTP_11;
162
163        Ok(request)
164    }
165}
166
167/// The callback trait.
168///
169/// The callback is called when the server receives an incoming WebSocket
170/// handshake request from the client. Specifying a callback allows you to analyze incoming headers
171/// and add additional headers to the response that server sends to the client and/or reject the
172/// connection based on the incoming headers.
173pub trait Callback: Sized {
174    /// Called whenever the server read the request from the client and is ready to reply to it.
175    /// May return additional reply headers.
176    /// Returning an error resulting in rejecting the incoming connection.
177    fn on_request(
178        self,
179        request: &Request,
180        response: Response,
181    ) -> StdResult<Response, ErrorResponse>;
182}
183
184impl<F> Callback for F
185where
186    F: FnOnce(&Request, Response) -> StdResult<Response, ErrorResponse>,
187{
188    fn on_request(
189        self,
190        request: &Request,
191        response: Response,
192    ) -> StdResult<Response, ErrorResponse> {
193        self(request, response)
194    }
195}
196
197/// Stub for callback that does nothing.
198#[derive(Clone, Copy, Debug)]
199pub struct NoCallback;
200
201impl Callback for NoCallback {
202    fn on_request(
203        self,
204        _request: &Request,
205        response: Response,
206    ) -> StdResult<Response, ErrorResponse> {
207        Ok(response)
208    }
209}
210
211/// Server handshake role.
212#[allow(missing_copy_implementations)]
213#[derive(Debug)]
214pub struct ServerHandshake<S, C> {
215    /// Callback which is called whenever the server read the request from the client and is ready
216    /// to reply to it. The callback returns an optional headers which will be added to the reply
217    /// which the server sends to the user.
218    callback: Option<C>,
219    /// WebSocket configuration.
220    config: Option<WebSocketConfig>,
221    /// Error code/flag. If set, an error will be returned after sending response to the client.
222    error_response: Option<ErrorResponse>,
223    /// Internal stream type.
224    _marker: PhantomData<S>,
225}
226
227impl<S: Read + Write, C: Callback> ServerHandshake<S, C> {
228    /// Start server handshake. `callback` specifies a custom callback which the user can pass to
229    /// the handshake, this callback will be called when the a websocket client connects to the
230    /// server, you can specify the callback if you want to add additional header to the client
231    /// upon join based on the incoming headers.
232    pub fn start(stream: S, callback: C, config: Option<WebSocketConfig>) -> MidHandshake<Self> {
233        trace!("Server handshake initiated.");
234        MidHandshake {
235            machine: HandshakeMachine::start_read(stream),
236            role: ServerHandshake {
237                callback: Some(callback),
238                config,
239                error_response: None,
240                _marker: PhantomData,
241            },
242        }
243    }
244}
245
246impl<S: Read + Write, C: Callback> HandshakeRole for ServerHandshake<S, C> {
247    type IncomingData = Request;
248    type InternalStream = S;
249    type FinalResult = WebSocket<S>;
250
251    fn stage_finished(
252        &mut self,
253        finish: StageResult<Self::IncomingData, Self::InternalStream>,
254    ) -> Result<ProcessingResult<Self::InternalStream, Self::FinalResult>> {
255        Ok(match finish {
256            StageResult::DoneReading { stream, result, tail } => {
257                if !tail.is_empty() {
258                    return Err(Error::Protocol(ProtocolError::JunkAfterRequest));
259                }
260
261                let response = create_response(&result)?;
262                let callback_result = if let Some(callback) = self.callback.take() {
263                    callback.on_request(&result, response)
264                } else {
265                    Ok(response)
266                };
267
268                match callback_result {
269                    Ok(response) => {
270                        let mut output = vec![];
271                        write_response(&mut output, &response)?;
272                        ProcessingResult::Continue(HandshakeMachine::start_write(stream, output))
273                    }
274
275                    Err(resp) => {
276                        if resp.status().is_success() {
277                            return Err(Error::Protocol(ProtocolError::CustomResponseSuccessful));
278                        }
279
280                        self.error_response = Some(resp);
281                        let resp = self.error_response.as_ref().unwrap();
282
283                        let mut output = vec![];
284                        write_response(&mut output, resp)?;
285
286                        if let Some(body) = resp.body() {
287                            output.extend_from_slice(body.as_bytes());
288                        }
289
290                        ProcessingResult::Continue(HandshakeMachine::start_write(stream, output))
291                    }
292                }
293            }
294
295            StageResult::DoneWriting(stream) => {
296                if let Some(err) = self.error_response.take() {
297                    debug!("Server handshake failed.");
298
299                    let (parts, body) = err.into_parts();
300                    let body = body.map(|b| b.as_bytes().to_vec());
301                    return Err(Error::Http(http::Response::from_parts(parts, body).into()));
302                } else {
303                    debug!("Server handshake done.");
304                    let websocket = WebSocket::from_raw_socket(stream, Role::Server, self.config);
305                    ProcessingResult::Done(websocket)
306                }
307            }
308        })
309    }
310}
311
312#[cfg(test)]
313mod tests {
314    use super::{super::machine::TryParse, create_response, Request};
315    use crate::error::{Error, ProtocolError};
316
317    fn request_with_key(key: &str) -> Request {
318        let data = format!(
319            "\
320            GET /script.ws HTTP/1.1\r\n\
321            Host: foo.com\r\n\
322            Connection: upgrade\r\n\
323            Upgrade: websocket\r\n\
324            Sec-WebSocket-Version: 13\r\n\
325            Sec-WebSocket-Key: {key}\r\n\
326            \r\n"
327        );
328
329        let (_, req) = Request::try_parse(data.as_bytes()).unwrap().unwrap();
330        req
331    }
332
333    fn assert_invalid_sec_websocket_key(key: &str) {
334        let req = request_with_key(key);
335        let err = create_response(&req).unwrap_err();
336        assert!(matches!(err, Error::Protocol(ProtocolError::InvalidSecWebSocketKey)));
337    }
338
339    #[test]
340    fn request_parsing() {
341        const DATA: &[u8] = b"GET /script.ws HTTP/1.1\r\nHost: foo.com\r\n\r\n";
342        let (_, req) = Request::try_parse(DATA).unwrap().unwrap();
343        assert_eq!(req.uri().path(), "/script.ws");
344        assert_eq!(req.headers().get("Host").unwrap(), &b"foo.com"[..]);
345    }
346
347    #[test]
348    fn request_replying() {
349        const DATA: &[u8] = b"\
350            GET /script.ws HTTP/1.1\r\n\
351            Host: foo.com\r\n\
352            Connection: upgrade\r\n\
353            Upgrade: websocket\r\n\
354            Sec-WebSocket-Version: 13\r\n\
355            Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n\
356            \r\n";
357        let (_, req) = Request::try_parse(DATA).unwrap().unwrap();
358        let response = create_response(&req).unwrap();
359
360        assert_eq!(
361            response.headers().get("Sec-WebSocket-Accept").unwrap(),
362            b"s3pPLMBiTxaQ9kYGzzhZRbK+xOo=".as_ref()
363        );
364    }
365
366    #[test]
367    fn test_invalid_websocket_key_empty() {
368        assert_invalid_sec_websocket_key("");
369    }
370
371    #[test]
372    fn test_invalid_websocket_key_too_long() {
373        assert_invalid_sec_websocket_key("dGhlIHNhbXBsZSBub25jZQ==AAAAAAAAAA");
374    }
375
376    #[test]
377    fn test_invalid_websocket_key_base64_symbol() {
378        assert_invalid_sec_websocket_key("dGhlIHNhbXBsZSBub25jZQ!!");
379    }
380
381    #[test]
382    fn test_invalid_websocket_key_decoded_length() {
383        assert_invalid_sec_websocket_key("AAAAAAAAAAAAAAAAAAAAAAAA");
384    }
385
386    #[test]
387    fn test_valid_websocket_key() {
388        let req = request_with_key("dGhlIHNhbXBsZSBub25jZQ==");
389        assert!(create_response(&req).is_ok());
390    }
391}
tungstenite/handshake/server.rs

tungstenite/handshake/
server.rs
