Properly follow the spec in WebSocket::Constructor

components/net/fetch/methods.rs

@@ -617,7 +617,7 @@ fn should_block_nosniff(request: &Request, response: &Response) -> bool {
 }
 
 /// https://fetch.spec.whatwg.org/#block-bad-port
-fn should_be_blocked_due_to_bad_port(url: &ServoUrl) -> bool {
+pub fn should_be_blocked_due_to_bad_port(url: &ServoUrl) -> bool {
     // Step 1 is not applicable, this function just takes the URL directly.
 
     // Step 2.

components/net/websocket_loader.rs

@@ -2,11 +2,13 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this
  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
 
+use cookie::Cookie;
 use cookie_storage::CookieStorage;
+use fetch::methods::should_be_blocked_due_to_bad_port;
 use http_loader;
-use hyper::header::Host;
-use net_traits::{WebSocketCommunicate, WebSocketConnectData, WebSocketDomAction, WebSocketNetworkEvent};
-use net_traits::MessageData;
+use hyper::header::{Host, SetCookie};
+use net_traits::{CookieSource, MessageData, WebSocketCommunicate};
+use net_traits::{WebSocketConnectData, WebSocketDomAction, WebSocketNetworkEvent};
 use net_traits::hosts::replace_hosts;
 use net_traits::unwrap_websocket_protocol;
 use servo_url::ServoUrl;
@@ -23,72 +25,90 @@ use websocket::sender::Sender;
 use websocket::stream::WebSocketStream;
 use websocket::ws::receiver::Receiver as WSReceiver;
 use websocket::ws::sender::Sender as Sender_Object;
-use websocket::ws::util::url::parse_url;
 
-/// *Establish a WebSocket Connection* as defined in RFC 6455.
-fn establish_a_websocket_connection(resource_url: &ServoUrl, net_url: (Host, String, bool),
-                                    origin: String, protocols: Vec<String>,
+// https://fetch.spec.whatwg.org/#concept-websocket-establish
+fn establish_a_websocket_connection(resource_url: &ServoUrl,
+                                    origin: String,
+                                    protocols: Vec<String>,
                                     cookie_jar: Arc<RwLock<CookieStorage>>)
-    -> WebSocketResult<(Headers, Sender<WebSocketStream>, Receiver<WebSocketStream>)> {
-    let host = Host {
-        hostname: resource_url.host_str().unwrap().to_owned(),
-        port: resource_url.port_or_known_default(),
-    };
+                                    -> WebSocketResult<(Headers, Sender<WebSocketStream>, Receiver<WebSocketStream>)> {
+    // Steps 1-2 are not really applicable here, given we don't exactly go
+    // through the same infrastructure as the Fetch spec.
 
-    let mut request = try!(Client::connect(net_url));
-    request.headers.set(Origin(origin));
-    request.headers.set(host);
+    if should_be_blocked_due_to_bad_port(resource_url) {
+        // Subset of steps 11-12, we inline the bad port check here from the
+        // main fetch algorithm for the same reason steps 1-2 are not
+        // applicable.
+        return Err(WebSocketError::RequestError("Request should be blocked due to bad port."));
+    }
+
+    // Steps 3-7.
+    let net_url = replace_hosts(resource_url);
+    let mut request = try!(Client::connect(net_url.as_url()));
+
+    // Client::connect sets the Host header to the host of the URL that is
+    // passed to it, so we need to reset it afterwards to the correct one.
+    request.headers.set(Host {
+        hostname: resource_url.host_str().unwrap().to_owned(),
+        port: resource_url.port(),
+    });
+
+    // Step 8.
     if !protocols.is_empty() {
         request.headers.set(WebSocketProtocol(protocols.clone()));
-    };
+    }
 
+    // Steps 9-10.
+    // TODO: support for permessage-deflate extension.
+
+    // Subset of step 11.
+    // See step 2 of https://fetch.spec.whatwg.org/#concept-fetch.
+    request.headers.set(Origin(origin));
+
+    // Transitive subset of step 11.
+    // See step 17.1 of https://fetch.spec.whatwg.org/#concept-http-network-or-cache-fetch.
     http_loader::set_request_cookies(&resource_url, &mut request.headers, &cookie_jar);
 
+    // Step 11, somewhat.
     let response = try!(request.send());
+
+    // Step 12, 14.
     try!(response.validate());
 
-    {
-       let protocol_in_use = unwrap_websocket_protocol(response.protocol());
-        if let Some(protocol_name) = protocol_in_use {
-                if !protocols.is_empty() && !protocols.iter().any(|p| (&**p).eq_ignore_ascii_case(protocol_name)) {
-                    return Err(WebSocketError::ProtocolError("Protocol in Use not in client-supplied protocol list"));
-            };
+    // Step 13 and transitive subset of step 14.
+    // See step 6 of http://tools.ietf.org/html/rfc6455#section-4.1.
+    if let Some(protocol_name) = unwrap_websocket_protocol(response.protocol()) {
+            if !protocols.is_empty() && !protocols.iter().any(|p| (&**p).eq_ignore_ascii_case(protocol_name)) {
+                return Err(WebSocketError::ProtocolError("Protocol in Use not in client-supplied protocol list"));
         };
+    };
+
+    // Transitive subset of step 11.
+    // See step 15 of https://fetch.spec.whatwg.org/#http-network-fetch.
+    if let Some(cookies) = response.headers.get::<SetCookie>() {
+        let mut jar = cookie_jar.write().unwrap();
+        for cookie in &**cookies {
+            if let Some(cookie) = Cookie::new_wrapped(cookie.clone(), resource_url, CookieSource::HTTP) {
+                jar.push(cookie, resource_url, CookieSource::HTTP);
+            }
+        }
     }
 
     let headers = response.headers.clone();
     let (sender, receiver) = response.begin().split();
     Ok((headers, sender, receiver))
-
 }
 
 pub fn init(connect: WebSocketCommunicate, connect_data: WebSocketConnectData, cookie_jar: Arc<RwLock<CookieStorage>>) {
     thread::Builder::new().name(format!("WebSocket connection to {}", connect_data.resource_url)).spawn(move || {
-        // Step 8: Protocols.
-
-        // Step 9.
-
-        // URL that we actually fetch from the network, after applying the replacements
-        // specified in the hosts file.
-        let net_url_result = parse_url(replace_hosts(&connect_data.resource_url).as_url());
-        let net_url = match net_url_result {
-            Ok(net_url) => net_url,
-            Err(e) => {
-                debug!("Failed to establish a WebSocket connection: {:?}", e);
-                let _ = connect.event_sender.send(WebSocketNetworkEvent::Fail);
-                return;
-            }
-        };
         let channel = establish_a_websocket_connection(&connect_data.resource_url,
-                                                       net_url,
                                                        connect_data.origin,
-                                                       connect_data.protocols.clone(),
+                                                       connect_data.protocols,
                                                        cookie_jar);
-        let (_, ws_sender, mut receiver) = match channel {
-            Ok(channel) => {
-                let _ = connect.event_sender.send(WebSocketNetworkEvent::ConnectionEstablished(channel.0.clone(),
-                                                                                               connect_data.protocols));
-                channel
+        let (ws_sender, mut receiver) = match channel {
+            Ok((headers, sender, receiver)) => {
+                let _ = connect.event_sender.send(WebSocketNetworkEvent::ConnectionEstablished(headers));
+                (sender, receiver)
             },
             Err(e) => {
                 debug!("Failed to establish a WebSocket connection: {:?}", e);