Pin local daemon WS server to the first connecting extension

crates/bsk-cli/src/daemon/extension_pin.rs

@@ -0,0 +1,127 @@
+//! Pins the local WS server to the single browser-extension origin that
+//! first connects to it, closing the gap documented in `ws.rs::origin_allowed`:
+//! that check only validates that an `Origin` header is *shaped* like a
+//! Chrome extension origin (`chrome-extension://<32 a-p chars>`), not that it
+//! is *the* BrowserSkill extension. Any other extension installed in the same
+//! browser (malicious, compromised, or side-loaded) is shaped identically and
+//! would otherwise be accepted, giving it full control of the daemon and, in
+//! turn, of the user's already-logged-in browser sessions.
+//!
+//! Pinning is deliberately simple: the first origin to complete a WS upgrade
+//! after the store is empty is trusted and persisted; every later connection
+//! must match it exactly. `bsk daemon reset-pin` clears the file so the user
+//! can re-pair after reinstalling or switching extensions.
+
+use std::fs;
+use std::io::Write;
+use std::path::PathBuf;
+
+use anyhow::{Context, Result};
+use serde::{Deserialize, Serialize};
+
+#[derive(Default, Serialize, Deserialize)]
+struct State {
+    pinned_origin: Option<String>,
+}
+
+#[derive(Debug, Clone)]
+pub struct ExtensionPinStore {
+    /// `None` when the daemon home directory couldn't be resolved (mirrors
+    /// `AuditStore`'s degrade-gracefully convention). No pin can be persisted
+    /// in that case, so every connection is treated as unpinned for that
+    /// process's lifetime — a narrower window than the pre-fix behavior
+    /// (which never pinned at all), not a regression.
+    path: Option<PathBuf>,
+}
+
+impl ExtensionPinStore {
+    pub fn new(home: Option<PathBuf>) -> Self {
+        Self {
+            path: home.map(|home| home.join("extension-pin.json")),
+        }
+    }
+
+    fn read_state(&self) -> Result<State> {
+        let Some(path) = self.path.as_deref() else {
+            return Ok(State::default());
+        };
+        match fs::read(path) {
+            Ok(bytes) => serde_json::from_slice(&bytes).context("invalid extension pin store"),
+            Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(State::default()),
+            Err(err) => Err(err.into()),
+        }
+    }
+
+    fn write_state(&self, state: &State) -> Result<()> {
+        let Some(path) = self.path.as_deref() else {
+            return Ok(());
+        };
+        let parent = path.parent().context("extension pin directory missing")?;
+        fs::create_dir_all(parent)?;
+        let mut file = tempfile::NamedTempFile::new_in(parent)?;
+        serde_json::to_writer(file.as_file_mut(), state)?;
+        file.as_file_mut().flush()?;
+        file.as_file().sync_all()?;
+        file.persist(path).map_err(|err| err.error)?;
+        Ok(())
+    }
+
+    /// The currently pinned origin, if any.
+    pub fn get(&self) -> Result<Option<String>> {
+        Ok(self.read_state()?.pinned_origin)
+    }
+
+    /// Pin `origin` if nothing is pinned yet. Returns `true` if this call did
+    /// the pinning, `false` if an origin was already pinned (regardless of
+    /// whether it matches `origin` — the caller is expected to have already
+    /// rejected a mismatch before calling this).
+    pub fn pin_if_empty(&self, origin: &str) -> Result<bool> {
+        let mut state = self.read_state()?;
+        if state.pinned_origin.is_some() {
+            return Ok(false);
+        }
+        state.pinned_origin = Some(origin.to_string());
+        self.write_state(&state)?;
+        Ok(true)
+    }
+
+    /// Clear the pin so the next connecting extension is trusted anew.
+    pub fn reset(&self) -> Result<()> {
+        self.write_state(&State::default())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn first_connection_pins_and_later_ones_read_it_back() {
+        let dir = tempfile::tempdir().unwrap();
+        let store = ExtensionPinStore::new(Some(dir.path().to_path_buf()));
+        assert_eq!(store.get().unwrap(), None);
+
+        assert!(store.pin_if_empty("chrome-extension://aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa").unwrap());
+        assert_eq!(
+            store.get().unwrap().as_deref(),
+            Some("chrome-extension://aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")
+        );
+
+        // A second extension trying to claim the pin is a no-op.
+        assert!(!store.pin_if_empty("chrome-extension://bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb").unwrap());
+        assert_eq!(
+            store.get().unwrap().as_deref(),
+            Some("chrome-extension://aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")
+        );
+    }
+
+    #[test]
+    fn reset_clears_the_pin() {
+        let dir = tempfile::tempdir().unwrap();
+        let store = ExtensionPinStore::new(Some(dir.path().to_path_buf()));
+        store.pin_if_empty("chrome-extension://aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa").unwrap();
+        store.reset().unwrap();
+        assert_eq!(store.get().unwrap(), None);
+        assert!(store.pin_if_empty("chrome-extension://cccccccccccccccccccccccccccccccc").unwrap());
+    }
+}

crates/bsk-cli/src/daemon/mod.rs

@@ -4,6 +4,7 @@ pub mod abort;
 pub mod audit;
 pub mod browsers;
 mod cancel_forward;
+pub mod extension_pin;
 pub mod file_transfer;
 pub mod inflight;
 pub mod info;

crates/bsk-cli/src/daemon/state.rs

@@ -56,6 +56,11 @@ pub struct DaemonState {
     /// Operation-scoped local file staging. The extension only sees paths
     /// minted here; agent-facing RPCs use opaque transfer ids.
     pub transfers: Arc<TransferRegistry>,
+    /// Pins the local WS server to the extension origin that first connects,
+    /// so any other browser extension shaped like an extension origin (see
+    /// `ws::origin_allowed`) can't also drive the daemon. See
+    /// `extension_pin` module docs.
+    pub extension_pin: Arc<super::extension_pin::ExtensionPinStore>,
 }
 
 impl DaemonState {
@@ -76,6 +81,9 @@ impl DaemonState {
         let abort_registry = Arc::new(AbortRegistry::new());
         let session_interrupts = Arc::new(SessionInterruptRegistry::new());
         let transfers = Arc::new(TransferRegistry::new().expect("initialise transfer staging"));
+        let extension_pin = Arc::new(super::extension_pin::ExtensionPinStore::new(
+            super::paths::bsk_home().ok(),
+        ));
         Self {
             audit,
             config,
@@ -86,6 +94,7 @@ impl DaemonState {
             tool_inflight,
             session_interrupts,
             transfers,
+            extension_pin,
         }
     }
 }

crates/bsk-cli/src/daemon/ws.rs

@@ -35,13 +35,14 @@ use super::state::{
 
 /// Result of the optional Origin allow-list check.
 ///
-/// TODO(M10/M12): pair v0.1 GA with an actual extension-id allow-list
-/// (`DaemonConfig::allowed_extension_ids: HashSet<String>`) populated
-/// from a config file / pairing flow, rather than accepting any
-/// extension-shaped origin. Review M4/M5 I8 — acceptable defense-in-
-/// depth gap for now because pairing happens through the popup, but
-/// a side-loaded extension on the same machine currently passes the
-/// gate.
+/// This only validates that `origin` is *shaped* like a Chrome extension
+/// origin (`chrome-extension://<32 a-p chars>`) — any extension installed in
+/// the same browser passes it identically, not just the BrowserSkill
+/// extension. `handle_connection` closes that gap afterward by pinning the
+/// daemon to the first extension origin that completes a WS upgrade
+/// (`state.extension_pin`) and rejecting any other one from then on, so a
+/// side-loaded or otherwise-installed extension can no longer piggyback on
+/// this shape check alone to drive the daemon.
 pub(super) fn origin_allowed(origin: &str, allow_any: bool) -> bool {
     if allow_any {
         return true;
@@ -152,7 +153,7 @@ async fn handle_connection(
         Ok(resp)
     };
 
-    let ws = match tokio_tungstenite::accept_hdr_async(stream, callback).await {
+    let mut ws = match tokio_tungstenite::accept_hdr_async(stream, callback).await {
         Ok(ws) => ws,
         Err(err) => {
             debug!(?peer, %err, "ws handshake rejected");
@@ -162,6 +163,41 @@ async fn handle_connection(
 
     let origin = captured_origin.lock().unwrap().clone().unwrap_or_default();
     debug!(?peer, %origin, "ws connection upgraded");
+
+    // `origin_allowed` above only checks that `origin` is *shaped* like a
+    // Chrome extension origin, not that it's the BrowserSkill extension —
+    // any other extension in the same browser passes that check equally.
+    // Pin the daemon to the first extension that connects and reject any
+    // other one from then on, closing that gap without requiring the full
+    // pairing-UI allow-list this was originally deferred to (see the
+    // `origin_allowed` doc comment).
+    if !allow_any {
+        match state.extension_pin.get() {
+            Ok(Some(pinned)) if pinned != origin => {
+                warn!(
+                    ?origin,
+                    pinned = %pinned,
+                    "rejecting ws connection from an extension origin other than the pinned one"
+                );
+                let _ = ws
+                    .send(Message::Close(Some(CloseFrame {
+                        code: tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode::Policy,
+                        reason: "extension origin does not match the pinned extension".into(),
+                    })))
+                    .await;
+                return Ok(());
+            }
+            Ok(_) => {
+                if let Err(err) = state.extension_pin.pin_if_empty(&origin) {
+                    warn!(%err, "failed to persist extension pin; continuing unpinned for this connection");
+                }
+            }
+            Err(err) => {
+                warn!(%err, "failed to read extension pin store; continuing unpinned for this connection");
+            }
+        }
+    }
+
     drive_connection(state, ws, None).await
 }
 

crates/bsk-cli/tests/extension_pin_ws.rs

@@ -0,0 +1,83 @@
+//! Live end-to-end proof that the daemon's local WS server pins itself to
+//! the first connecting extension origin and rejects any other one, closing
+//! the gap `daemon::ws::origin_allowed`'s own doc comment used to describe:
+//! any extension-shaped origin (not just the real BrowserSkill extension)
+//! passed that check.
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use bsk::daemon::extension_pin::ExtensionPinStore;
+use bsk::daemon::ws::WsServer;
+use bsk::daemon::{DaemonConfig, DaemonState};
+use futures_util::StreamExt;
+use tokio_tungstenite::tungstenite::client::IntoClientRequest;
+
+const EXT_A: &str = "chrome-extension://aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";
+const EXT_B: &str = "chrome-extension://bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb";
+
+/// Complete the WS opening handshake (which only checks `origin_allowed`'s
+/// *shape* rule) and then read the next frame, which is where the pin check
+/// actually rejects a mismatched origin -- as a server-initiated Close, not
+/// as an HTTP-level upgrade failure. `Ok(true)` means the server accepted
+/// the origin (kept the socket open long enough to matter); `Ok(false)`
+/// means it closed the connection right after upgrade.
+async fn try_connect(addr: std::net::SocketAddr, origin: &str) -> bool {
+    let mut request = format!("ws://{addr}/").into_client_request().unwrap();
+    request
+        .headers_mut()
+        .insert("Origin", origin.parse().unwrap());
+    let (mut ws, _) = tokio_tungstenite::connect_async(request)
+        .await
+        .expect("ws opening handshake (shape check) should always succeed for a valid origin");
+    match tokio::time::timeout(Duration::from_secs(2), ws.next()).await {
+        // Server closed the socket right after upgrade: rejected.
+        Ok(Some(Ok(tokio_tungstenite::tungstenite::Message::Close(_)))) => false,
+        Ok(Some(Err(_))) | Ok(None) => false,
+        // Anything else (no immediate close, e.g. the connection stays open
+        // waiting for our system.handshake) means the pin check let it through.
+        Ok(Some(Ok(_))) => true,
+        Err(_) => true,
+    }
+}
+
+#[tokio::test]
+async fn second_different_extension_origin_is_rejected_after_the_first_pins() {
+    tokio::time::timeout(Duration::from_secs(20), async {
+        let temp = tempfile::tempdir().unwrap();
+        let mut state = DaemonState::new(DaemonConfig::new(0));
+        state.extension_pin = Arc::new(ExtensionPinStore::new(Some(temp.path().to_path_buf())));
+        let state = Arc::new(state);
+        let server = WsServer::new(Arc::clone(&state))
+            .bind("127.0.0.1:0".parse().unwrap())
+            .await
+            .unwrap();
+
+        // Nothing pinned yet: EXT_A completes the upgrade and becomes the pin.
+        assert!(
+            try_connect(server.local_addr, EXT_A).await,
+            "first extension (unpinned state) should be accepted"
+        );
+        assert_eq!(
+            state.extension_pin.get().unwrap().as_deref(),
+            Some(EXT_A),
+            "the first successful connection must have pinned its origin"
+        );
+
+        // EXT_B is shaped exactly like a valid extension origin -- before the
+        // fix, `origin_allowed` alone gated this and EXT_B would also be
+        // accepted. It must now be rejected because it isn't the pinned one.
+        assert!(
+            !try_connect(server.local_addr, EXT_B).await,
+            "a second, different extension origin must be rejected once a pin exists"
+        );
+
+        // The legitimate, pinned extension can still reconnect freely.
+        assert!(
+            try_connect(server.local_addr, EXT_A).await,
+            "the pinned extension must still be able to reconnect"
+        );
+    })
+    .await
+    .expect("test timed out");
+}