Read `access_token` from WS header (#44)

Cargo.lock

@@ -302,7 +302,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 
 [[package]]
 name = "flodgatt"
-version = "0.2.0"
+version = "0.2.1"
 dependencies = [
  "dotenv 0.14.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "futures 0.1.26 (registry+https://github.com/rust-lang/crates.io-index)",

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "flodgatt"
 description = "A blazingly fast drop-in replacement for the Mastodon streaming api server"
-version = "0.2.0"
+version = "0.2.1"
 authors = ["Daniel Long Sockwell <daniel@codesections.com", "Julian Laubstein <contact@julianlaubstein.de>"]
 edition = "2018"
 

src/parse_client_request/query.rs

@@ -71,7 +71,7 @@ pub fn optional_media_query() -> BoxedFilter<(Media,)> {
 pub struct OptionalAccessToken;
 
 impl OptionalAccessToken {
-    pub fn from_header() -> warp::filters::BoxedFilter<(Option<String>,)> {
+    pub fn from_sse_header() -> warp::filters::BoxedFilter<(Option<String>,)> {
         let from_header = warp::header::header::<String>("authorization").map(|auth: String| {
             match auth.split(' ').nth(1) {
                 Some(s) => Some(s.to_string()),
@@ -80,6 +80,13 @@ impl OptionalAccessToken {
         });
         let no_token = warp::any().map(|| None);
 
+        from_header.or(no_token).unify().boxed()
+    }
+    pub fn from_ws_header() -> warp::filters::BoxedFilter<(Option<String>,)> {
+        let from_header =
+            warp::header::header::<String>("Sec-Websocket-Protocol").map(|auth: String| Some(auth));
+        let no_token = warp::any().map(|| None);
+
         from_header.or(no_token).unify().boxed()
     }
 }

src/parse_client_request/sse.rs

@@ -63,7 +63,7 @@ pub fn extract_user_or_reject() -> BoxedFilter<(User,)> {
     )
     // because SSE requests place their `access_token` in the header instead of in a query
     // parameter, we need to update our Query if the header has a token
-    .and(query::OptionalAccessToken::from_header())
+    .and(query::OptionalAccessToken::from_sse_header())
     .and_then(Query::update_access_token)
     .and_then(User::from_query)
     .boxed()

src/parse_client_request/ws.rs

@@ -30,7 +30,11 @@ fn parse_query() -> BoxedFilter<(Query,)> {
 }
 
 pub fn extract_user_or_reject() -> BoxedFilter<(User,)> {
-    parse_query().and_then(User::from_query).boxed()
+    parse_query()
+        .and(query::OptionalAccessToken::from_ws_header())
+        .and_then(Query::update_access_token)
+        .and_then(User::from_query)
+        .boxed()
 }
 #[cfg(test)]
 mod test {