From aa9bc1f78531d071896d41ea52d48cd1ea2a19d7 Mon Sep 17 00:00:00 2001
From: BlackDex <black.dex@gmail.com>
Date: Sun, 27 Aug 2023 22:28:35 +0200
Subject: [PATCH] Allow Authorization header for Web Sockets

Some clients (Thirdparty) might use the `Authorization` header instead
of a query param. We didn't supported this since all the official
clients do not seem to use this way of working. But Bitwarden does check
both ways.

This PR adds an extra check for this header which can be optional.

Fixes #3776
---
 src/api/notifications.rs | 12 ++++++++++--
 src/auth.rs              | 23 +++++++++++++++++++++++
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/src/api/notifications.rs b/src/api/notifications.rs
index 7e76021b..3df91e91 100644
--- a/src/api/notifications.rs
+++ b/src/api/notifications.rs
@@ -20,7 +20,7 @@ use tokio_tungstenite::{
 };
 
 use crate::{
-    auth::ClientIp,
+    auth::{ClientIp, WsAccessTokenHeader},
     db::{
         models::{Cipher, Folder, Send as DbSend, User},
         DbConn,
@@ -111,11 +111,19 @@ fn websockets_hub<'r>(
     ws: rocket_ws::WebSocket,
     data: WsAccessToken,
     ip: ClientIp,
+    header_token: WsAccessTokenHeader,
 ) -> Result<rocket_ws::Stream!['r], Error> {
     let addr = ip.ip;
     info!("Accepting Rocket WS connection from {addr}");
 
-    let Some(token) = data.access_token else { err_code!("Invalid claim", 401) };
+    let token = if let Some(token) = data.access_token {
+        token
+    } else if let Some(token) = header_token.access_token {
+        token
+    } else {
+        err_code!("Invalid claim", 401)
+    };
+
     let Ok(claims) = crate::auth::decode_login(&token) else { err_code!("Invalid token", 401) };
 
     let (mut rx, guard) = {
diff --git a/src/auth.rs b/src/auth.rs
index 6879bb6e..e23aa32d 100644
--- a/src/auth.rs
+++ b/src/auth.rs
@@ -825,3 +825,26 @@ impl<'r> FromRequest<'r> for ClientIp {
         })
     }
 }
+
+pub struct WsAccessTokenHeader {
+    pub access_token: Option<String>,
+}
+
+#[rocket::async_trait]
+impl<'r> FromRequest<'r> for WsAccessTokenHeader {
+    type Error = ();
+
+    async fn from_request(request: &'r Request<'_>) -> Outcome<Self, Self::Error> {
+        let headers = request.headers();
+
+        // Get access_token
+        let access_token = match headers.get_one("Authorization") {
+            Some(a) => a.rsplit("Bearer ").next().map(String::from),
+            None => None,
+        };
+
+        Outcome::Success(Self {
+            access_token,
+        })
+    }
+}