mirrors/pikasTech-unidesk
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add provider ssh tcp data pool

Browse Source
This commit is contained in:
Codex
2026-06-07 02:02:38 +00:00
parent d4b7fc95f9
commit fe1b91dcbd
32 changed files with 988 additions and 350 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/components/backend-core/Dockerfile
+3 -1
View File
@@ -1,9 +1,11 @@
FROM rust:1-bookworm AS build
WORKDIR /app/src/components/backend-core
ARG CARGO_BUILD_JOBS=1
ENV CARGO_BUILD_JOBS=${CARGO_BUILD_JOBS}
COPY src/components/backend-core/Cargo.toml ./Cargo.toml
COPY src/components/backend-core/Cargo.lock ./Cargo.lock
COPY src/components/backend-core/src ./src
RUN CARGO_BUILD_JOBS=1 cargo build --release
RUN cargo build --release --jobs "${CARGO_BUILD_JOBS}"
FROM postgres:16-bookworm
RUN apt-get update \
src/components/backend-core/src/config.rs
+1
View File
@@ -135,6 +135,7 @@ pub fn read_config() -> anyhow::Result<RuntimeConfig> {
Ok(RuntimeConfig {
port: number_env("PORT")?,
provider_port: number_env("PROVIDER_PORT")?,
provider_data_port: optional_number_env("PROVIDER_DATA_PORT", 8082)?,
database_url: database_url.clone(),
identity: runtime_identity(&database_url),
provider_token: required_env("PROVIDER_TOKEN")?,
src/components/backend-core/src/main.rs
+8 -1
View File
@@ -10,6 +10,7 @@ mod overview;
mod performance;
mod provider_registry;
mod scheduler;
mod ssh_data_channel;
mod ssh_bridge;
mod state;
mod task_dispatcher;
@@ -36,6 +37,7 @@ use crate::performance::record_request_performance;
use crate::provider_registry::{mark_stale_providers_offline, provider_ws};
use crate::scheduler::{recover_scheduled_runs, run_due_scheduled_tasks};
use crate::ssh_bridge::ssh_ws;
use crate::ssh_data_channel::spawn_ssh_data_listener;
use crate::state::AppState;
use crate::task_dispatcher::mark_stale_tasks_failed;
@@ -71,6 +73,7 @@ async fn main() -> anyhow::Result<()> {
let api_addr = SocketAddr::from(([0, 0, 0, 0], state.config.port));
let provider_addr = SocketAddr::from(([0, 0, 0, 0], state.config.provider_port));
let provider_data_addr = SocketAddr::from(([0, 0, 0, 0], state.config.provider_data_port));
let api = Router::new()
.route("/ws/ssh", get(ssh_ws_handler))
.fallback(api_handler)
@@ -84,19 +87,22 @@ async fn main() -> anyhow::Result<()> {
let api_listener = TcpListener::bind(api_addr).await?;
let provider_listener = TcpListener::bind(provider_addr).await?;
let provider_data_listener = TcpListener::bind(provider_data_addr).await?;
state.log(
"info",
"server_listening",
Some(serde_json::json!({
"apiUrl": format!("http://0.0.0.0:{}", state.config.port),
"providerIngressUrl": format!("ws://0.0.0.0:{}/ws/provider", state.config.provider_port),
"providerDataTcpUrl": format!("tcp://0.0.0.0:{}", state.config.provider_data_port),
"logFile": state.config.log_file,
})),
);
let api_server = axum::serve(api_listener, api);
let provider_server = axum::serve(provider_listener, provider);
tokio::try_join!(api_server, provider_server)?;
let provider_data_server = spawn_ssh_data_listener(state.clone(), provider_data_listener);
tokio::try_join!(api_server, provider_server, provider_data_server)?;
Ok(())
}
@@ -185,6 +191,7 @@ async fn provider_health_handler(State(state): State<Arc<AppState>>) -> Response
"ok": true,
"service": "unidesk-provider-ingress",
"activeSocketCount": state.active_provider_count().await,
"providerDataPort": state.config.provider_data_port,
}),
200,
)
src/components/backend-core/src/provider-registry.ts
+1 -6
View File
@@ -72,12 +72,7 @@ export async function handleProviderMessage(ws: ProviderSocket, raw: string | Bu
ws.data.providerId = message.providerId;
ctx.activeProviders.set(message.providerId, ws);
if (
message.type === "host_ssh_opened" ||
message.type === "host_ssh_data" ||
message.type === "host_ssh_exit" ||
message.type === "host_ssh_error"
) {
if (message.type === "host_ssh_error") {
forwardSshProviderMessage(message);
return;
}
src/components/backend-core/src/provider_registry.rs
+6 -1
View File
@@ -657,7 +657,7 @@ async fn handle_provider_text(
.insert(provider_id.clone(), connection.clone());
let message_type = message.get("type").and_then(Value::as_str).unwrap_or("");
match message_type {
"host_ssh_opened" | "host_ssh_data" | "host_ssh_exit" | "host_ssh_error" => {
"host_ssh_error" => {
forward_ssh_provider_message(state, &message).await;
}
"http_tunnel_response" => {
@@ -820,6 +820,11 @@ async fn handle_provider_text(
pub async fn mark_provider_offline(state: &Arc<AppState>, provider_id: &str) -> anyhow::Result<()> {
state.active_providers.write().await.remove(provider_id);
state
.active_ssh_data_channels
.lock()
.await
.retain(|_, channel| channel.provider_id != provider_id);
if !state.db_ready() {
return Ok(());
}
src/components/backend-core/src/ssh-bridge.ts
+16 -84
View File
@@ -1,14 +1,6 @@
import type { Server } from "bun";
import type {
CoreHostSshCloseMessage,
CoreHostSshEofMessage,
CoreHostSshInputMessage,
CoreHostSshOpenMessage,
CoreHostSshResizeMessage,
ProviderHostSshDataMessage,
ProviderHostSshErrorMessage,
ProviderHostSshExitMessage,
ProviderHostSshOpenedMessage,
} from "../../shared/src/index";
import { ctx, sql, config, logger } from "./context";
import type { ProviderSocket, WsData } from "./types";
@@ -33,26 +25,13 @@ export function closeSshClient(sessionId: string, code = 1000, reason = "ssh ses
}
export function forwardSshProviderMessage(
message: ProviderHostSshOpenedMessage | ProviderHostSshDataMessage | ProviderHostSshExitMessage | ProviderHostSshErrorMessage,
message: ProviderHostSshErrorMessage,
): void {
const client = sshClientFor(message.sessionId);
if (client === null) {
logger("warn", "ssh_client_missing", { providerId: message.providerId, sessionId: message.sessionId, type: message.type });
return;
}
if (message.type === "host_ssh_opened") {
wsSendJson(client, { type: "ssh.opened", providerId: message.providerId, sessionId: message.sessionId });
return;
}
if (message.type === "host_ssh_data") {
wsSendJson(client, { type: "ssh.data", stream: message.stream, data: message.data, encoding: message.encoding });
return;
}
if (message.type === "host_ssh_exit") {
wsSendJson(client, { type: "ssh.exit", exitCode: message.exitCode, signal: message.signal });
setTimeout(() => closeSshClient(message.sessionId), 50);
return;
}
wsSendJson(client, { type: "ssh.error", message: message.message });
setTimeout(() => closeSshClient(message.sessionId, 1011, "ssh session error"), 50);
}
@@ -95,70 +74,23 @@ export async function handleSshClientMessage(ws: ProviderSocket, raw: string | B
wsSendJson(ws, { type: "ssh.error", message: `provider does not declare host.ssh capability: ${providerId}` });
return;
}
const sessionId = safeSessionId();
ws.data.channel = "ssh";
ws.data.providerId = providerId;
ws.data.sessionId = sessionId;
ctx.activeSshClients.set(sessionId, ws);
const openMessage: CoreHostSshOpenMessage = {
type: "host_ssh_open",
sessionId,
cols: numberFromUnknown(message.cols, 100, 20, 300),
rows: numberFromUnknown(message.rows, 30, 8, 120),
};
if (typeof message.cwd === "string" && message.cwd.length > 0) openMessage.cwd = message.cwd;
if (typeof message.command === "string" && message.command.length > 0) openMessage.command = message.command;
if (typeof message.tty === "boolean") openMessage.tty = message.tty;
provider.send(JSON.stringify(openMessage));
wsSendJson(ws, { type: "ssh.dispatched", providerId, sessionId });
logger("info", "ssh_session_dispatched", { providerId, sessionId, hasCommand: typeof openMessage.command === "string" });
return;
}
const sessionId = ws.data.sessionId;
const providerId = ws.data.providerId;
if (sessionId === undefined || providerId === undefined) {
if (message.type === "ssh.eof" || message.type === "ssh.close") return;
wsSendJson(ws, { type: "ssh.error", message: "ssh session is not open" });
return;
}
const provider = providerForSsh(providerId);
if (provider === null) {
wsSendJson(ws, { type: "ssh.error", message: `provider went offline: ${providerId}` });
return;
}
if (message.type === "ssh.input") {
if (typeof message.data !== "string" || message.encoding !== "base64") {
wsSendJson(ws, { type: "ssh.error", message: "ssh.input requires base64 data" });
if (!(await providerSupports(providerId, "host.ssh.tcp-pool"))) {
wsSendJson(ws, {
type: "ssh.error",
message: `provider gateway must be upgraded before host.ssh can run: ${providerId} does not declare host.ssh.tcp-pool`,
failureKind: "provider-gateway-upgrade-required",
providerId,
requiredCapability: "host.ssh.tcp-pool",
});
return;
}
const inputMessage: CoreHostSshInputMessage = { type: "host_ssh_input", sessionId, data: message.data, encoding: "base64" };
provider.send(JSON.stringify(inputMessage));
return;
}
if (message.type === "ssh.resize") {
const resizeMessage: CoreHostSshResizeMessage = {
type: "host_ssh_resize",
sessionId,
cols: numberFromUnknown(message.cols, 100, 20, 300),
rows: numberFromUnknown(message.rows, 30, 8, 120),
};
provider.send(JSON.stringify(resizeMessage));
return;
}
if (message.type === "ssh.eof") {
const eofMessage: CoreHostSshEofMessage = { type: "host_ssh_eof", sessionId };
provider.send(JSON.stringify(eofMessage));
return;
}
if (message.type === "ssh.close") {
const closeMessage: CoreHostSshCloseMessage = { type: "host_ssh_close", sessionId };
provider.send(JSON.stringify(closeMessage));
closeSshClient(sessionId);
wsSendJson(ws, {
type: "ssh.error",
message: "legacy TypeScript backend-core does not implement host.ssh.tcp-pool data channels; use the Rust backend-core",
failureKind: "backend-core-upgrade-required",
providerId,
requiredCapability: "host.ssh.tcp-pool",
});
return;
}
src/components/backend-core/src/ssh_bridge.rs
+101 -49
View File
@@ -3,6 +3,7 @@ use std::sync::Arc;
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::http::{HeaderMap, Uri};
use axum::response::Response;
use base64::Engine;
use futures_util::{SinkExt, StreamExt};
use rand::Rng;
use serde_json::{json, Value};
@@ -10,7 +11,11 @@ use tokio::sync::mpsc;
use crate::db::provider_supports;
use crate::http::json_response;
use crate::state::{AppState, SshClientConnection};
use crate::ssh_data_channel::{
active_ssh_data_channel_count, claim_ssh_data_channel, close_ssh_data_session,
idle_ssh_data_channel_count, send_ssh_data_frame,
};
use crate::state::{AppState, SshClientConnection, SshDataFrame};
pub async fn ssh_ws(
state: Arc<AppState>,
@@ -65,7 +70,7 @@ async fn ssh_socket_task(state: Arc<AppState>, socket: WebSocket) {
}
}
if let Some(session_id) = session_id {
state.active_ssh_clients.write().await.remove(&session_id);
close_ssh_client(&state, &session_id, Some("ssh client disconnected")).await;
}
send_task.abort();
}
@@ -104,18 +109,53 @@ async fn handle_ssh_client_text(
let _ = client_sender.send(Message::Text(json!({ "type": "ssh.error", "message": format!("provider does not declare host.ssh capability: {provider_id}") }).to_string()));
return Ok(());
}
if !provider_supports(state, provider_id, "host.ssh.tcp-pool")
.await
.unwrap_or(false)
{
let _ = client_sender.send(Message::Text(json!({
"type": "ssh.error",
"message": format!("provider gateway must be upgraded before host.ssh can run: {provider_id} does not declare host.ssh.tcp-pool"),
"failureKind": "provider-gateway-upgrade-required",
"providerId": provider_id,
"requiredCapability": "host.ssh.tcp-pool",
}).to_string()));
return Ok(());
}
let new_session_id = safe_session_id();
let Some((data_channel_id, _data_writer)) =
claim_ssh_data_channel(state, provider_id, &new_session_id).await
else {
let ready = idle_ssh_data_channel_count(state, provider_id).await;
let total = active_ssh_data_channel_count(state, provider_id).await;
let _ = client_sender.send(Message::Text(json!({
"type": "ssh.error",
"message": format!("provider ssh tcp data pool has no idle channel: {provider_id}"),
"failureKind": "provider-data-pool-exhausted",
"providerId": provider_id,
"dataPool": {
"ready": ready,
"total": total,
"required": 1
},
"requiredCapability": "host.ssh.tcp-pool",
}).to_string()));
return Ok(());
};
*session_id = Some(new_session_id.clone());
state.active_ssh_clients.write().await.insert(
new_session_id.clone(),
SshClientConnection {
provider_id: provider_id.to_string(),
sender: client_sender.clone(),
data_channel_id: data_channel_id.clone(),
},
);
let mut open_message = json!({
"type": "host_ssh_open",
"sessionId": new_session_id,
"transport": "tcp-pool",
"dataChannelId": data_channel_id,
"cols": number_from_unknown(message.get("cols"), 100, 20, 300),
"rows": number_from_unknown(message.get("rows"), 30, 8, 120),
});
@@ -139,8 +179,8 @@ async fn handle_ssh_client_text(
let _ = provider
.sender
.send(Message::Text(open_message.to_string()));
let _ = client_sender.send(Message::Text(json!({ "type": "ssh.dispatched", "providerId": provider_id, "sessionId": new_session_id }).to_string()));
state.log("info", "ssh_session_dispatched", Some(json!({ "providerId": provider_id, "sessionId": session_id, "hasCommand": open_message.get("command").is_some() })));
let _ = client_sender.send(Message::Text(json!({ "type": "ssh.dispatched", "providerId": provider_id, "sessionId": new_session_id, "transport": "tcp-pool", "dataChannelId": open_message.get("dataChannelId").cloned().unwrap_or(Value::Null) }).to_string()));
state.log("info", "ssh_session_dispatched", Some(json!({ "providerId": provider_id, "sessionId": new_session_id, "transport": "tcp-pool", "dataChannelId": open_message.get("dataChannelId").cloned().unwrap_or(Value::Null), "hasCommand": open_message.get("command").is_some() })));
return Ok(());
}
@@ -153,28 +193,15 @@ async fn handle_ssh_client_text(
));
return Ok(());
};
let provider_id = {
let session = {
let clients = state.active_ssh_clients.read().await;
clients
.get(&current_session_id)
.map(|client| client.provider_id.clone())
.map(|client| (client.provider_id.clone(), client.data_channel_id.clone()))
};
let provider_id = provider_id.unwrap_or_else(|| {
message
.get("providerId")
.and_then(Value::as_str)
.unwrap_or("")
.to_string()
});
let provider = state
.active_providers
.read()
.await
.get(&provider_id)
.cloned();
let Some(provider) = provider else {
let Some((provider_id, data_channel_id)) = session else {
let _ = client_sender.send(Message::Text(
json!({ "type": "ssh.error", "message": "provider went offline" }).to_string(),
json!({ "type": "ssh.error", "message": "ssh session disappeared" }).to_string(),
));
return Ok(());
};
@@ -190,28 +217,60 @@ async fn handle_ssh_client_text(
));
return Ok(());
}
let _ = provider.sender.send(Message::Text(json!({ "type": "host_ssh_input", "sessionId": current_session_id, "data": data, "encoding": "base64" }).to_string()));
let bytes = match base64::engine::general_purpose::STANDARD.decode(data) {
Ok(bytes) => bytes,
Err(error) => {
let _ = client_sender.send(Message::Text(
json!({ "type": "ssh.error", "message": format!("invalid ssh.input base64: {error}") })
.to_string(),
));
return Ok(());
}
};
if !send_ssh_data_frame(
state,
&provider_id,
&data_channel_id,
SshDataFrame {
header: json!({ "type": "input", "sessionId": current_session_id }),
payload: bytes,
},
)
.await
{
let _ = client_sender.send(Message::Text(json!({ "type": "ssh.error", "message": "ssh tcp data channel is not available", "failureKind": "provider-data-channel-missing", "providerId": provider_id, "dataChannelId": data_channel_id }).to_string()));
}
}
"ssh.resize" => {
let _ = provider.sender.send(Message::Text(
json!({
"type": "host_ssh_resize",
let _ = send_ssh_data_frame(
state,
&provider_id,
&data_channel_id,
SshDataFrame {
header: json!({
"type": "resize",
"sessionId": current_session_id,
"cols": number_from_unknown(message.get("cols"), 100, 20, 300),
"rows": number_from_unknown(message.get("rows"), 30, 8, 120),
})
.to_string(),
));
}),
payload: Vec::new(),
},
)
.await;
}
"ssh.eof" => {
let _ = provider.sender.send(Message::Text(
json!({ "type": "host_ssh_eof", "sessionId": current_session_id }).to_string(),
));
let _ = send_ssh_data_frame(
state,
&provider_id,
&data_channel_id,
SshDataFrame {
header: json!({ "type": "eof", "sessionId": current_session_id }),
payload: Vec::new(),
},
)
.await;
}
"ssh.close" => {
let _ = provider.sender.send(Message::Text(
json!({ "type": "host_ssh_close", "sessionId": current_session_id }).to_string(),
));
close_ssh_client(state, &current_session_id, None).await;
}
_ => {
@@ -232,23 +291,15 @@ pub async fn forward_ssh_provider_message(state: &Arc<AppState>, message: &Value
return;
};
let message_type = message.get("type").and_then(Value::as_str).unwrap_or("");
let outbound = match message_type {
"host_ssh_opened" => {
json!({ "type": "ssh.opened", "providerId": message.get("providerId").cloned().unwrap_or(Value::Null), "sessionId": session_id })
}
"host_ssh_data" => {
json!({ "type": "ssh.data", "stream": message.get("stream").cloned().unwrap_or(Value::Null), "data": message.get("data").cloned().unwrap_or(Value::Null), "encoding": message.get("encoding").cloned().unwrap_or(Value::Null) })
}
"host_ssh_exit" => {
json!({ "type": "ssh.exit", "exitCode": message.get("exitCode").cloned().unwrap_or(Value::Null), "signal": message.get("signal").cloned().unwrap_or(Value::Null) })
}
_ => {
json!({ "type": "ssh.error", "message": message.get("message").and_then(Value::as_str).unwrap_or("ssh session error") })
}
};
let outbound = json!({
"type": "ssh.error",
"message": message.get("message").and_then(Value::as_str).unwrap_or("ssh session error"),
"transport": "tcp-pool",
"controlFallback": true,
});
let _ = client.sender.send(Message::Text(outbound.to_string()));
drop(clients);
if message_type == "host_ssh_exit" || message_type == "host_ssh_error" {
if message_type == "host_ssh_error" {
let state = state.clone();
let session_id = session_id.to_string();
tokio::spawn(async move {
@@ -259,6 +310,7 @@ pub async fn forward_ssh_provider_message(state: &Arc<AppState>, message: &Value
}
async fn close_ssh_client(state: &Arc<AppState>, session_id: &str, message: Option<&str>) {
close_ssh_data_session(state, session_id).await;
let client = state.active_ssh_clients.write().await.remove(session_id);
if let Some(client) = client {
let close = Message::Close(Some(axum::extract::ws::CloseFrame {
src/components/backend-core/src/ssh_data_channel.rs
+358
View File
@@ -0,0 +1,358 @@
use std::io::{Error, ErrorKind};
use std::sync::Arc;
use axum::extract::ws::Message;
use base64::Engine;
use serde_json::{json, Value};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{tcp::OwnedReadHalf, tcp::OwnedWriteHalf, TcpListener, TcpStream};
use tokio::sync::mpsc;
use crate::state::{AppState, SshDataChannel, SshDataFrame};
const SSH_DATA_PROTOCOL: &str = "unidesk-host-ssh-tcp-pool-v1";
const SSH_DATA_MAX_HEADER_BYTES: usize = 64 * 1024;
const SSH_DATA_MAX_PAYLOAD_BYTES: usize = 16 * 1024 * 1024;
fn channel_key(provider_id: &str, channel_id: &str) -> String {
format!("{provider_id}:{channel_id}")
}
fn header_string(header: &Value, key: &str) -> String {
header
.get(key)
.and_then(Value::as_str)
.unwrap_or("")
.to_string()
}
async fn read_data_frame(reader: &mut OwnedReadHalf) -> std::io::Result<Option<(Value, Vec<u8>)>> {
let mut header_len_bytes = [0_u8; 4];
if let Err(error) = reader.read_exact(&mut header_len_bytes).await {
if error.kind() == ErrorKind::UnexpectedEof {
return Ok(None);
}
return Err(error);
}
let header_len = u32::from_be_bytes(header_len_bytes) as usize;
if header_len == 0 || header_len > SSH_DATA_MAX_HEADER_BYTES {
return Err(Error::new(
ErrorKind::InvalidData,
format!("invalid ssh data header length: {header_len}"),
));
}
let mut header_bytes = vec![0_u8; header_len];
reader.read_exact(&mut header_bytes).await?;
let header: Value = serde_json::from_slice(&header_bytes)
.map_err(|error| Error::new(ErrorKind::InvalidData, error.to_string()))?;
let payload_len = header
.get("length")
.and_then(Value::as_u64)
.unwrap_or(0) as usize;
if payload_len > SSH_DATA_MAX_PAYLOAD_BYTES {
return Err(Error::new(
ErrorKind::InvalidData,
format!("ssh data payload too large: {payload_len}"),
));
}
let mut payload = vec![0_u8; payload_len];
if payload_len > 0 {
reader.read_exact(&mut payload).await?;
}
Ok(Some((header, payload)))
}
async fn write_data_frame(
writer: &mut OwnedWriteHalf,
frame: SshDataFrame,
) -> std::io::Result<()> {
let mut header = frame.header;
if let Some(object) = header.as_object_mut() {
object.insert(
"length".to_string(),
Value::Number(serde_json::Number::from(frame.payload.len())),
);
}
let header_bytes = serde_json::to_vec(&header)
.map_err(|error| Error::new(ErrorKind::InvalidData, error.to_string()))?;
if header_bytes.len() > SSH_DATA_MAX_HEADER_BYTES {
return Err(Error::new(
ErrorKind::InvalidData,
format!("ssh data header too large: {}", header_bytes.len()),
));
}
writer
.write_all(&(header_bytes.len() as u32).to_be_bytes())
.await?;
writer.write_all(&header_bytes).await?;
if !frame.payload.is_empty() {
writer.write_all(&frame.payload).await?;
}
writer.flush().await?;
Ok(())
}
pub async fn claim_ssh_data_channel(
state: &Arc<AppState>,
provider_id: &str,
session_id: &str,
) -> Option<(String, mpsc::UnboundedSender<SshDataFrame>)> {
let mut channels = state.active_ssh_data_channels.lock().await;
let channel = channels
.values_mut()
.filter(|channel| channel.provider_id == provider_id && channel.active_session_id.is_none())
.min_by(|left, right| left.channel_id.cmp(&right.channel_id))?;
channel.active_session_id = Some(session_id.to_string());
Some((channel.channel_id.clone(), channel.writer.clone()))
}
pub async fn send_ssh_data_frame(
state: &Arc<AppState>,
provider_id: &str,
channel_id: &str,
frame: SshDataFrame,
) -> bool {
let channels = state.active_ssh_data_channels.lock().await;
channels
.get(&channel_key(provider_id, channel_id))
.is_some_and(|channel| channel.writer.send(frame).is_ok())
}
pub async fn release_ssh_data_channel(
state: &Arc<AppState>,
provider_id: &str,
channel_id: &str,
) {
let mut channels = state.active_ssh_data_channels.lock().await;
if let Some(channel) = channels.get_mut(&channel_key(provider_id, channel_id)) {
channel.active_session_id = None;
}
}
pub async fn close_ssh_data_session(state: &Arc<AppState>, session_id: &str) {
let client = {
let clients = state.active_ssh_clients.read().await;
clients.get(session_id).map(|client| {
(
client.provider_id.clone(),
client.data_channel_id.clone(),
)
})
};
let Some((provider_id, channel_id)) = client else {
return;
};
let _ = send_ssh_data_frame(
state,
&provider_id,
&channel_id,
SshDataFrame {
header: json!({ "type": "close", "sessionId": session_id }),
payload: Vec::new(),
},
)
.await;
release_ssh_data_channel(state, &provider_id, &channel_id).await;
}
pub async fn active_ssh_data_channel_count(state: &Arc<AppState>, provider_id: &str) -> usize {
let channels = state.active_ssh_data_channels.lock().await;
channels
.values()
.filter(|channel| channel.provider_id == provider_id)
.count()
}
pub async fn idle_ssh_data_channel_count(state: &Arc<AppState>, provider_id: &str) -> usize {
let channels = state.active_ssh_data_channels.lock().await;
channels
.values()
.filter(|channel| channel.provider_id == provider_id && channel.active_session_id.is_none())
.count()
}
async fn forward_provider_frame(
state: &Arc<AppState>,
provider_id: &str,
channel_id: &str,
header: Value,
payload: Vec<u8>,
) {
let session_id = header_string(&header, "sessionId");
if session_id.is_empty() {
state.log(
"warn",
"ssh_data_frame_missing_session",
Some(json!({ "providerId": provider_id, "channelId": channel_id, "header": header })),
);
return;
}
let clients = state.active_ssh_clients.read().await;
let Some(client) = clients.get(&session_id) else {
state.log(
"warn",
"ssh_data_client_missing",
Some(json!({ "providerId": provider_id, "channelId": channel_id, "sessionId": session_id, "header": header })),
);
return;
};
let frame_type = header_string(&header, "type");
let client_sender = client.sender.clone();
let outbound = match frame_type.as_str() {
"opened" => json!({
"type": "ssh.opened",
"providerId": provider_id,
"sessionId": session_id,
"transport": "tcp-pool",
"dataChannelId": channel_id,
}),
"data" => json!({
"type": "ssh.data",
"stream": header.get("stream").cloned().unwrap_or_else(|| Value::String("stdout".to_string())),
"data": base64::engine::general_purpose::STANDARD.encode(payload),
"encoding": "base64",
"transport": "tcp-pool",
"dataChannelId": channel_id,
}),
"exit" => json!({
"type": "ssh.exit",
"exitCode": header.get("exitCode").cloned().unwrap_or(Value::Null),
"signal": header.get("signal").cloned().unwrap_or(Value::Null),
"transport": "tcp-pool",
"dataChannelId": channel_id,
}),
"error" => json!({
"type": "ssh.error",
"message": header.get("message").and_then(Value::as_str).unwrap_or("ssh data channel error"),
"transport": "tcp-pool",
"dataChannelId": channel_id,
}),
_ => json!({
"type": "ssh.error",
"message": format!("unsupported ssh data frame: {frame_type}"),
"transport": "tcp-pool",
"dataChannelId": channel_id,
}),
};
let _ = client_sender.send(Message::Text(outbound.to_string()));
drop(clients);
if frame_type == "exit" || frame_type == "error" {
let close = Message::Close(Some(axum::extract::ws::CloseFrame {
code: axum::extract::ws::close_code::NORMAL,
reason: "ssh session closed".to_string().into(),
}));
let _ = client_sender.send(close);
release_ssh_data_channel(state, provider_id, channel_id).await;
let state = state.clone();
tokio::spawn(async move {
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
state.active_ssh_clients.write().await.remove(&session_id);
});
}
}
async fn handle_ssh_data_stream(state: Arc<AppState>, stream: TcpStream) {
let peer_addr = stream.peer_addr().ok();
let (mut reader, mut writer) = stream.into_split();
let hello = match read_data_frame(&mut reader).await {
Ok(Some((header, _))) => header,
Ok(None) => return,
Err(error) => {
state.log("warn", "ssh_data_hello_read_failed", Some(json!({ "peer": peer_addr.map(|addr| addr.to_string()), "error": error.to_string() })));
return;
}
};
let provider_id = header_string(&hello, "providerId");
let channel_id = header_string(&hello, "channelId");
let token = header_string(&hello, "token");
let protocol = header_string(&hello, "protocol");
if provider_id.is_empty()
|| channel_id.is_empty()
|| token != state.config.provider_token
|| protocol != SSH_DATA_PROTOCOL
{
state.log(
"warn",
"ssh_data_hello_rejected",
Some(json!({ "providerId": provider_id, "channelId": channel_id, "protocol": protocol, "peer": peer_addr.map(|addr| addr.to_string()) })),
);
return;
}
if !state.active_providers.read().await.contains_key(&provider_id) {
state.log(
"warn",
"ssh_data_provider_not_registered",
Some(json!({ "providerId": provider_id, "channelId": channel_id })),
);
return;
}
let (tx, mut rx) = mpsc::unbounded_channel::<SshDataFrame>();
state.active_ssh_data_channels.lock().await.insert(
channel_key(&provider_id, &channel_id),
SshDataChannel {
provider_id: provider_id.clone(),
channel_id: channel_id.clone(),
writer: tx,
active_session_id: None,
},
);
state.log(
"info",
"ssh_data_channel_ready",
Some(json!({ "providerId": provider_id, "channelId": channel_id, "peer": peer_addr.map(|addr| addr.to_string()) })),
);
let writer_state = state.clone();
let writer_provider_id = provider_id.clone();
let writer_channel_id = channel_id.clone();
let writer_task = tokio::spawn(async move {
while let Some(frame) = rx.recv().await {
if let Err(error) = write_data_frame(&mut writer, frame).await {
writer_state.log(
"warn",
"ssh_data_write_failed",
Some(json!({ "providerId": writer_provider_id, "channelId": writer_channel_id, "error": error.to_string() })),
);
break;
}
}
});
loop {
match read_data_frame(&mut reader).await {
Ok(Some((header, payload))) => {
forward_provider_frame(&state, &provider_id, &channel_id, header, payload).await;
}
Ok(None) => break,
Err(error) => {
state.log(
"warn",
"ssh_data_read_failed",
Some(json!({ "providerId": provider_id, "channelId": channel_id, "error": error.to_string() })),
);
break;
}
}
}
state
.active_ssh_data_channels
.lock()
.await
.remove(&channel_key(&provider_id, &channel_id));
writer_task.abort();
state.log(
"warn",
"ssh_data_channel_closed",
Some(json!({ "providerId": provider_id, "channelId": channel_id })),
);
}
pub async fn spawn_ssh_data_listener(
state: Arc<AppState>,
listener: TcpListener,
) -> std::io::Result<()> {
loop {
let (stream, _) = listener.accept().await?;
tokio::spawn(handle_ssh_data_stream(state.clone(), stream));
}
}
src/components/backend-core/src/state.rs
+15
View File
@@ -23,6 +23,19 @@ pub struct ProviderConnection {
pub struct SshClientConnection {
pub provider_id: String,
pub sender: mpsc::UnboundedSender<Message>,
pub data_channel_id: String,
}
pub struct SshDataFrame {
pub header: JsonValue,
pub payload: Vec<u8>,
}
pub struct SshDataChannel {
pub provider_id: String,
pub channel_id: String,
pub writer: mpsc::UnboundedSender<SshDataFrame>,
pub active_session_id: Option<String>,
}
pub struct HttpTunnelResponse {
@@ -43,6 +56,7 @@ pub struct AppState {
pub service_started_at: chrono::DateTime<chrono::Utc>,
pub active_providers: RwLock<HashMap<String, Arc<ProviderConnection>>>,
pub active_ssh_clients: RwLock<HashMap<String, SshClientConnection>>,
pub active_ssh_data_channels: Mutex<HashMap<String, SshDataChannel>>,
pub http_tunnel_waiters: Mutex<HashMap<String, HttpTunnelWaiter>>,
pub task_terminal_waiters: Mutex<HashMap<String, Vec<TaskTerminalWaiter>>>,
pub active_egress_tcp: Mutex<HashMap<String, crate::egress_tcp::EgressTcpConnection>>,
@@ -77,6 +91,7 @@ impl AppState {
service_started_at: chrono::Utc::now(),
active_providers: RwLock::new(HashMap::new()),
active_ssh_clients: RwLock::new(HashMap::new()),
active_ssh_data_channels: Mutex::new(HashMap::new()),
http_tunnel_waiters: Mutex::new(HashMap::new()),
task_terminal_waiters: Mutex::new(HashMap::new()),
active_egress_tcp: Mutex::new(HashMap::new()),
src/components/backend-core/src/types.rs
+1
View File
@@ -9,6 +9,7 @@ pub type JsonValue = Value;
pub struct RuntimeConfig {
pub port: u16,
pub provider_port: u16,
pub provider_data_port: u16,
pub database_url: String,
pub identity: RuntimeIdentity,
pub provider_token: String,