feat: rshc: add sshserve

This commit is contained in:
200mill
2026-07-04 02:32:54 +09:00
parent 4e515ba964
commit 45554836d8
6 changed files with 1860 additions and 134 deletions

View File

@@ -32,3 +32,5 @@ rand = { workspace = true }
tracing = { workspace = true }
tracing-subscriber = { workspace = true }
bytes = { workspace = true }
russh = "0.62.1"

View File

@@ -4,3 +4,4 @@ pub mod connect;
pub mod keys;
pub mod watch;
pub mod shell;
pub mod serve_ssh;

View File

@@ -0,0 +1,290 @@
use crate::auth::AuthedClient;
use crate::cmd::connection;
use crate::config::Config;
use crate::ui;
use anyhow::{anyhow, Result};
use rsh_types::{AttachIOFrame, OpReq, OpResp};
use russh::server::{Auth, Session};
use russh::ChannelId;
use std::collections::HashMap;
use std::sync::Arc;
use bytes::Bytes;
use tokio::net::TcpListener;
pub async fn run(
cfg: Config,
client: AuthedClient,
session_name: String,
listen_host: String,
listen_port: u16,
connection_id: Option<u64>,
shell: Option<String>,
) -> Result<()> {
// Validate that we have connections before starting the server.
// Also resolve connection_id if not explicitly provided but only 1 exists.
let conns = connection::fetch(&client, Some(session_name.clone())).await?;
if conns.is_empty() {
return Err(anyhow!("no connections for session '{}'", session_name));
}
let conn_id = match connection_id {
Some(id) => {
conns
.iter()
.find(|c| c.connection_id == id)
.ok_or_else(|| anyhow!("no connection {} in session '{}'", id, session_name))?
.connection_id
}
None => {
if conns.len() == 1 {
conns[0].connection_id
} else {
return Err(anyhow!(
"multiple connections found. please specify --connection <id>"
));
}
}
};
let mut config = russh::server::Config {
inactivity_timeout: Some(std::time::Duration::from_secs(3600)),
auth_rejection_time: std::time::Duration::from_secs(1),
auth_rejection_time_initial: Some(std::time::Duration::from_secs(0)),
..Default::default()
};
// Use a static ephemeral key for the SSH server
let key_str = "-----BEGIN OPENSSH PRIVATE KEY-----\nb3BlbnNzaC1rZXktdjEAAAAABG5vbmUAAAAEbm9uZQAAAAAAAAABAAAAMwAAAAtzc2gtZW\nQyNTUxOQAAACBxYVtR5Xn/kk9PM26c//B9JlvqBBv4EbkWr3jRI8CLfQAAAJjAfItSwHyL\nUgAAAAtzc2gtZWQyNTUxOQAAACBxYVtR5Xn/kk9PM26c//B9JlvqBBv4EbkWr3jRI8CLfQ\nAAAEA3iPbbQVsorPV8V36bsOZWFHyYzqAx4SIPBQsVWW4RdHFhW1Hlef+ST08zbpz/8H0m\nW+oEG/gRuRaveNEjwIt9AAAAD2V0eEBldHgtZGVza3RvcAECAwQFBg==\n-----END OPENSSH PRIVATE KEY-----";
let key = russh::keys::PrivateKey::from_openssh(key_str)
.map_err(|e| anyhow!("Failed to parse SSH host key: {}", e))?;
config.keys.push(key);
let config = Arc::new(config);
let bind_addr = format!("{}:{}", listen_host, listen_port);
let listener = TcpListener::bind(&bind_addr).await.map_err(|e| anyhow!(e))?;
ui::print_info(&format!("serving SSH on {}", bind_addr));
loop {
let (stream, _) = match listener.accept().await {
Ok(res) => res,
Err(_) => continue,
};
let handler = SessionHandler {
cfg: cfg.clone(),
session_name: session_name.clone(),
connection_id: conn_id,
shell: shell.clone(),
clients: HashMap::new(),
pty_reqs: HashMap::new(),
};
let config_clone = config.clone();
tokio::spawn(async move {
let _ = russh::server::run_stream(config_clone, stream, handler).await;
});
}
}
struct SessionHandler {
cfg: Config,
session_name: String,
connection_id: u64,
shell: Option<String>,
clients: HashMap<ChannelId, Arc<AuthedClient>>,
pty_reqs: HashMap<ChannelId, (u16, u16)>,
}
impl SessionHandler {
async fn spawn_shell(
&mut self,
channel: ChannelId,
session: &mut Session,
shell_cmd: Option<String>,
) -> Result<(), anyhow::Error> {
let pty = self.pty_reqs.contains_key(&channel);
let (cols, rows) = self.pty_reqs.get(&channel).cloned().unwrap_or((80, 24));
let client = AuthedClient::connect(&self.cfg).await?;
let spawned = client
.req(OpReq::SpawnShell {
session: self.session_name.clone(),
connection_id: Some(self.connection_id),
shell: shell_cmd,
pty,
cols,
rows,
})
.await?;
let shell_id = match spawned {
OpResp::ShellSpawned { shell_id, .. } => shell_id,
OpResp::Err(e) => return Err(anyhow!(e)),
other => return Err(anyhow!("unexpected: {:?}", other)),
};
let (attach_id, mut resps) = client
.req_stream(OpReq::Attach {
session: self.session_name.clone(),
connection_id: Some(self.connection_id),
shell_id: Some(shell_id),
pty,
cols,
rows,
})
.await?;
let ready = resps.recv().await.ok_or_else(|| anyhow!("attach: no response"))?;
match ready {
OpResp::AttachReady { .. } => {}
OpResp::Err(e) => {
client.drop_stream(attach_id).await;
return Err(anyhow!(e));
}
other => {
client.drop_stream(attach_id).await;
return Err(anyhow!("unexpected: {:?}", other));
}
}
let client = Arc::new(client);
self.clients.insert(channel, client.clone());
// Spawn a task to pump output from rsh-backend to russh
let handle = session.handle();
tokio::spawn(async move {
while let Some(resp) = resps.recv().await {
match resp {
OpResp::Stdout(b) => {
let _ = handle.data(channel, Bytes::from(b)).await;
}
OpResp::Stderr(b) => {
let _ = handle.extended_data(channel, 1, Bytes::from(b)).await;
}
OpResp::Exited { code } => {
let _ = handle.exit_status_request(channel, code.unwrap_or(0) as u32).await;
let _ = handle.close(channel).await;
break;
}
OpResp::Err(_) => {
let _ = handle.close(channel).await;
break;
}
_ => {}
}
}
client.drop_stream(attach_id).await;
});
Ok(())
}
}
impl russh::server::Handler for SessionHandler {
type Error = anyhow::Error;
async fn auth_publickey(
&mut self,
_: &str,
_: &russh::keys::PublicKey,
) -> Result<Auth, Self::Error> {
Ok(Auth::Accept)
}
async fn auth_password(&mut self, _: &str, _: &str) -> Result<Auth, Self::Error> {
Ok(Auth::Accept)
}
async fn auth_none(&mut self, _: &str) -> Result<Auth, Self::Error> {
Ok(Auth::Accept)
}
async fn pty_request(
&mut self,
channel: ChannelId,
_term: &str,
col_width: u32,
row_height: u32,
_pix_width: u32,
_pix_height: u32,
_modes: &[(russh::Pty, u32)],
_session: &mut Session,
) -> Result<(), Self::Error> {
self.pty_reqs.insert(channel, (col_width as u16, row_height as u16));
Ok(())
}
async fn shell_request(
&mut self,
channel: ChannelId,
session: &mut Session,
) -> Result<(), Self::Error> {
self.spawn_shell(channel, session, self.shell.clone()).await?;
Ok(())
}
async fn exec_request(
&mut self,
channel: ChannelId,
data: &[u8],
session: &mut Session,
) -> Result<(), Self::Error> {
let command = String::from_utf8_lossy(data).to_string();
self.spawn_shell(channel, session, Some(command)).await?;
Ok(())
}
async fn window_change_request(
&mut self,
channel: ChannelId,
col_width: u32,
row_height: u32,
_pix_width: u32,
_pix_height: u32,
_session: &mut Session,
) -> Result<(), Self::Error> {
if let Some(client) = self.clients.get(&channel) {
let _ = client.send_attach_io(AttachIOFrame::Resize {
cols: col_width as u16,
rows: row_height as u16,
}).await;
}
Ok(())
}
async fn data(
&mut self,
channel: ChannelId,
data: &[u8],
_session: &mut Session,
) -> Result<(), Self::Error> {
if let Some(client) = self.clients.get(&channel) {
let _ = client.send_attach_io(AttachIOFrame::Stdin(data.to_vec())).await;
}
Ok(())
}
async fn channel_eof(
&mut self,
channel: ChannelId,
_session: &mut Session,
) -> Result<(), Self::Error> {
if let Some(client) = self.clients.get(&channel) {
let _ = client.send_attach_io(AttachIOFrame::Eof).await;
}
Ok(())
}
async fn channel_close(
&mut self,
channel: ChannelId,
_session: &mut Session,
) -> Result<(), Self::Error> {
if let Some(client) = self.clients.remove(&channel) {
let _ = client.send_attach_io(AttachIOFrame::Kill).await;
}
self.pty_reqs.remove(&channel);
Ok(())
}
}

View File

@@ -26,6 +26,7 @@ enum Cmd {
Connect(ConnectArgs),
#[command(alias = "sh")]
Shell(ShellArgs),
ServeSsh(ServeSshArgs),
Keys(KeysCmd),
}
@@ -97,6 +98,20 @@ struct ShellArgs {
no_pty: bool,
}
#[derive(Args, Debug)]
struct ServeSshArgs {
session: String,
#[arg(short = 'h', long, default_value = "127.0.0.1")]
listen_host: String,
#[arg(short = 'p', long, default_value_t = 2222)]
listen_port: u16,
#[arg(long)]
connection: Option<u64>,
#[arg(long)]
shell: Option<String>,
}
#[derive(Args, Debug)]
struct KeysCmd {
#[command(subcommand)]
@@ -163,6 +178,8 @@ async fn run() -> Result<()> {
},
Cmd::Connect(a) => cmd::connect::run(&client, a.session, a.connection_id, a.no_pty).await,
Cmd::Shell(a) => cmd::shell::run(&client, a.session, a.connection, a.shell, a.no_pty).await,
Cmd::ServeSsh(a) => cmd::serve_ssh::run(cfg, client, a.session, a.listen_host, a.listen_port, a.connection, a.shell).await,
Cmd::Keys(k) => match k.sub {
KeysSub::Append { key, file, url } => cmd::keys::append(&client, key, file, url).await,
KeysSub::Rm { key, file, url } => cmd::keys::remove(&client, key, file, url).await,