add support for upgrades

Author: jbr

Cargo.toml

@@ -14,7 +14,7 @@ edition = "2018"
 [dependencies]
 httparse = "1.3.3"
 async-std = { version = "1.6.0", features = ["unstable"] }
-http-types = "2.0.0"
+http-types = { version = "2.0.0", features = ["unstable"] }
 byte-pool = "0.2.1"
 lazy_static = "1.4.0"
 futures-core = "0.3.1"

src/server/mod.rs

@@ -3,9 +3,10 @@
 use std::time::Duration;
 
 use async_std::future::{timeout, Future, TimeoutError};
-use async_std::io::{self};
-use async_std::io::{Read, Write};
-use http_types::{Request, Response};
+use async_std::io::{self, Read, Write};
+use http_types::headers::{CONNECTION, UPGRADE};
+use http_types::upgrade::Connection;
+use http_types::{Request, Response, StatusCode};
 
 mod decode;
 mod encode;
@@ -70,14 +71,32 @@ where
             }
         };
 
+        let upgrade_requested = match (req.header(UPGRADE), req.header(CONNECTION)) {
+            (Some(_), Some(upgrade)) if upgrade.as_str().eq_ignore_ascii_case("upgrade") => true,
+            _ => false,
+        };
+
         let method = req.method();
+
         // Pass the request to the endpoint and encode the response.
-        let res = endpoint(req).await?;
+        let mut res = endpoint(req).await?;
+
+        let upgrade_provided = res.status() == StatusCode::SwitchingProtocols && res.has_upgrade();
+
+        let upgrade_sender = if upgrade_requested && upgrade_provided {
+            Some(res.send_upgrade())
+        } else {
+            None
+        };
 
         let mut encoder = Encoder::new(res, method);
 
         // Stream the response to the writer.
         io::copy(&mut encoder, &mut io).await?;
+
+        if let Some(upgrade_sender) = upgrade_sender {
+            upgrade_sender.send(Connection::new(io.clone())).await;
+        }
     }
 
     Ok(())