net/
websocket_loader.rs

1/* This Source Code Form is subject to the terms of the Mozilla Public
2 * License, v. 2.0. If a copy of the MPL was not distributed with this
3 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
4
5//! The websocket handler has three main responsibilities:
6//! 1) initiate the initial HTTP connection and process the response
7//! 2) ensure any DOM requests for sending/closing are propagated to the network
8//! 3) transmit any incoming messages/closing to the DOM
9//!
10//! In order to accomplish this, the handler uses a long-running loop that selects
11//! over events from the network and events from the DOM, using async/await to avoid
12//! the need for a dedicated thread per websocket.
13
14use std::sync::Arc;
15use std::sync::atomic::{AtomicBool, Ordering};
16
17use async_tungstenite::WebSocketStream;
18use async_tungstenite::tokio::{ConnectStream, client_async_tls_with_connector_and_config};
19use futures::stream::StreamExt;
20use headers::{
21    Authorization, Connection, HeaderMapExt, SecWebsocketKey, SecWebsocketVersion, Upgrade,
22};
23use http::HeaderMap;
24use http::header::{self, HeaderName, HeaderValue};
25use ipc_channel::ipc::IpcSender;
26use log::{debug, trace, warn};
27use net_traits::request::{RequestBuilder, RequestMode};
28use net_traits::{CookieSource, MessageData, WebSocketDomAction, WebSocketNetworkEvent};
29use servo_base::generic_channel::CallbackSetter;
30use servo_url::ServoUrl;
31use tokio::net::TcpStream;
32use tokio::select;
33use tokio::sync::mpsc::{UnboundedReceiver, unbounded_channel};
34use tokio_rustls::TlsConnector;
35use tungstenite::error::{Error, ProtocolError, UrlError};
36use tungstenite::handshake::client::Response;
37use tungstenite::protocol::CloseFrame;
38use tungstenite::{ClientRequestBuilder, Message};
39
40use crate::async_runtime::spawn_task;
41use crate::connector::TlsConfig;
42use crate::cookie::ServoCookie;
43use crate::hosts::replace_host;
44use crate::http_loader::HttpState;
45
46/// Create a Request object for the initial HTTP request.
47/// This request contains `Origin`, `Sec-WebSocket-Protocol`, `Authorization`,
48/// and `Cookie` headers as appropriate.
49/// Returns an error if any header values are invalid or tungstenite cannot create
50/// the desired request.
51pub fn create_handshake_request(
52    request: RequestBuilder,
53    http_state: Arc<HttpState>,
54) -> Result<net_traits::request::Request, Error> {
55    let origin = request.url.origin();
56
57    let mut headers = HeaderMap::new();
58    headers.insert(
59        "Origin",
60        HeaderValue::from_str(&request.url.origin().ascii_serialization())?,
61    );
62
63    let host = format!(
64        "{}",
65        origin
66            .host()
67            .ok_or_else(|| Error::Url(UrlError::NoHostName))?
68    );
69    headers.insert("Host", HeaderValue::from_str(&host)?);
70
71    // https://websockets.spec.whatwg.org/#concept-websocket-establish
72    // 3. Append (`Upgrade`, `websocket`) to request’s header list.
73    headers.typed_insert(Upgrade::websocket());
74
75    // 4. Append (`Connection`, `Upgrade`) to request’s header list.
76    headers.typed_insert(Connection::upgrade());
77
78    // 5. Let keyValue be a nonce consisting of a randomly selected 16-byte value that has been
79    // forgiving-base64-encoded and isomorphic encoded.
80    let mut nonce: [u8; 16] = [0; 16];
81    rand::fill(&mut nonce);
82    let sec_websocket_key_header: SecWebsocketKey = nonce.into();
83
84    // 6. Append (`Sec-WebSocket-Key`, keyValue) to request’s header list.
85    headers.typed_insert(sec_websocket_key_header);
86
87    // 7. Append (`Sec-WebSocket-Version`, `13`) to request’s header list.
88    headers.typed_insert(SecWebsocketVersion::V13);
89
90    // 8. For each protocol in protocols, combine (`Sec-WebSocket-Protocol`, protocol) in request’s
91    // header list.
92    let protocols = match request.mode {
93        RequestMode::WebSocket {
94            ref protocols,
95            original_url: _,
96        } => protocols,
97        _ => unreachable!("How did we get here?"),
98    };
99    if !protocols.is_empty() {
100        let protocols = protocols.join(",");
101        headers.insert("Sec-WebSocket-Protocol", HeaderValue::from_str(&protocols)?);
102    }
103
104    let mut cookie_jar = http_state.cookie_jar.write();
105    cookie_jar.remove_expired_cookies_for_url(&request.url);
106    if let Some(cookie_list) = cookie_jar.cookies_for_url(&request.url, CookieSource::HTTP) {
107        headers.insert("Cookie", HeaderValue::from_str(&cookie_list)?);
108    }
109
110    if request.url.password().is_some() || request.url.username() != "" {
111        headers.typed_insert(Authorization::basic(
112            request.url.username(),
113            request.url.password().unwrap_or(""),
114        ));
115    }
116    Ok(request.headers(headers).build())
117}
118
119/// Process an HTTP response resulting from a WS handshake.
120/// This ensures that any `Cookie` or HSTS headers are recognized.
121/// Returns an error if the protocol selected by the handshake doesn't
122/// match the list of provided protocols in the original request.
123fn process_ws_response(
124    http_state: &HttpState,
125    response: &Response,
126    resource_url: &ServoUrl,
127    protocols: &[String],
128) -> Result<Option<String>, Error> {
129    trace!("processing websocket http response for {}", resource_url);
130    let mut protocol_in_use = None;
131    if let Some(protocol_name) = response.headers().get("Sec-WebSocket-Protocol") {
132        let protocol_name = protocol_name.to_str().unwrap_or("");
133        if !protocols.is_empty() && !protocols.iter().any(|p| protocol_name == (*p)) {
134            return Err(Error::Protocol(ProtocolError::InvalidHeader(Box::new(
135                HeaderName::from_static("sec-websocket-protocol"),
136            ))));
137        }
138        protocol_in_use = Some(protocol_name.to_string());
139    }
140
141    let mut jar = http_state.cookie_jar.write();
142    // TODO(eijebong): Replace thise once typed headers settled on a cookie impl
143    for cookie in response.headers().get_all(header::SET_COOKIE) {
144        let cookie_bytes = cookie.as_bytes();
145        if !ServoCookie::is_valid_name_or_value(cookie_bytes) {
146            continue;
147        }
148        if let Ok(s) = std::str::from_utf8(cookie_bytes) &&
149            let Some(cookie) =
150                ServoCookie::from_cookie_string(s, resource_url, CookieSource::HTTP)
151        {
152            jar.push(cookie, resource_url, CookieSource::HTTP);
153        }
154    }
155
156    http_state
157        .hsts_list
158        .write()
159        .update_hsts_list_from_response(resource_url, response.headers());
160
161    Ok(protocol_in_use)
162}
163
164#[derive(Debug)]
165enum DomMsg {
166    Send(Message),
167    Close(Option<(u16, String)>),
168}
169
170/// Initialize a listener for DOM actions. These are routed from the IPC channel
171/// to a tokio channel that the main WS client task uses to receive them.
172fn setup_dom_listener(
173    dom_action_receiver: CallbackSetter<WebSocketDomAction>,
174    initiated_close: Arc<AtomicBool>,
175) -> UnboundedReceiver<DomMsg> {
176    let (sender, receiver) = unbounded_channel();
177
178    dom_action_receiver.set_callback(move |message| {
179        let dom_action = message.expect("Ws dom_action message to deserialize");
180        trace!("handling WS DOM action: {:?}", dom_action);
181        match dom_action {
182            WebSocketDomAction::SendMessage(MessageData::Text(data)) => {
183                if let Err(e) = sender.send(DomMsg::Send(Message::Text(data.into()))) {
184                    warn!("Error sending websocket message: {:?}", e);
185                }
186            },
187            WebSocketDomAction::SendMessage(MessageData::Binary(data)) => {
188                if let Err(e) = sender.send(DomMsg::Send(Message::Binary(data.into()))) {
189                    warn!("Error sending websocket message: {:?}", e);
190                }
191            },
192            WebSocketDomAction::Close(code, reason) => {
193                if initiated_close.fetch_or(true, Ordering::SeqCst) {
194                    return;
195                }
196                let frame = code.map(move |c| (c, reason.unwrap_or_default()));
197                if let Err(e) = sender.send(DomMsg::Close(frame)) {
198                    warn!("Error closing websocket: {:?}", e);
199                }
200            },
201        }
202    });
203
204    receiver
205}
206
207/// Listen for WS events from the DOM and the network until one side
208/// closes the connection or an error occurs. Since this is an async
209/// function that uses the select operation, it will run as a task
210/// on the WS tokio runtime.
211async fn run_ws_loop(
212    mut dom_receiver: UnboundedReceiver<DomMsg>,
213    resource_event_sender: IpcSender<WebSocketNetworkEvent>,
214    mut stream: WebSocketStream<ConnectStream>,
215) {
216    loop {
217        select! {
218            dom_msg = dom_receiver.recv() => {
219                trace!("processing dom msg: {:?}", dom_msg);
220                let dom_msg = match dom_msg {
221                    Some(msg) => msg,
222                    None => break,
223                };
224                match dom_msg {
225                    DomMsg::Send(m) => {
226                        if let Err(e) = stream.send(m).await {
227                            warn!("error sending websocket message: {:?}", e);
228                        }
229                    },
230                    DomMsg::Close(frame) => {
231                        if let Err(e) = stream.close(frame.map(|(code, reason)| {
232                            CloseFrame {
233                                code: code.into(),
234                                reason: reason.into(),
235                            }
236                        })).await {
237                            warn!("error closing websocket: {:?}", e);
238                        }
239                    },
240                }
241            }
242            ws_msg = stream.next() => {
243                trace!("processing WS stream: {:?}", ws_msg);
244                let msg = match ws_msg {
245                    Some(Ok(msg)) => msg,
246                    Some(Err(e)) => {
247                        warn!("Error in WebSocket communication: {:?}", e);
248                        let _ = resource_event_sender.send(WebSocketNetworkEvent::Fail);
249                        break;
250                    },
251                    None => {
252                        warn!("Error in WebSocket communication");
253                        let _ = resource_event_sender.send(WebSocketNetworkEvent::Fail);
254                        break;
255                    }
256                };
257                match msg {
258                    Message::Text(s) => {
259                        let message = MessageData::Text(s.as_str().to_owned());
260                        if let Err(e) = resource_event_sender
261                            .send(WebSocketNetworkEvent::MessageReceived(message))
262                        {
263                            warn!("Error sending websocket notification: {:?}", e);
264                            break;
265                        }
266                    }
267
268                    Message::Binary(v) => {
269                        let message = MessageData::Binary(v.to_vec());
270                        if let Err(e) = resource_event_sender
271                            .send(WebSocketNetworkEvent::MessageReceived(message))
272                        {
273                            warn!("Error sending websocket notification: {:?}", e);
274                            break;
275                        }
276                    }
277
278                    Message::Ping(_) | Message::Pong(_) => {}
279
280                    Message::Close(frame) => {
281                        let (reason, code) = match frame {
282                            Some(frame) => (frame.reason, Some(frame.code.into())),
283                            None => ("".into(), None),
284                        };
285                        debug!("Websocket connection closing due to ({:?}) {}", code, reason);
286                        let _ = resource_event_sender.send(WebSocketNetworkEvent::Close(
287                            code,
288                            reason.to_string(),
289                        ));
290                        break;
291                    }
292
293                    Message::Frame(_) => {
294                        warn!("Unexpected websocket frame message");
295                    }
296                }
297            }
298        }
299    }
300}
301
302/// Initiate a new async WS connection. Returns an error if the connection fails
303/// for any reason, or if the response isn't valid. Otherwise, the endless WS
304/// listening loop will be started.
305pub(crate) async fn start_websocket(
306    http_state: Arc<HttpState>,
307    resource_event_sender: IpcSender<WebSocketNetworkEvent>,
308    protocols: &[String],
309    client: &net_traits::request::Request,
310    tls_config: TlsConfig,
311    dom_action_receiver: CallbackSetter<WebSocketDomAction>,
312) -> Result<Response, Error> {
313    trace!("starting WS connection to {}", client.url());
314
315    let initiated_close = Arc::new(AtomicBool::new(false));
316    let dom_receiver = setup_dom_listener(dom_action_receiver, initiated_close.clone());
317
318    let url = client.url();
319    let host = replace_host(url.host_str().expect("URL has no host"));
320    let mut net_url = client.url().into_url();
321    net_url
322        .set_host(Some(&host))
323        .map_err(|e| Error::Url(UrlError::UnableToConnect(e.to_string())))?;
324
325    let domain = net_url
326        .host()
327        .ok_or_else(|| Error::Url(UrlError::NoHostName))?;
328    let port = net_url
329        .port_or_known_default()
330        .ok_or_else(|| Error::Url(UrlError::UnableToConnect("Unknown port".into())))?;
331
332    let try_socket = TcpStream::connect((&*domain.to_string(), port)).await;
333    let socket = try_socket.map_err(Error::Io)?;
334    let connector = TlsConnector::from(Arc::new(tls_config));
335
336    // TODO(pylbrecht): move request conversion to a separate function
337    let mut original_url = client.original_url();
338    if original_url.scheme() == "ws" && url.scheme() == "https" {
339        original_url.as_mut_url().set_scheme("wss").unwrap();
340    }
341    let mut builder =
342        ClientRequestBuilder::new(original_url.as_str().parse().expect("unable to parse URI"));
343    for (key, value) in client.headers.iter() {
344        builder = builder.with_header(
345            key.as_str(),
346            value
347                .to_str()
348                .expect("unable to convert header value to string"),
349        );
350    }
351
352    let (stream, response) =
353        client_async_tls_with_connector_and_config(builder, socket, Some(connector), None).await?;
354
355    let protocol_in_use = process_ws_response(&http_state, &response, &url, protocols)?;
356
357    if !initiated_close.load(Ordering::SeqCst) {
358        if resource_event_sender
359            .send(WebSocketNetworkEvent::ConnectionEstablished { protocol_in_use })
360            .is_err()
361        {
362            return Ok(response);
363        }
364
365        trace!("about to start ws loop for {}", url);
366        spawn_task(run_ws_loop(dom_receiver, resource_event_sender, stream));
367    } else {
368        trace!("client closed connection for {}, not running loop", url);
369    }
370    Ok(response)
371}
net/websocket_loader.rs

net/
websocket_loader.rs
