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,
src/components/provider-gateway/package.json
+2 -2
View File
@@ -1,10 +1,10 @@
{
"name": "@unidesk/provider-gateway",
"version": "0.2.27",
"version": "0.2.28",
"private": true,
"type": "module",
"scripts": {
"start": "bun run src/index.ts",
"check": "tsc -p tsconfig.json --noEmit"
"check": "tsc -b ../shared/tsconfig.json && tsc -p tsconfig.json --noEmit"
}
}
src/components/provider-gateway/src/index.ts
+329 -77
View File
@@ -1,16 +1,13 @@
import { spawnSync } from "node:child_process";
import { existsSync, readFileSync, readdirSync } from "node:fs";
import { createConnection, type Socket } from "node:net";
import {
type CoreEgressTcpCloseMessage,
type CoreEgressTcpDataMessage,
type CoreEgressTcpOpenedMessage,
type CoreDispatchMessage,
type CoreHttpTunnelRequestMessage,
type CoreHostSshCloseMessage,
type CoreHostSshEofMessage,
type CoreHostSshInputMessage,
type CoreHostSshOpenMessage,
type CoreHostSshResizeMessage,
type DockerContainerSummary,
type DockerImageSummary,
type DockerNetworkSummary,
@@ -50,6 +47,10 @@ interface RuntimeConfig {
hostSshKey: string | null;
hostRemoteCwd: string | null;
hostLoginShell: string | null;
providerDataHost: string;
providerDataPort: number;
providerDataPoolSize: number;
providerDataConnectTimeoutMs: number;
egressProxyEnabled: boolean;
egressProxyListenHost: string;
egressProxyPort: number;
@@ -87,6 +88,20 @@ interface HostSshSession {
proc: ReturnType<typeof Bun.spawn>;
openedAt: number;
tty: boolean;
dataChannelId?: string;
}
interface SshDataChannel {
id: string;
socket: Socket;
status: "connecting" | "ready" | "claimed" | "closed";
buffer: Buffer;
openedAt: number;
lastReadyAt: number | null;
activeSessionId: string | null;
bytesIn: number;
bytesOut: number;
lastError: string | null;
}
interface HostSshStdin {
@@ -127,6 +142,13 @@ const gatewayMetadata = readGatewayMetadata();
const defaultMasterServer = "http://74.48.78.17/";
const defaultProviderToken = "unidesk-dev-token-change-me";
const microserviceHttpMaxBodyTextLength = 8 * 1024 * 1024;
const sshDataProtocol = "unidesk-host-ssh-tcp-pool-v1";
const sshDataMaxHeaderBytes = 64 * 1024;
const sshDataMaxPayloadBytes = 16 * 1024 * 1024;
const sshDataChannels = new Map<string, SshDataChannel>();
let sshDataPoolDesired = false;
let sshDataChannelSeq = 0;
let sshDataLastError: string | null = null;
const microserviceForwardRequestHeaders = [
"accept",
"content-type",
@@ -247,6 +269,18 @@ function readProviderServerUrl(): string {
return providerServerUrlFromMaster(readEnv("UNIDESK_MASTER_SERVER", "UNIDESK_MASTER_SERVER_URL", "MASTER_SERVER_URL", "MASTER_SERVER_IP") ?? defaultMasterServer);
}
function deriveProviderDataEndpoint(serverUrl: string): { host: string; port: number } {
const explicitHost = readOptionalStringEnv("PROVIDER_DATA_HOST", "UNIDESK_PROVIDER_DATA_HOST");
const explicitPort = readOptionalNumberEnv("PROVIDER_DATA_PORT", "UNIDESK_PROVIDER_DATA_PORT");
const url = new URL(serverUrl);
const host = explicitHost ?? url.hostname;
if (explicitPort !== null) return { host, port: explicitPort };
const controlPort = Number(url.port || (url.protocol === "wss:" ? 443 : 80));
if (controlPort === 8081) return { host, port: 8082 };
if (controlPort === 18082) return { host, port: 18084 };
return { host, port: controlPort + 2 };
}
function readProviderId(): string {
return requiredEnv("PROVIDER_ID", "UNIDESK_PROVIDER_ID");
}
@@ -459,6 +493,8 @@ function defaultHostRemoteCwd(user: string | null): string | null {
function readConfig(): RuntimeConfig {
const providerId = readProviderId();
const serverUrl = readProviderServerUrl();
const providerDataEndpoint = deriveProviderDataEndpoint(serverUrl);
const workspacePath = readOptionalStringEnv("PROVIDER_UPGRADE_WORKSPACE_PATH") ?? "/workspace";
const hostProjectRoot = readOptionalStringEnv("PROVIDER_UPGRADE_HOST_PROJECT_ROOT")
?? discoverDockerMountSource(workspacePath)
@@ -468,7 +504,7 @@ function readConfig(): RuntimeConfig {
const mountedHostSshKey = existsSync(defaultHostSshKey);
const hostSshUser = readOptionalStringEnv("HOST_SSH_USER") ?? (mountedHostSshKey ? inferredSshUser : null);
return {
serverUrl: readProviderServerUrl(),
serverUrl,
token: readEnv("PROVIDER_TOKEN", "UNIDESK_PROVIDER_TOKEN") ?? defaultProviderToken,
providerId,
providerName: readEnv("PROVIDER_NAME", "UNIDESK_PROVIDER_NAME") ?? providerId,
@@ -491,6 +527,10 @@ function readConfig(): RuntimeConfig {
hostSshKey: readOptionalStringEnv("HOST_SSH_KEY") ?? (mountedHostSshKey ? defaultHostSshKey : null),
hostRemoteCwd: readOptionalStringEnv("HOST_REMOTE_CWD") ?? defaultHostRemoteCwd(hostSshUser),
hostLoginShell: readOptionalStringEnv("HOST_LOGIN_SHELL") ?? (hostSshUser === null ? null : "/bin/bash"),
providerDataHost: providerDataEndpoint.host,
providerDataPort: providerDataEndpoint.port,
providerDataPoolSize: Math.max(1, Math.min(64, Math.floor(readNumberEnv("PROVIDER_DATA_POOL_SIZE", 10, "UNIDESK_PROVIDER_DATA_POOL_SIZE")))),
providerDataConnectTimeoutMs: Math.max(250, Math.min(30_000, Math.floor(readNumberEnv("PROVIDER_DATA_CONNECT_TIMEOUT_MS", 5000, "UNIDESK_PROVIDER_DATA_CONNECT_TIMEOUT_MS")))),
egressProxyEnabled: readBooleanEnv("PROVIDER_EGRESS_PROXY_ENABLED", true, "UNIDESK_PROVIDER_EGRESS_PROXY_ENABLED"),
egressProxyListenHost: readOptionalStringEnv("PROVIDER_EGRESS_PROXY_LISTEN_HOST", "UNIDESK_PROVIDER_EGRESS_PROXY_LISTEN_HOST") ?? "0.0.0.0",
egressProxyPort: readNumberEnv("PROVIDER_EGRESS_PROXY_PORT", 18789, "UNIDESK_PROVIDER_EGRESS_PROXY_PORT"),
@@ -530,6 +570,7 @@ function currentLabels(): ProviderLabels {
const hostSshConfigured = isHostSshConfigured();
const containerGuard = refreshSelfContainerGuard();
const egressProxyStatus = providerEgressProxy?.status() ?? null;
const dataPoolStatus = sshDataPoolStatus();
return {
...config.labels,
dockerSocketPresent: existsSync(config.dockerSocketPath),
@@ -546,6 +587,15 @@ function currentLabels(): ProviderLabels {
providerGatewayEgressProxyPort: config.egressProxyPort,
providerGatewayEgressProxyConnected: egressProxyStatus === null ? false : egressProxyStatus.connected === true,
providerGatewayEgressProxyActiveTunnels: egressProxyStatus === null ? 0 : egressProxyStatus.activeTunnels ?? 0,
providerGatewaySshDataTransport: "tcp-pool",
providerGatewaySshDataHost: config.providerDataHost,
providerGatewaySshDataPort: config.providerDataPort,
providerGatewaySshDataPoolDesired: config.providerDataPoolSize,
providerGatewaySshDataPoolTotal: dataPoolStatus.total,
providerGatewaySshDataPoolReady: dataPoolStatus.ready,
providerGatewaySshDataPoolClaimed: dataPoolStatus.claimed,
providerGatewaySshDataPoolConnecting: dataPoolStatus.connecting,
providerGatewaySshDataPoolLastError: dataPoolStatus.lastError,
providerGatewayDockerRestartGuard: true,
providerGatewayContainerId: containerGuard.containerId,
providerGatewayContainerName: containerGuard.containerName,
@@ -572,9 +622,247 @@ function sendJsonOk(value: unknown): boolean {
return true;
}
function sshDataPoolStatus(): { total: number; ready: number; claimed: number; connecting: number; lastError: string | null } {
let ready = 0;
let claimed = 0;
let connecting = 0;
for (const channel of sshDataChannels.values()) {
if (channel.status === "ready") ready += 1;
if (channel.status === "claimed") claimed += 1;
if (channel.status === "connecting") connecting += 1;
}
return { total: sshDataChannels.size, ready, claimed, connecting, lastError: sshDataLastError };
}
function nextSshDataChannelId(): string {
sshDataChannelSeq += 1;
return `sshdata_${safeProviderSlug(config.providerId)}_${Date.now()}_${sshDataChannelSeq}`;
}
function writeSshDataFrame(channel: SshDataChannel, header: Record<string, unknown>, payload: Buffer = Buffer.alloc(0)): boolean {
if (channel.status === "closed" || channel.socket.destroyed) return false;
const fullHeader = { ...header, length: payload.length };
const headerBytes = Buffer.from(JSON.stringify(fullHeader), "utf8");
if (headerBytes.length > sshDataMaxHeaderBytes) {
channel.lastError = `ssh data header too large: ${headerBytes.length}`;
sshDataLastError = channel.lastError;
channel.socket.destroy(new Error(channel.lastError));
return false;
}
const lengthBytes = Buffer.allocUnsafe(4);
lengthBytes.writeUInt32BE(headerBytes.length, 0);
channel.socket.write(lengthBytes);
channel.socket.write(headerBytes);
if (payload.length > 0) channel.socket.write(payload);
channel.bytesOut += payload.length;
return true;
}
function sendSshDataSessionFrame(sessionId: string, header: Record<string, unknown>, payload: Buffer = Buffer.alloc(0)): boolean {
const session = hostSshSessions.get(sessionId);
if (session?.dataChannelId === undefined) return false;
const channel = sshDataChannels.get(session.dataChannelId);
if (channel === undefined) return false;
return writeSshDataFrame(channel, { ...header, sessionId }, payload);
}
function handleSshDataFrame(channel: SshDataChannel, header: Record<string, unknown>, payload: Buffer): void {
const type = typeof header.type === "string" ? header.type : "";
const sessionId = typeof header.sessionId === "string" ? header.sessionId : "";
if (sessionId.length === 0) {
logger("warn", "ssh_data_frame_missing_session", { channelId: channel.id, type });
return;
}
const session = hostSshSessions.get(sessionId);
if (session === undefined) {
logger("warn", "ssh_data_session_missing", { channelId: channel.id, sessionId, type });
return;
}
try {
if (type === "input") {
const stdin = session.proc.stdin as HostSshStdin | undefined;
if (stdin === undefined) throw new Error("host SSH stdin is not available");
if (!session.tty && payload.length === 1 && payload[0] === 4) return;
stdin.write(payload);
return;
}
if (type === "eof") {
const stdin = session.proc.stdin as HostSshStdin | undefined;
if (stdin !== undefined) stdin.end();
return;
}
if (type === "resize") {
logger("debug", "host_ssh_resize_requested", {
sessionId,
cols: typeof header.cols === "number" || typeof header.cols === "string" ? header.cols : null,
rows: typeof header.rows === "number" || typeof header.rows === "string" ? header.rows : null,
transport: "tcp-pool",
});
return;
}
if (type === "close") {
closeHostSshSession(sessionId);
return;
}
logger("warn", "ssh_data_frame_unsupported", { channelId: channel.id, sessionId, type });
} catch (error) {
sendSshDataSessionFrame(sessionId, { type: "error", message: error instanceof Error ? error.message : String(error) });
}
}
function consumeSshDataBytes(channel: SshDataChannel, chunk: Buffer): void {
channel.buffer = Buffer.concat([channel.buffer, chunk]);
while (channel.buffer.length >= 4) {
const headerLength = channel.buffer.readUInt32BE(0);
if (headerLength <= 0 || headerLength > sshDataMaxHeaderBytes) {
channel.lastError = `invalid ssh data header length: ${headerLength}`;
sshDataLastError = channel.lastError;
channel.socket.destroy(new Error(channel.lastError));
return;
}
if (channel.buffer.length < 4 + headerLength) return;
const headerBytes = channel.buffer.subarray(4, 4 + headerLength);
let header: Record<string, unknown>;
try {
header = JSON.parse(headerBytes.toString("utf8")) as Record<string, unknown>;
} catch (error) {
channel.lastError = `invalid ssh data header json: ${error instanceof Error ? error.message : String(error)}`;
sshDataLastError = channel.lastError;
channel.socket.destroy(new Error(channel.lastError));
return;
}
const payloadLength = typeof header.length === "number" && Number.isFinite(header.length)
? Math.max(0, Math.floor(header.length))
: 0;
if (payloadLength > sshDataMaxPayloadBytes) {
channel.lastError = `ssh data payload too large: ${payloadLength}`;
sshDataLastError = channel.lastError;
channel.socket.destroy(new Error(channel.lastError));
return;
}
if (channel.buffer.length < 4 + headerLength + payloadLength) return;
const payload = channel.buffer.subarray(4 + headerLength, 4 + headerLength + payloadLength);
channel.buffer = channel.buffer.subarray(4 + headerLength + payloadLength);
channel.bytesIn += payload.length;
handleSshDataFrame(channel, header, Buffer.from(payload));
}
}
function closeSshDataChannel(channel: SshDataChannel, reason: string): void {
channel.status = "closed";
sshDataChannels.delete(channel.id);
if (channel.activeSessionId !== null) {
const sessionId = channel.activeSessionId;
const session = hostSshSessions.get(sessionId);
if (session !== undefined) {
sendHostSshError(sessionId, `ssh tcp data channel closed: ${reason}`);
try {
session.proc.kill("SIGTERM");
} catch {
/* process may already be gone */
}
}
}
if (sshDataPoolDesired) {
setTimeout(ensureSshDataPool, 250);
}
}
function openSshDataChannel(): void {
const channelId = nextSshDataChannelId();
const socket = createConnection({ host: config.providerDataHost, port: config.providerDataPort });
const channel: SshDataChannel = {
id: channelId,
socket,
status: "connecting",
buffer: Buffer.alloc(0),
openedAt: Date.now(),
lastReadyAt: null,
activeSessionId: null,
bytesIn: 0,
bytesOut: 0,
lastError: null,
};
sshDataChannels.set(channelId, channel);
const timeout = setTimeout(() => {
if (channel.status === "connecting") {
channel.lastError = `provider data tcp connect timeout after ${config.providerDataConnectTimeoutMs}ms`;
sshDataLastError = channel.lastError;
socket.destroy(new Error(channel.lastError));
}
}, config.providerDataConnectTimeoutMs);
socket.on("connect", () => {
clearTimeout(timeout);
channel.status = "ready";
channel.lastReadyAt = Date.now();
sshDataLastError = null;
writeSshDataFrame(channel, {
type: "hello",
protocol: sshDataProtocol,
providerId: config.providerId,
channelId,
token: config.token,
});
logger("info", "ssh_data_channel_ready", { providerId: config.providerId, channelId, host: config.providerDataHost, port: config.providerDataPort });
sendHeartbeat();
});
socket.on("data", (chunk) => consumeSshDataBytes(channel, Buffer.from(chunk)));
socket.on("error", (error) => {
channel.lastError = error.message;
sshDataLastError = error.message;
logger("warn", "ssh_data_channel_error", { providerId: config.providerId, channelId, error: error.message });
});
socket.on("close", () => {
clearTimeout(timeout);
logger("warn", "ssh_data_channel_close", { providerId: config.providerId, channelId, lastError: channel.lastError });
closeSshDataChannel(channel, channel.lastError ?? "socket closed");
});
}
function ensureSshDataPool(): void {
if (!sshDataPoolDesired || !isHostSshConfigured()) return;
const status = sshDataPoolStatus();
const needed = config.providerDataPoolSize - status.total;
for (let index = 0; index < needed; index += 1) {
openSshDataChannel();
}
}
function closeSshDataPool(reason: string): void {
sshDataPoolDesired = false;
for (const channel of sshDataChannels.values()) {
channel.status = "closed";
try {
channel.socket.destroy(new Error(reason));
} catch {
/* socket may already be closed */
}
}
sshDataChannels.clear();
}
function acquireSshDataChannel(sessionId: string, channelId: string): SshDataChannel | null {
const channel = sshDataChannels.get(channelId);
if (channel === undefined || channel.status !== "ready" || channel.activeSessionId !== null) {
return null;
}
channel.status = "claimed";
channel.activeSessionId = sessionId;
return channel;
}
function releaseSshDataChannel(channelId: string, sessionId: string): void {
const channel = sshDataChannels.get(channelId);
if (channel === undefined) return;
if (channel.activeSessionId !== sessionId) return;
channel.activeSessionId = null;
if (channel.status !== "closed" && !channel.socket.destroyed) channel.status = "ready";
ensureSshDataPool();
}
function sendRegister(): void {
const capabilities = ["heartbeat", "system.status", "docker.status", "docker.ps", "provider.upgrade", "microservice.http", "microservice.http.cache", "microservice.http.tunnel", "echo"];
if (isHostSshConfigured()) capabilities.push("host.ssh");
if (isHostSshConfigured()) capabilities.push("host.ssh", "host.ssh.tcp-pool");
if (config.egressProxyEnabled) capabilities.push("network.egress-proxy");
sendJson({
type: "register",
@@ -1499,21 +1787,15 @@ async function pumpHostSshOutput(sessionId: string, streamName: "stdout" | "stde
while (true) {
const chunk = await reader.read();
if (chunk.done) return;
sendJson({
type: "host_ssh_data",
providerId: config.providerId,
sessionId,
stream: streamName,
data: Buffer.from(chunk.value).toString("base64"),
encoding: "base64",
at: new Date().toISOString(),
});
const sent = sendSshDataSessionFrame(sessionId, { type: "data", stream: streamName, at: new Date().toISOString() }, Buffer.from(chunk.value));
if (!sent) throw new Error(`ssh tcp data channel is not available for ${sessionId}`);
}
}
function sendHostSshError(sessionId: string, error: unknown): void {
const message = error instanceof Error ? error.message : String(error);
logger("error", "host_ssh_session_error", { sessionId, error: message });
if (sendSshDataSessionFrame(sessionId, { type: "error", message, at: new Date().toISOString() })) return;
sendJson({
type: "host_ssh_error",
providerId: config.providerId,
@@ -1528,6 +1810,15 @@ function startHostSshSession(message: CoreHostSshOpenMessage): void {
sendHostSshError(message.sessionId, `host SSH session already exists: ${message.sessionId}`);
return;
}
if (message.transport !== "tcp-pool" || typeof message.dataChannelId !== "string" || message.dataChannelId.length === 0) {
sendHostSshError(message.sessionId, "host SSH requires tcp-pool transport; upgrade backend-core/provider-gateway together");
return;
}
const dataChannel = acquireSshDataChannel(message.sessionId, message.dataChannelId);
if (dataChannel === null) {
sendHostSshError(message.sessionId, `requested ssh tcp data channel is not ready: ${message.dataChannelId}`);
return;
}
try {
const command = typeof message.command === "string" && message.command.length > 0 ? message.command : null;
const allocateTty = typeof message.tty === "boolean" ? message.tty : command === null;
@@ -1537,73 +1828,43 @@ function startHostSshSession(message: CoreHostSshOpenMessage): void {
stdout: "pipe",
stderr: "pipe",
});
hostSshSessions.set(message.sessionId, { proc, openedAt: Date.now(), tty: allocateTty });
sendJson({
type: "host_ssh_opened",
providerId: config.providerId,
hostSshSessions.set(message.sessionId, { proc, openedAt: Date.now(), tty: allocateTty, dataChannelId: message.dataChannelId });
writeSshDataFrame(dataChannel, {
type: "opened",
sessionId: message.sessionId,
providerId: config.providerId,
at: new Date().toISOString(),
});
const stdoutDone = pumpHostSshOutput(message.sessionId, "stdout", proc.stdout);
const stderrDone = pumpHostSshOutput(message.sessionId, "stderr", proc.stderr);
Promise.all([stdoutDone, stderrDone, proc.exited])
.then(([, , exitCode]) => {
hostSshSessions.delete(message.sessionId);
sendJson({
type: "host_ssh_exit",
providerId: config.providerId,
sessionId: message.sessionId,
sendSshDataSessionFrame(message.sessionId, {
type: "exit",
exitCode,
signal: null,
at: new Date().toISOString(),
});
hostSshSessions.delete(message.sessionId);
releaseSshDataChannel(message.dataChannelId, message.sessionId);
})
.catch((error) => {
hostSshSessions.delete(message.sessionId);
sendHostSshError(message.sessionId, error);
hostSshSessions.delete(message.sessionId);
releaseSshDataChannel(message.dataChannelId, message.sessionId);
});
logger("info", "host_ssh_session_started", { sessionId: message.sessionId, hasCommand: command !== null, tty: allocateTty, cwd: message.cwd ?? null });
logger("info", "host_ssh_session_started", { sessionId: message.sessionId, hasCommand: command !== null, tty: allocateTty, cwd: message.cwd ?? null, transport: "tcp-pool", dataChannelId: message.dataChannelId });
} catch (error) {
sendHostSshError(message.sessionId, error);
releaseSshDataChannel(message.dataChannelId, message.sessionId);
}
}
function writeHostSshInput(message: CoreHostSshInputMessage): void {
const session = hostSshSessions.get(message.sessionId);
if (session === undefined) {
sendHostSshError(message.sessionId, `unknown host SSH session: ${message.sessionId}`);
return;
}
try {
const stdin = session.proc.stdin as HostSshStdin | undefined;
if (stdin === undefined) throw new Error("host SSH stdin is not available");
const chunk = Buffer.from(message.data, "base64");
if (!session.tty && chunk.length === 1 && chunk[0] === 4) return;
stdin.write(chunk);
} catch (error) {
sendHostSshError(message.sessionId, error);
}
}
function resizeHostSshSession(message: CoreHostSshResizeMessage): void {
logger("debug", "host_ssh_resize_requested", { sessionId: message.sessionId, cols: message.cols, rows: message.rows });
}
function eofHostSshSession(message: CoreHostSshEofMessage): void {
const session = hostSshSessions.get(message.sessionId);
function closeHostSshSession(sessionId: string): void {
const session = hostSshSessions.get(sessionId);
if (session === undefined) return;
try {
const stdin = session.proc.stdin as HostSshStdin | undefined;
if (stdin !== undefined) stdin.end();
} catch (error) {
sendHostSshError(message.sessionId, error);
}
}
function closeHostSshSession(message: CoreHostSshCloseMessage): void {
const session = hostSshSessions.get(message.sessionId);
if (session === undefined) return;
hostSshSessions.delete(message.sessionId);
hostSshSessions.delete(sessionId);
if (session.dataChannelId !== undefined) releaseSshDataChannel(session.dataChannelId, sessionId);
session.proc.kill("SIGTERM");
setTimeout(() => {
try {
@@ -2309,22 +2570,6 @@ function handleMessage(raw: MessageEvent<string>): void {
startHostSshSession(parsed as CoreHostSshOpenMessage);
return;
}
if (parsed.type === "host_ssh_input") {
writeHostSshInput(parsed as CoreHostSshInputMessage);
return;
}
if (parsed.type === "host_ssh_resize") {
resizeHostSshSession(parsed as CoreHostSshResizeMessage);
return;
}
if (parsed.type === "host_ssh_eof") {
eofHostSshSession(parsed as CoreHostSshEofMessage);
return;
}
if (parsed.type === "host_ssh_close") {
closeHostSshSession(parsed as CoreHostSshCloseMessage);
return;
}
if (
parsed.type === "egress_tcp_opened" ||
parsed.type === "egress_tcp_data" ||
@@ -2373,6 +2618,7 @@ function enterUpgradeSleep(durationMs: number): void {
upgradeSleepUntil = Date.now() + sleepMs;
logger("warn", "upgrade_sleep_enter", { providerId: config.providerId, sleepMs, wakeAt: new Date(upgradeSleepUntil).toISOString() });
clearConnectionTimers();
closeSshDataPool("provider upgrade sleep");
try {
socket?.close(4000, "provider upgrade sleep");
} catch (error) {
@@ -2407,6 +2653,10 @@ function connect(): void {
reconnectAttempt = 0;
logger("info", "connect_open", { providerId: config.providerId });
sendRegister();
if (isHostSshConfigured()) {
sshDataPoolDesired = true;
ensureSshDataPool();
}
sendHeartbeat();
sendSystemStatus().catch((error) => logger("error", "system_status_initial_failed", { error: String(error) }));
sendDockerStatus().catch((error) => logger("error", "docker_status_initial_failed", { error: String(error) }));
@@ -2425,6 +2675,7 @@ function connect(): void {
socket.addEventListener("close", (event) => {
logger("warn", "connect_close", { code: event.code, reason: event.reason });
clearConnectionTimers();
closeSshDataPool("provider control socket closed");
providerEgressProxy?.closeAll("provider-gateway core socket closed");
scheduleReconnect();
});
@@ -2445,6 +2696,7 @@ process.on("SIGTERM", () => {
}
hostSshSessions.clear();
providerEgressProxy?.close();
closeSshDataPool("provider shutdown");
socket?.close(1000, "provider shutdown");
process.exit(0);
});
src/components/shared/src/index.ts
+2 -60
View File
@@ -153,6 +153,8 @@ export interface CoreDispatchMessage {
export interface CoreHostSshOpenMessage {
type: "host_ssh_open";
sessionId: string;
transport: "tcp-pool";
dataChannelId: string;
cwd?: string;
command?: string;
tty?: boolean;
@@ -160,30 +162,6 @@ export interface CoreHostSshOpenMessage {
rows?: number;
}
export interface CoreHostSshInputMessage {
type: "host_ssh_input";
sessionId: string;
data: string;
encoding: "base64";
}
export interface CoreHostSshResizeMessage {
type: "host_ssh_resize";
sessionId: string;
cols: number;
rows: number;
}
export interface CoreHostSshCloseMessage {
type: "host_ssh_close";
sessionId: string;
}
export interface CoreHostSshEofMessage {
type: "host_ssh_eof";
sessionId: string;
}
export interface CoreAcknowledgeMessage {
type: "ack";
requestId: string;
@@ -191,32 +169,6 @@ export interface CoreAcknowledgeMessage {
message: string;
}
export interface ProviderHostSshOpenedMessage {
type: "host_ssh_opened";
providerId: string;
sessionId: string;
at: string;
}
export interface ProviderHostSshDataMessage {
type: "host_ssh_data";
providerId: string;
sessionId: string;
stream: "stdout" | "stderr";
data: string;
encoding: "base64";
at: string;
}
export interface ProviderHostSshExitMessage {
type: "host_ssh_exit";
providerId: string;
sessionId: string;
exitCode: number | null;
signal: string | null;
at: string;
}
export interface ProviderHostSshErrorMessage {
type: "host_ssh_error";
providerId: string;
@@ -289,9 +241,6 @@ export type ProviderToCoreMessage =
| ProviderSystemStatusMessage
| ProviderDockerStatusMessage
| ProviderTaskStatusMessage
| ProviderHostSshOpenedMessage
| ProviderHostSshDataMessage
| ProviderHostSshExitMessage
| ProviderHostSshErrorMessage
| ProviderEgressTcpOpenMessage
| ProviderEgressTcpDataMessage
@@ -302,10 +251,6 @@ export type CoreToProviderMessage =
| CoreDispatchMessage
| CoreHttpTunnelRequestMessage
| CoreHostSshOpenMessage
| CoreHostSshInputMessage
| CoreHostSshResizeMessage
| CoreHostSshCloseMessage
| CoreHostSshEofMessage
| CoreEgressTcpOpenedMessage
| CoreEgressTcpDataMessage
| CoreEgressTcpCloseMessage
@@ -386,9 +331,6 @@ export function isProviderToCoreMessage(value: unknown): value is ProviderToCore
msg.type === "system_status" ||
msg.type === "docker_status" ||
msg.type === "task_status" ||
msg.type === "host_ssh_opened" ||
msg.type === "host_ssh_data" ||
msg.type === "host_ssh_exit" ||
msg.type === "host_ssh_error" ||
msg.type === "egress_tcp_open" ||
msg.type === "egress_tcp_data" ||