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,8 @@ 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 +51,16 @@ 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(
+                    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

@@ -1,6 +1,6 @@
 use super::preview_request;
 use crate::commands::dev::utils::{get_path_as_str, rewrite_redirect};
-use crate::commands::dev::{tls, Protocol, ServerConfig};
+use crate::commands::dev::{Protocol, ServerConfig, tls};
 use crate::terminal::emoji;
 use crate::terminal::message::{Message, StdOut};
 use std::sync::{Arc, Mutex};
@@ -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,7 @@ 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 +54,16 @@ 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(
+                    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

@@ -4,11 +4,13 @@ mod https;
 pub use self::http::http;
 pub use self::https::https;
 
-use crate::commands::dev::utils::get_path_as_str;
 use crate::commands::dev::Protocol;
+use crate::commands::dev::utils::get_path_as_str;
 
 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,24 @@ 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

@@ -72,6 +72,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!(