Add token verification for the wasm sdk websocket connection

Author: thlsrms

Cargo.lock

@@ -12,7 +12,9 @@ default = []
 web = [
     "dep:getrandom",
     "dep:gloo-console",
+    "dep:gloo-net",
     "dep:gloo-storage",
+    "dep:js-sys",
     "dep:rustls-pki-types",
     "dep:tokio-tungstenite-wasm",
     "dep:wasm-bindgen",
@@ -42,7 +44,9 @@ rand.workspace = true
 
 getrandom = { version = "0.3.2", features = ["wasm_js"], optional = true }
 gloo-console = { version = "0.3.0", optional = true }
+gloo-net = { version = "0.6.0", optional = true }
 gloo-storage = { version = "0.3.0", optional = true }
+js-sys = { version = "0.3", optional = true }
 rustls-pki-types = { version = "1.12.0", features = ["web"], optional = true }
 tokio-tungstenite-wasm = { version = "0.6.0", optional = true }
 wasm-bindgen = { version = "0.2.100", optional = true }

crates/sdk/Cargo.toml

@@ -1013,7 +1013,6 @@ but you must call one of them, or else the connection will never progress.
     /// If the passed token is invalid or rejected by the host,
     /// the connection will fail asynchrnonously.
     // FIXME: currently this causes `disconnect` to be called rather than `on_connect_error`.
-    #[cfg(not(target_arch = "wasm32"))]
     pub fn with_token(mut self, token: Option<impl Into<String>>) -> Self {
         self.token = token.map(|token| token.into());
         self

crates/sdk/src/db_connection.rs

@@ -98,6 +98,10 @@ pub enum WsError {
 
     #[error("Unrecognized compression scheme: {scheme:#x}")]
     UnknownCompressionScheme { scheme: u8 },
+
+    #[cfg(feature = "web")]
+    #[error("Token verification error: {0}")]
+    TokenVerification(String),
 }
 
 pub(crate) struct WsConnection {
@@ -131,7 +135,29 @@ pub(crate) struct WsParams {
     pub light: bool,
 }
 
+#[cfg(not(feature = "web"))]
 fn make_uri(host: Uri, db_name: &str, connection_id: ConnectionId, params: WsParams) -> Result<Uri, UriError> {
+    make_uri_impl(host, db_name, connection_id, params, None)
+}
+
+#[cfg(feature = "web")]
+fn make_uri(
+    host: Uri,
+    db_name: &str,
+    connection_id: ConnectionId,
+    params: WsParams,
+    token: Option<&str>,
+) -> Result<Uri, UriError> {
+    make_uri_impl(host, db_name, connection_id, params, token)
+}
+
+fn make_uri_impl(
+    host: Uri,
+    db_name: &str,
+    connection_id: ConnectionId,
+    params: WsParams,
+    token: Option<&str>,
+) -> Result<Uri, UriError> {
     let mut parts = host.into_parts();
     let scheme = parse_scheme(parts.scheme.take())?;
     parts.scheme = Some(scheme);
@@ -171,6 +197,11 @@ fn make_uri(host: Uri, db_name: &str, connection_id: ConnectionId, params: WsPar
         path.push_str("&light=true");
     }
 
+    // Specify the `token` param if needed
+    if let Some(token) = token {
+        path.push_str(&format!("&token={token}"));
+    }
+
     parts.path_and_query = Some(path.parse().map_err(|source: InvalidUri| UriError::InvalidUri {
         source: Arc::new(source),
     })?);
@@ -222,6 +253,53 @@ fn request_insert_auth_header(req: &mut http::Request<()>, token: Option<&str>)
     }
 }
 
+#[cfg(feature = "web")]
+async fn fetch_ws_token(host: &Uri, auth_token: &str) -> Result<String, WsError> {
+    use gloo_net::http::{Method, RequestBuilder};
+    use js_sys::{Reflect, JSON};
+    use wasm_bindgen::{JsCast, JsValue};
+
+    let url = format!("{}v1/identity/websocket-token", host);
+
+    // helpers to convert gloo_net::Error or JsValue into WsError::TokenVerification
+    let gloo_to_ws_err = |e: gloo_net::Error| match e {
+        gloo_net::Error::JsError(js_err) => WsError::TokenVerification(js_err.message.into()),
+        gloo_net::Error::SerdeError(e) => WsError::TokenVerification(e.to_string()),
+        gloo_net::Error::GlooError(msg) => WsError::TokenVerification(msg),
+    };
+    let js_to_ws_err = |e: JsValue| {
+        if let Some(err) = e.dyn_ref::<js_sys::Error>() {
+            WsError::TokenVerification(err.message().into())
+        } else if let Some(s) = e.as_string() {
+            WsError::TokenVerification(s)
+        } else {
+            WsError::TokenVerification(format!("{:?}", e))
+        }
+    };
+
+    let res = RequestBuilder::new(&url)
+        .method(Method::POST)
+        .header("Authorization", &format!("Bearer {auth_token}"))
+        .send()
+        .await
+        .map_err(gloo_to_ws_err)?;
+
+    if !res.ok() {
+        return Err(WsError::TokenVerification(format!(
+            "HTTP error: {} {}",
+            res.status(),
+            res.status_text()
+        )));
+    }
+
+    let body = res.text().await.map_err(gloo_to_ws_err)?;
+    let json = JSON::parse(&body).map_err(js_to_ws_err)?;
+    let token_js = Reflect::get(&json, &JsValue::from_str("token")).map_err(js_to_ws_err)?;
+    token_js
+        .as_string()
+        .ok_or_else(|| WsError::TokenVerification("`token` parsing failed".into()))
+}
+
 impl WsConnection {
     #[cfg(not(feature = "web"))]
     pub(crate) async fn connect(
@@ -260,11 +338,17 @@ impl WsConnection {
     pub(crate) async fn connect(
         host: Uri,
         db_name: &str,
-        _token: Option<&str>,
+        token: Option<&str>,
         connection_id: ConnectionId,
         params: WsParams,
     ) -> Result<Self, WsError> {
-        let uri = make_uri(host, db_name, connection_id, params)?;
+        let token = if let Some(auth_token) = token {
+            Some(fetch_ws_token(&host, auth_token).await?)
+        } else {
+            None
+        };
+
+        let uri = make_uri(host, db_name, connection_id, params, token.as_deref())?;
         let sock = tokio_tungstenite_wasm::connect_with_protocols(&uri.to_string(), &[BIN_PROTOCOL])
             .await
             .map_err(|source| WsError::Tungstenite {