Proxy websocket connections when using authenticated (realish) preview

Previously, Wrangler would return a "101 Switching Protocols" request
and then immediately close the TCP connection. This changes it to instead
continue proxying the connection to the remote worker.

This is simpler than `wrangler dev --inspect` (and uses a different codepath) because it only
proxies the TCP connection directly rather than trying to inspect each websocket message.

 # How do I use this?

1. Create a worker that uses WebSockets. I used https://github.com/cloudflare/websocket-template with the `protocol` in template.js changed from `wss` to `ws`.
2. Run `wrangler dev http http`.

Note that `wrangler dev` (i.e. with the default of `https` for the upstream connection) requires you
to first deploy the worker at least once. Without deploying, the remote worker won't receive the
`Connection: Upgrade` and `Upgrade: Websocket` headers.

Unauthenticated preview is not currently supported.

src/commands/dev/edge/server/http.rs

@@ -8,6 +8,7 @@ use std::sync::{Arc, Mutex};
 use anyhow::Result;
 use chrono::prelude::*;
 use hyper::service::{make_service_fn, service_fn};
+use hyper::upgrade::OnUpgrade;
 use hyper::{Body, Client as HyperClient, Server};
 use hyper_rustls::HttpsConnector;
 use tokio::sync::oneshot::{Receiver, Sender};
@@ -34,6 +35,11 @@ pub async fn http(
 
         async move {
             Ok::<_, anyhow::Error>(service_fn(move |req| {
+                let is_websocket = req
+                    .headers()
+                    .get("upgrade")
+                    .map_or(false, |h| h.as_bytes() == b"websocket");
+
                 let client = client.to_owned();
                 let preview_token = preview_token.lock().unwrap().to_owned();
                 let host = host.to_owned();
@@ -48,14 +54,17 @@ pub async fn http(
                 let now: DateTime<Local> = Local::now();
                 let path = get_path_as_str(&parts.uri);
                 async move {
-                    let mut resp = client.request(preview_request(
-                        parts, body,
+                    let mut req = preview_request(
+                        parts,
+                        body,
                         preview_token.to_owned(),
                         host.clone(),
                         upstream_protocol,
-                    ))
-                    .await?;
+                    );
+                    let client_on_upgrade = req.extensions_mut().remove::<OnUpgrade>();
 
+                    let mut resp = client.request(req).await?;
+                    super::maybe_proxy_websocket(is_websocket, client_on_upgrade, &mut resp);
                     rewrite_redirect(&mut resp, &host, &local_host, false);
 
                     println!(

src/commands/dev/edge/server/https.rs

@@ -10,6 +10,7 @@ use chrono::prelude::*;
 use futures_util::{stream::StreamExt, FutureExt};
 
 use hyper::service::{make_service_fn, service_fn};
+use hyper::upgrade::OnUpgrade;
 use hyper::{Body, Client as HyperClient, Server};
 use hyper_rustls::HttpsConnector;
 use tokio::net::TcpListener;
@@ -38,6 +39,10 @@ pub async fn https(
 
         async move {
             Ok::<_, anyhow::Error>(service_fn(move |req| {
+                let is_websocket = req
+                    .headers()
+                    .get("upgrade")
+                    .map_or(false, |h| h.as_bytes() == b"websocket");
                 let client = client.to_owned();
                 let preview_token = preview_token.lock().unwrap().to_owned();
                 let host = host.to_owned();
@@ -52,13 +57,17 @@ pub async fn https(
                 let now: DateTime<Local> = Local::now();
                 let path = get_path_as_str(&parts.uri);
                 async move {
-                    let mut resp = client.request(preview_request(
-                        parts, body,
+                    let mut req = preview_request(
+                        parts,
+                        body,
                         preview_token.to_owned(),
                         host.clone(),
-                        Protocol::Https,
-                    ))
-                    .await?;
+                        Protocol::Http,
+                    );
+
+                    let client_on_upgrade = req.extensions_mut().remove::<OnUpgrade>();
+                    let mut resp = client.request(req).await?;
+                    super::maybe_proxy_websocket(is_websocket, client_on_upgrade, &mut resp);
 
                     rewrite_redirect(&mut resp, &host, &local_host, true);
 

src/commands/dev/edge/server/mod.rs

@@ -8,7 +8,9 @@ use crate::commands::dev::utils::get_path_as_str;
 use crate::commands::dev::Protocol;
 
 use hyper::header::{HeaderName, HeaderValue};
+use hyper::upgrade::OnUpgrade;
 use hyper::{Body, Request};
+use tokio::io::copy_bidirectional;
 
 fn preview_request(
     mut parts: ::http::request::Parts,
@@ -38,3 +40,29 @@ fn preview_request(
 
     Request::from_parts(parts, body)
 }
+
+fn maybe_proxy_websocket(
+    is_websocket: bool,
+    client_on_upgrade: Option<OnUpgrade>,
+    resp: &mut ::http::Response<Body>,
+) {
+    if is_websocket && resp.status() == 101 {
+        if let (Some(client_on_upgrade), Some(upstream_on_upgrade)) = (
+            client_on_upgrade,
+            resp.extensions_mut().remove::<OnUpgrade>(),
+        ) {
+            tokio::spawn(async move {
+                match tokio::try_join!(client_on_upgrade, upstream_on_upgrade) {
+                    Ok((mut client_upgraded, mut server_upgraded)) => {
+                        let proxy_future =
+                            copy_bidirectional(&mut client_upgraded, &mut server_upgraded);
+                        if let Err(err) = proxy_future.await {
+                            log::warn!("could not proxy WebSocket: {}", err);
+                        }
+                    }
+                    Err(e) => log::warn!("could not proxy WebSocket: {}", e),
+                }
+            });
+        }
+    }
+}

src/commands/dev/gcs/server/http.rs

@@ -55,11 +55,9 @@ pub async fn http(server_config: ServerConfig, preview_id: Arc<Mutex<String>>) -
 
                 async move {
                     // send the request to the preview service
-                    let resp = client.request(preview_request(
-                        parts, body,
-                        preview_id.to_owned(),
-                    ))
-                    .await?;
+                    let resp = client
+                        .request(preview_request(parts, body, preview_id.to_owned()))
+                        .await?;
                     let (mut parts, body) = resp.into_parts();
 
                     // format the response for the user
@@ -72,6 +70,8 @@ pub async fn http(server_config: ServerConfig, preview_id: Arc<Mutex<String>>) -
                         false,
                     );
 
+                    // TODO: proxy websocket
+
                     // print information about the response
                     // [2020-04-20 15:25:54] GET example.com/ HTTP/1.1 200 OK
                     println!(

src/commands/dev/gcs/server/https.rs

@@ -60,11 +60,9 @@ pub async fn https(server_config: ServerConfig, preview_id: Arc<Mutex<String>>)
 
                 async move {
                     // send the request to the preview service
-                    let resp = client.request(preview_request(
-                        parts, body,
-                        preview_id.to_owned(),
-                    ))
-                    .await?;
+                    let resp = client
+                        .request(preview_request(parts, body, preview_id.to_owned()))
+                        .await?;
                     let (mut parts, body) = resp.into_parts();
 
                     // format the response for the user