initial commit

crates/rsh-backend/src/state.rs

@@ -0,0 +1,71 @@
+use crate::config::Config;
+use dashmap::DashMap;
+use rsh_types::{BackendOpMsg, BackendStubMsg, OpEvent, SessionRecord, StubInfo};
+use ssh_key::PublicKey;
+use std::collections::HashMap;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::Arc;
+use tokio::sync::{broadcast, mpsc, Mutex, RwLock};
+
+pub struct ConnHandle {
+    pub info: StubInfo,
+    pub to_stub: mpsc::Sender<BackendStubMsg>,
+    pub attach: Mutex<Option<AttachSink>>,
+    pub connected_at: i64,
+}
+
+pub struct AttachSink {
+    pub req_id: u64,
+    pub sender: mpsc::Sender<BackendOpMsg>,
+}
+
+pub struct AppState {
+    pub cfg: Config,
+    pub sessions: RwLock<HashMap<String, SessionRecord>>,
+    pub connections: DashMap<(String, u64), Arc<ConnHandle>>,
+    pub next_conn_id: DashMap<String, AtomicU64>,
+    pub authorized_keys: RwLock<Vec<PublicKey>>,
+    pub event_bus: broadcast::Sender<OpEvent>,
+}
+
+impl AppState {
+    pub fn new(cfg: Config) -> Self {
+        let (tx, _) = broadcast::channel(256);
+        Self {
+            cfg,
+            sessions: RwLock::new(HashMap::new()),
+            connections: DashMap::new(),
+            next_conn_id: DashMap::new(),
+            authorized_keys: RwLock::new(Vec::new()),
+            event_bus: tx,
+        }
+    }
+
+    pub fn alloc_conn_id(&self, session: &str) -> u64 {
+        let entry = self
+            .next_conn_id
+            .entry(session.to_string())
+            .or_insert_with(|| AtomicU64::new(1));
+        entry.fetch_add(1, Ordering::Relaxed)
+    }
+
+    pub fn connection_count(&self, session: &str) -> u32 {
+        self.connections
+            .iter()
+            .filter(|kv| kv.key().0 == session)
+            .count() as u32
+    }
+
+    pub fn list_connections(&self, filter: Option<&str>) -> Vec<rsh_types::ConnectionView> {
+        self.connections
+            .iter()
+            .filter(|kv| filter.map_or(true, |s| kv.key().0 == s))
+            .map(|kv| rsh_types::ConnectionView {
+                session_id: kv.key().0.clone(),
+                connection_id: kv.key().1,
+                info: kv.value().info.clone(),
+                connected_at: kv.value().connected_at,
+            })
+            .collect()
+    }
+}