297 lines
9.3 KiB
JavaScript
297 lines
9.3 KiB
JavaScript
import http from "node:http";
|
|
import { randomUUID } from "node:crypto";
|
|
import { WebSocketServer } from "ws";
|
|
|
|
function parseJsonBody(req) {
|
|
return new Promise((resolve, reject) => {
|
|
let raw = "";
|
|
req.on("data", (chunk) => {
|
|
raw += chunk.toString("utf8");
|
|
});
|
|
req.on("end", () => {
|
|
try {
|
|
resolve(raw === "" ? {} : JSON.parse(raw));
|
|
} catch (error) {
|
|
reject(error);
|
|
}
|
|
});
|
|
req.on("error", reject);
|
|
});
|
|
}
|
|
|
|
function jsonResponse(res, statusCode, body) {
|
|
res.writeHead(statusCode, { "content-type": "application/json" });
|
|
res.end(JSON.stringify(body));
|
|
}
|
|
|
|
export function createBrokerServer(options = {}) {
|
|
const authMode = options.authMode || process.env.EDGE_AUTH_MODE || "stub";
|
|
const sharedSecret = options.sharedSecret ?? process.env.EDGE_BROKER_SHARED_SECRET ?? "";
|
|
const commandTimeoutMs = options.commandTimeoutMs ?? 10000;
|
|
const agents = new Map();
|
|
const pendingCommands = new Map();
|
|
const browserSessions = new Map();
|
|
|
|
const validateAgent = options.validateAgent || (async ({ gatewayId }) => ({ id: gatewayId }));
|
|
const validateShellSession = options.validateShellSession || (async ({ token }) => ({ id: token, gateway_id: 1, reason: "stub" }));
|
|
const closeShellSession = options.closeShellSession || (async () => ({}));
|
|
|
|
const server = http.createServer(async (req, res) => {
|
|
try {
|
|
const url = new URL(req.url, "http://localhost");
|
|
if (req.method === "POST" && /^\/api\/gateways\/\d+\/commands$/.test(url.pathname)) {
|
|
if (sharedSecret && req.headers["x-edge-broker-secret"] !== sharedSecret) {
|
|
jsonResponse(res, 403, { error: "Forbidden" });
|
|
return;
|
|
}
|
|
|
|
const gatewayId = url.pathname.split("/")[3];
|
|
const agent = agents.get(String(gatewayId));
|
|
if (!agent || agent.readyState !== 1) {
|
|
jsonResponse(res, 503, { error: "Gateway agent is offline" });
|
|
return;
|
|
}
|
|
|
|
const body = await parseJsonBody(req);
|
|
const commandId = randomUUID();
|
|
const promise = new Promise((resolve, reject) => {
|
|
const timeout = setTimeout(() => {
|
|
pendingCommands.delete(commandId);
|
|
reject(new Error("Agent command timed out"));
|
|
}, commandTimeoutMs);
|
|
pendingCommands.set(commandId, {
|
|
resolve,
|
|
reject,
|
|
timeout,
|
|
});
|
|
});
|
|
|
|
agent.send(JSON.stringify({
|
|
type: "COMMAND",
|
|
commandId,
|
|
commandType: body.commandType,
|
|
payload: body.payload || {},
|
|
}));
|
|
|
|
try {
|
|
const result = await promise;
|
|
jsonResponse(res, 200, result);
|
|
} catch (error) {
|
|
jsonResponse(res, 504, { ok: false, error: error instanceof Error ? error.message : String(error) });
|
|
}
|
|
return;
|
|
}
|
|
|
|
jsonResponse(res, 404, { error: "Not found" });
|
|
} catch (error) {
|
|
jsonResponse(res, 500, { error: error instanceof Error ? error.message : String(error) });
|
|
}
|
|
});
|
|
|
|
const wss = new WebSocketServer({ noServer: true });
|
|
server.on("upgrade", async (req, socket, head) => {
|
|
const url = new URL(req.url, "http://localhost");
|
|
try {
|
|
if (url.pathname === "/ws/agent") {
|
|
const gatewayId = String(url.searchParams.get("gatewayId") || "");
|
|
const token = String(url.searchParams.get("token") || "");
|
|
if (gatewayId === "" || token === "") {
|
|
socket.destroy();
|
|
return;
|
|
}
|
|
if (authMode !== "stub") {
|
|
await validateAgent({ gatewayId, token, headers: req.headers });
|
|
}
|
|
wss.handleUpgrade(req, socket, head, (ws) => {
|
|
ws.gatewayId = gatewayId;
|
|
agents.set(gatewayId, ws);
|
|
wss.emit("connection", ws, req);
|
|
});
|
|
return;
|
|
}
|
|
|
|
if (url.pathname === "/ws/browser-shell") {
|
|
const token = String(url.searchParams.get("token") || "");
|
|
if (token === "") {
|
|
socket.destroy();
|
|
return;
|
|
}
|
|
const session = authMode === "stub"
|
|
? await validateShellSession({ token })
|
|
: await validateShellSession({ token, headers: req.headers });
|
|
|
|
wss.handleUpgrade(req, socket, head, (ws) => {
|
|
ws.sessionToken = token;
|
|
ws.sessionInfo = session;
|
|
browserSessions.set(String(session.id), { ws, session, transcript: "" });
|
|
const agent = agents.get(String(session.gateway_id));
|
|
if (agent && agent.readyState === 1) {
|
|
agent.send(JSON.stringify({
|
|
type: "OPEN_ROOT_SHELL",
|
|
payload: {
|
|
sessionId: String(session.id),
|
|
reason: session.reason,
|
|
},
|
|
}));
|
|
}
|
|
wss.emit("connection", ws, req);
|
|
});
|
|
return;
|
|
}
|
|
} catch {
|
|
socket.destroy();
|
|
return;
|
|
}
|
|
|
|
socket.destroy();
|
|
});
|
|
|
|
wss.on("connection", (ws) => {
|
|
ws.on("message", async (raw) => {
|
|
const message = JSON.parse(raw.toString());
|
|
|
|
if (ws.gatewayId) {
|
|
if (message.type === "COMMAND_RESULT") {
|
|
const pending = pendingCommands.get(message.commandId);
|
|
if (!pending) {
|
|
return;
|
|
}
|
|
clearTimeout(pending.timeout);
|
|
pendingCommands.delete(message.commandId);
|
|
pending.resolve({
|
|
ok: Boolean(message.ok),
|
|
payload: message.payload,
|
|
error: message.error,
|
|
});
|
|
return;
|
|
}
|
|
|
|
if (["SHELL_OUTPUT", "SHELL_OPENED", "SHELL_EXIT"].includes(message.type)) {
|
|
const sessionRecord = browserSessions.get(String(message.sessionId));
|
|
if (!sessionRecord) {
|
|
return;
|
|
}
|
|
if (message.type === "SHELL_OUTPUT") {
|
|
sessionRecord.transcript += String(message.data || "");
|
|
sessionRecord.ws.send(JSON.stringify({ type: "output", data: String(message.data || "") }));
|
|
}
|
|
if (message.type === "SHELL_OPENED") {
|
|
sessionRecord.ws.send(JSON.stringify({ type: "opened" }));
|
|
}
|
|
if (message.type === "SHELL_EXIT") {
|
|
sessionRecord.ws.send(JSON.stringify({ type: "closed", code: message.code ?? 0 }));
|
|
await closeShellSession(sessionRecord.session.id, sessionRecord.ws.sessionToken, sessionRecord.transcript, "agent_exit");
|
|
browserSessions.delete(String(message.sessionId));
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
if (ws.sessionInfo) {
|
|
const sessionId = String(ws.sessionInfo.id);
|
|
const agent = agents.get(String(ws.sessionInfo.gateway_id));
|
|
if (!agent || agent.readyState !== 1) {
|
|
return;
|
|
}
|
|
if (message.type === "input") {
|
|
agent.send(JSON.stringify({
|
|
type: "SHELL_INPUT",
|
|
payload: {
|
|
sessionId,
|
|
data: String(message.data || ""),
|
|
},
|
|
}));
|
|
}
|
|
if (message.type === "resize") {
|
|
agent.send(JSON.stringify({
|
|
type: "RESIZE_ROOT_SHELL",
|
|
payload: {
|
|
sessionId,
|
|
cols: Number(message.cols || 0),
|
|
rows: Number(message.rows || 0),
|
|
},
|
|
}));
|
|
}
|
|
if (message.type === "close") {
|
|
agent.send(JSON.stringify({
|
|
type: "CLOSE_ROOT_SHELL",
|
|
payload: { sessionId },
|
|
}));
|
|
}
|
|
}
|
|
});
|
|
|
|
ws.on("close", async () => {
|
|
if (ws.gatewayId) {
|
|
agents.delete(String(ws.gatewayId));
|
|
return;
|
|
}
|
|
|
|
if (ws.sessionInfo) {
|
|
const sessionId = String(ws.sessionInfo.id);
|
|
const agent = agents.get(String(ws.sessionInfo.gateway_id));
|
|
if (agent && agent.readyState === 1) {
|
|
agent.send(JSON.stringify({
|
|
type: "CLOSE_ROOT_SHELL",
|
|
payload: { sessionId },
|
|
}));
|
|
}
|
|
const sessionRecord = browserSessions.get(sessionId);
|
|
if (sessionRecord) {
|
|
await closeShellSession(sessionRecord.session.id, ws.sessionToken, sessionRecord.transcript, "browser_closed");
|
|
browserSessions.delete(sessionId);
|
|
}
|
|
}
|
|
});
|
|
});
|
|
|
|
return {
|
|
server,
|
|
listen(port = Number(process.env.PORT || 4300)) {
|
|
return new Promise((resolve) => {
|
|
server.listen(port, () => resolve(server.address()));
|
|
});
|
|
},
|
|
close() {
|
|
return new Promise((resolve, reject) => {
|
|
for (const agent of agents.values()) {
|
|
agent.terminate();
|
|
}
|
|
for (const session of browserSessions.values()) {
|
|
session.ws.terminate();
|
|
}
|
|
for (const pending of pendingCommands.values()) {
|
|
clearTimeout(pending.timeout);
|
|
pending.reject(new Error("Broker shutting down"));
|
|
}
|
|
pendingCommands.clear();
|
|
|
|
wss.close(() => {
|
|
server.close((error) => {
|
|
if (error) {
|
|
reject(error);
|
|
return;
|
|
}
|
|
resolve();
|
|
});
|
|
});
|
|
});
|
|
},
|
|
state: {
|
|
agents,
|
|
browserSessions,
|
|
pendingCommands,
|
|
},
|
|
};
|
|
}
|
|
|
|
if (import.meta.url === `file://${process.argv[1]}`) {
|
|
const broker = createBrokerServer();
|
|
broker.listen().then(() => {
|
|
console.log("TruckWash edge broker listening");
|
|
}).catch((error) => {
|
|
console.error(error instanceof Error ? error.message : String(error));
|
|
process.exitCode = 1;
|
|
});
|
|
}
|