1

crates/rsh-backend/src/state.rs

@@ -1,17 +1,20 @@
 use crate::config::Config;
+use crate::keys::parse_authorized_keys;
 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};
+use tokio::sync::{broadcast, mpsc, oneshot, Mutex, RwLock};
 
 pub struct ConnHandle {
     pub info: StubInfo,
     pub to_stub: mpsc::Sender<BackendStubMsg>,
     pub attach: Mutex<Option<AttachSink>>,
     pub connected_at: i64,
+    pub extra_shells: DashMap<u64, Mutex<Option<AttachSink>>>,
+    pub next_shell_id: AtomicU64,
 }
 
 pub struct AttachSink {
@@ -25,19 +28,28 @@ pub struct AppState {
     pub connections: DashMap<(String, u64), Arc<ConnHandle>>,
     pub next_conn_id: DashMap<String, AtomicU64>,
     pub authorized_keys: RwLock<Vec<PublicKey>>,
+    pub env_keys: Vec<PublicKey>,
     pub event_bus: broadcast::Sender<OpEvent>,
+    pub spawn_shell_pending: DashMap<(String, u64, u64), oneshot::Sender<()>>,
 }
 
 impl AppState {
     pub fn new(cfg: Config) -> Self {
         let (tx, _) = broadcast::channel(256);
+        let env_keys = cfg
+            .authorized_keys_env
+            .as_deref()
+            .map(parse_authorized_keys)
+            .unwrap_or_default();
         Self {
             cfg,
             sessions: RwLock::new(HashMap::new()),
             connections: DashMap::new(),
             next_conn_id: DashMap::new(),
             authorized_keys: RwLock::new(Vec::new()),
+            env_keys,
             event_bus: tx,
+            spawn_shell_pending: DashMap::new(),
         }
     }