mincomk/rsh
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

1

Browse Source
This commit is contained in:
minco
2026-05-12 23:19:12 +09:00
parent b19d9d25f7
commit cde37d516b
21 changed files with 634 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

131
crates/rsh-backend/src/ws_op.rs
View File

@@ -1,5 +1,6 @@
use crate::auth::{find_key, verify_signature};
use crate::keys::parse_authorized_keys;
use crate::keys::{merge_env_keys, parse_authorized_keys};
use std::sync::atomic::Ordering;
use crate::persist;
use crate::state::{AppState, AttachSink};
use axum::extract::ws::{Message, WebSocket};
@@ -51,7 +52,7 @@ async fn run(mut socket: WebSocket, state: Arc<AppState>) {
let _ = sink.close().await;
});
let mut attached: Option<(String, u64, u64)> = None;
let mut attached: Option<(String, u64, u64, Option<u64>)> = None;
while let Some(Ok(msg)) = stream.next().await {
let text = match msg {
@@ -82,10 +83,16 @@ async fn run(mut socket: WebSocket, state: Arc<AppState>) {
}
}
if let Some((s, c, _)) = attached {
if let Some((s, c, _, shell_id)) = attached {
if let Some(handle) = state.connections.get(&(s, c)) {
let mut a = handle.attach.lock().await;
*a = None;
match shell_id {
None => { *handle.attach.lock().await = None; }
Some(sid) => {
if let Some(slot) = handle.extra_shells.get(&sid) {
*slot.lock().await = None;
}
}
}
}
}
drop(out_tx);
@@ -119,7 +126,7 @@ async fn auth_handshake(socket: &mut WebSocket, state: &Arc<AppState>) -> Result
async fn handle_req(
state: &Arc<AppState>,
out_tx: &mpsc::Sender<BackendOpMsg>,
attached: &mut Option<(String, u64, u64)>,
attached: &mut Option<(String, u64, u64, Option<u64>)>,
req_id: u64,
body: OpReq,
) -> Option<OpResp> {
@@ -187,7 +194,7 @@ async fn handle_req(
OpReq::ConnectionList { session } => {
Some(OpResp::Connections(state.list_connections(session.as_deref())))
}
OpReq::Attach { session, connection_id, pty: _, cols, rows } => {
OpReq::Attach { session, connection_id, shell_id, pty: _, cols, rows } => {
let conn_id = match connection_id {
Some(c) => c,
None => {
@@ -209,42 +216,71 @@ async fn handle_req(
let Some(handle) = state.connections.get(&(session.clone(), conn_id)).map(|h| h.clone()) else {
return Some(OpResp::Err("connection not found".into()));
};
if let Some(sid) = shell_id {
let Some(slot) = handle.extra_shells.get(&sid) else {
return Some(OpResp::Err(format!("no shell {sid}")));
};
*slot.lock().await = Some(AttachSink { req_id, sender: out_tx.clone() });
let _ = handle.to_stub.send(BackendStubMsg::ShellResize { shell_id: sid, cols, rows }).await;
*attached = Some((session, conn_id, req_id, Some(sid)));
return Some(OpResp::AttachReady { connection_id: conn_id });
}
{
let mut a = handle.attach.lock().await;
*a = Some(AttachSink { req_id, sender: out_tx.clone() });
}
let _ = handle.to_stub.send(BackendStubMsg::Resize { cols, rows }).await;
*attached = Some((session, conn_id, req_id));
*attached = Some((session, conn_id, req_id, None));
Some(OpResp::AttachReady { connection_id: conn_id })
}
OpReq::AttachIO(frame) => {
let Some((session, conn_id, _)) = attached.clone() else {
let Some((session, conn_id, _, shell_id)) = attached.clone() else {
return Some(OpResp::Err("not attached".into()));
};
let Some(handle) = state.connections.get(&(session, conn_id)).map(|h| h.clone()) else {
return Some(OpResp::Err("connection gone".into()));
};
match frame {
AttachIOFrame::Stdin(b) => {
let _ = handle.to_stub.send(BackendStubMsg::Stdin(b)).await;
if let Some(sid) = shell_id {
match frame {
AttachIOFrame::Stdin(b) => {
let _ = handle.to_stub.send(BackendStubMsg::ShellStdin { shell_id: sid, data: b }).await;
}
AttachIOFrame::Resize { cols, rows } => {
let _ = handle.to_stub.send(BackendStubMsg::ShellResize { shell_id: sid, cols, rows }).await;
}
AttachIOFrame::Kill | AttachIOFrame::Eof => {
let _ = handle.to_stub.send(BackendStubMsg::ShellKill { shell_id: sid }).await;
}
}
AttachIOFrame::Resize { cols, rows } => {
let _ = handle.to_stub.send(BackendStubMsg::Resize { cols, rows }).await;
}
AttachIOFrame::Kill => {
let _ = handle.to_stub.send(BackendStubMsg::Kill).await;
}
AttachIOFrame::Eof => {
let _ = handle.to_stub.send(BackendStubMsg::Stdin(Vec::new())).await;
} else {
match frame {
AttachIOFrame::Stdin(b) => {
let _ = handle.to_stub.send(BackendStubMsg::Stdin(b)).await;
}
AttachIOFrame::Resize { cols, rows } => {
let _ = handle.to_stub.send(BackendStubMsg::Resize { cols, rows }).await;
}
AttachIOFrame::Kill => {
let _ = handle.to_stub.send(BackendStubMsg::Kill).await;
}
AttachIOFrame::Eof => {
let _ = handle.to_stub.send(BackendStubMsg::Stdin(Vec::new())).await;
}
}
}
None
}
OpReq::Detach => {
if let Some((s, c, _)) = attached.take() {
if let Some((s, c, _, shell_id)) = attached.take() {
if let Some(handle) = state.connections.get(&(s, c)) {
let mut a = handle.attach.lock().await;
*a = None;
match shell_id {
None => { *handle.attach.lock().await = None; }
Some(sid) => {
if let Some(slot) = handle.extra_shells.get(&sid) {
*slot.lock().await = None;
}
}
}
}
}
Some(OpResp::Ok)
@@ -309,6 +345,49 @@ async fn handle_req(
reload_authorized_keys(state, &content).await;
Some(OpResp::Ok)
}
OpReq::SpawnShell { session, connection_id, shell, pty, cols, rows } => {
let conn_id = match connection_id {
Some(c) => c,
None => {
let mut found = None;
for kv in state.connections.iter() {
if kv.key().0 == session {
if found.is_some() {
return Some(OpResp::Err("multiple connections; specify id".into()));
}
found = Some(kv.key().1);
}
}
match found {
Some(c) => c,
None => return Some(OpResp::Err("no connections".into())),
}
}
};
let Some(handle) = state.connections.get(&(session.clone(), conn_id)).map(|h| h.clone()) else {
return Some(OpResp::Err("connection not found".into()));
};
let shell_id = handle.next_shell_id.fetch_add(1, Ordering::Relaxed);
handle.extra_shells.insert(shell_id, tokio::sync::Mutex::new(None));
let (ready_tx, ready_rx) = tokio::sync::oneshot::channel::<()>();
state.spawn_shell_pending.insert((session.clone(), conn_id, shell_id), ready_tx);
if handle.to_stub.send(BackendStubMsg::SpawnShell { shell_id, shell, pty, cols, rows }).await.is_err() {
state.spawn_shell_pending.remove(&(session.clone(), conn_id, shell_id));
handle.extra_shells.remove(&shell_id);
return Some(OpResp::Err("connection lost".into()));
}
match tokio::time::timeout(std::time::Duration::from_secs(10), ready_rx).await {
Ok(Ok(())) => Some(OpResp::ShellSpawned { connection_id: conn_id, shell_id }),
_ => {
state.spawn_shell_pending.remove(&(session, conn_id, shell_id));
handle.extra_shells.remove(&shell_id);
Some(OpResp::Err("shell spawn timed out".into()))
}
}
}
OpReq::Watch { session } => {
let mut rx = state.event_bus.subscribe();
let tx = out_tx.clone();
@@ -358,9 +437,9 @@ async fn disconnect_session(state: &Arc<AppState>, name: &str) {
}
async fn reload_authorized_keys(state: &Arc<AppState>, text: &str) {
let parsed = parse_authorized_keys(text);
let mut k = state.authorized_keys.write().await;
*k = parsed;
let mut keys = parse_authorized_keys(text);
merge_env_keys(&mut keys, &state.env_keys);
*state.authorized_keys.write().await = keys;
}
async fn send(socket: &mut WebSocket, msg: &BackendOpMsg) -> Result<(), axum::Error> {