From 58c5879c6ce8e996792c9867a6bca6b0019b3752 Mon Sep 17 00:00:00 2001 From: Nicolas Meienberger Date: Wed, 3 Jun 2026 17:54:03 +0200 Subject: [PATCH] feat: enforce protocol version between agent and controller --- .../__tests__/controller-runtime.test.ts | 37 +++++++ .../modules/agents/__tests__/session.test.ts | 79 ++++++++++++++- .../modules/agents/controller/server.ts | 8 ++ .../modules/agents/controller/session.ts | 43 +++++++- apps/agent/src/controller-session.ts | 3 +- packages/contracts/src/agent-protocol.ts | 99 +++++++++++++++++++ 6 files changed, 266 insertions(+), 3 deletions(-) diff --git a/app/server/modules/agents/__tests__/controller-runtime.test.ts b/app/server/modules/agents/__tests__/controller-runtime.test.ts index 801646af..82a19a9f 100644 --- a/app/server/modules/agents/__tests__/controller-runtime.test.ts +++ b/app/server/modules/agents/__tests__/controller-runtime.test.ts @@ -207,6 +207,7 @@ test("websocket lifecycle updates agent connection status", async () => { expect(agentsServiceMocks.markAgentOnline).toHaveBeenCalledWith(LOCAL_AGENT_ID, expect.any(Number), { backup: true, protocolVersion: 1, + protocolCompatible: true, hostname: "host", platform: "linux", }); @@ -215,6 +216,40 @@ test("websocket lifecycle updates agent connection status", async () => { expect(stop).toHaveBeenCalledWith(true); }); +test("websocket protocol rejection forwards the event and closes the connection", async () => { + const serve = vi + .spyOn(Bun, "serve") + .mockReturnValue(fromPartial({ port: 3001, stop: vi.fn(() => Promise.resolve()) })); + const { runtime, onEvent } = await startRuntime(vi.fn()); + const websocket = serve.mock.calls[0]?.[0].websocket; + const socket = createSocket("connection-1"); + + await websocket?.open?.(fromPartial(socket)); + await websocket?.message?.( + fromPartial(socket), + JSON.stringify({ + type: "agent.ready", + payload: { + protocolVersion: 2, + hostname: "host", + platform: "linux", + }, + }), + ); + await Effect.runPromise(runtime.stop); + + expect(onEvent).toHaveBeenCalledWith( + expect.objectContaining({ + type: "agent.protocolRejected", + agentId: LOCAL_AGENT_ID, + agentName: LOCAL_AGENT_NAME, + payload: expect.objectContaining({ reason: "agent_too_new" }), + }), + ); + expect(agentsServiceMocks.markAgentOffline).toHaveBeenCalledWith(LOCAL_AGENT_ID); + expect(socket.close).toHaveBeenCalledWith(1002, "agent_too_new"); +}); + test("websocket restore events are forwarded with agent metadata", async () => { const serve = vi .spyOn(Bun, "serve") @@ -224,6 +259,8 @@ test("websocket restore events are forwarded with agent metadata", async () => { const socket = createSocket("connection-1"); await websocket?.open?.(fromPartial(socket)); + await websocket?.message?.(fromPartial(socket), createAgentMessage("agent.ready", readyPayload)); + onEvent.mockClear(); await websocket?.message?.( fromPartial(socket), createAgentMessage("restore.completed", { diff --git a/app/server/modules/agents/__tests__/session.test.ts b/app/server/modules/agents/__tests__/session.test.ts index 6c41a880..fb076077 100644 --- a/app/server/modules/agents/__tests__/session.test.ts +++ b/app/server/modules/agents/__tests__/session.test.ts @@ -2,7 +2,11 @@ import { Effect, Exit, Fiber, Scope } from "effect"; import { expect, test, vi } from "vitest"; import waitForExpect from "wait-for-expect"; import { fromPartial } from "@total-typescript/shoehorn"; -import { createAgentMessage, type AgentMessage } from "@zerobyte/contracts/agent-protocol"; +import { + createAgentMessage, + SUPPORTED_AGENT_PROTOCOL_MAX_VERSION, + type AgentMessage, +} from "@zerobyte/contracts/agent-protocol"; import type { Volume } from "@zerobyte/contracts/volumes"; import { LOCAL_AGENT_ID, LOCAL_AGENT_KIND, LOCAL_AGENT_NAME } from "../constants"; import { createControllerAgentSession } from "../controller/session"; @@ -133,6 +137,19 @@ test("invalid inbound messages are ignored", () => { const onEvent = vi.fn(() => Effect.void); const { session, close } = createSession(onEvent); + Effect.runSync( + session.handleMessage( + createAgentMessage("agent.ready", { + agentId: LOCAL_AGENT_ID, + protocolVersion: 1, + hostname: "host", + platform: "linux", + capabilities: { backup: true }, + }), + ), + ); + onEvent.mockClear(); + Effect.runSync(session.handleMessage("not json")); Effect.runSync(session.handleMessage(JSON.stringify({ type: "backup.progress", payload: {} }))); @@ -193,6 +210,19 @@ test("backup agent messages are forwarded unchanged", () => { }, } satisfies Extract; + Effect.runSync( + session.handleMessage( + createAgentMessage("agent.ready", { + agentId: LOCAL_AGENT_ID, + protocolVersion: 1, + hostname: "host", + platform: "linux", + capabilities: { backup: true }, + }), + ), + ); + onEvent.mockClear(); + Effect.runSync(session.handleMessage(createAgentMessage(message.type, message.payload))); expect(onEvent).toHaveBeenCalledWith( @@ -208,6 +238,53 @@ test("backup agent messages are forwarded unchanged", () => { close(); }); +test("unsupported agent protocol rejects startup and closes the session", () => { + const onEvent = vi.fn(() => Effect.void); + const { session, socket } = createSession(onEvent); + + Effect.runSync( + session.handleMessage( + JSON.stringify({ + type: "agent.ready", + payload: { + protocolVersion: SUPPORTED_AGENT_PROTOCOL_MAX_VERSION + 1, + hostname: "host", + platform: "linux", + }, + }), + ), + ); + + expect(Effect.runSync(session.isReady())).toBe(false); + expect(onEvent).toHaveBeenCalledWith({ + type: "agent.protocolRejected", + payload: expect.objectContaining({ + reason: "agent_too_new", + protocolVersion: SUPPORTED_AGENT_PROTOCOL_MAX_VERSION + 1, + hostname: "host", + platform: "linux", + }), + }); + expect(socket.close).toHaveBeenCalledWith(1002, "agent_too_new"); +}); + +test("pre-ready non-ready messages reject startup and close the session", () => { + const onEvent = vi.fn(() => Effect.void); + const { session, socket } = createSession(onEvent); + + Effect.runSync(session.handleMessage(createAgentMessage("heartbeat.pong", { sentAt: 123 }))); + + expect(Effect.runSync(session.isReady())).toBe(false); + expect(onEvent).toHaveBeenCalledWith({ + type: "agent.protocolRejected", + payload: expect.objectContaining({ + reason: "unexpected_startup_message", + messageType: "heartbeat.pong", + }), + }); + expect(socket.close).toHaveBeenCalledWith(1002, "unexpected_startup_message"); +}); + test("a dropped backup.cancel closes the session and reports a transport disconnect", async () => { const send = vi.fn(() => 0); const socket = createSocket({ send, close: vi.fn() }); diff --git a/app/server/modules/agents/controller/server.ts b/app/server/modules/agents/controller/server.ts index e5c431b0..65ef356b 100644 --- a/app/server/modules/agents/controller/server.ts +++ b/app/server/modules/agents/controller/server.ts @@ -3,6 +3,7 @@ import { logger } from "@zerobyte/core/node"; import { toMessage } from "@zerobyte/core/utils"; import type { AgentMessage, + AgentProtocolRejection, BackupCancelPayload, BackupRunPayload, RestoreCancelPayload, @@ -26,6 +27,7 @@ type AgentEventContext = { export type AgentManagerEvent = | (AgentEventContext & { type: "agent.disconnected" }) + | (AgentEventContext & { type: "agent.protocolRejected"; payload: AgentProtocolRejection }) | (AgentEventContext & AgentMessage); type ControllerAgentSessionHandle = { @@ -92,11 +94,17 @@ export function createAgentManagerRuntime(onEvent: (event: AgentManagerEvent) => await agentsService.markAgentOnline(agentId, at, { ...event.payload.capabilities, protocolVersion: event.payload.protocolVersion, + protocolCompatible: true, hostname: event.payload.hostname, platform: event.payload.platform, }); }); } + case "agent.protocolRejected": { + return Effect.sync(() => + onEvent({ type: "agent.protocolRejected", agentId, agentName, payload: event.payload }), + ); + } case "heartbeat.pong": { const at = Date.now(); return Effect.promise(() => agentsService.markAgentSeen(agentId, at)); diff --git a/app/server/modules/agents/controller/session.ts b/app/server/modules/agents/controller/session.ts index c17fafff..0f2522b1 100644 --- a/app/server/modules/agents/controller/session.ts +++ b/app/server/modules/agents/controller/session.ts @@ -3,7 +3,11 @@ import type { AgentKind } from "../../../db/schema"; import { createControllerMessage, parseAgentMessage, + parseAgentStartupMessage, + SUPPORTED_AGENT_PROTOCOL_MAX_VERSION, + SUPPORTED_AGENT_PROTOCOL_MIN_VERSION, type AgentMessage, + type AgentProtocolRejection, type BackupCancelPayload, type BackupRunPayload, type ControllerWireMessage, @@ -27,6 +31,7 @@ type AgentSocket = Bun.ServerWebSocket; type SessionState = { isReady: boolean; + protocolVersion: number | null; lastSeenAt: number | null; lastPongAt: number | null; }; @@ -35,6 +40,7 @@ type PendingCommand = { deferred: Deferred.Deferred + | { type: "agent.protocolRejected"; payload: AgentProtocolRejection } | { type: "agent.disconnected" }; export type ControllerAgentSession = { @@ -59,6 +65,7 @@ export const createControllerAgentSession = ( const pendingCommands = yield* Ref.make(new Map()); const state = yield* Ref.make({ isReady: false, + protocolVersion: null, lastSeenAt: null, lastPongAt: null, }); @@ -183,7 +190,12 @@ export const createControllerAgentSession = ( switch (message.type) { case "agent.ready": { const readyAt = Date.now(); - yield* updateState((current) => ({ ...current, isReady: true, lastSeenAt: readyAt })); + yield* updateState((current) => ({ + ...current, + isReady: true, + protocolVersion: message.payload.protocolVersion, + lastSeenAt: readyAt, + })); yield* logger.effect.info(`Agent "${socket.data.agentName}" (${socket.data.agentId}) is ready`); yield* onEvent(message); break; @@ -209,10 +221,30 @@ export const createControllerAgentSession = ( } }); + const rejectStartupMessage = (rejection: AgentProtocolRejection) => + Effect.gen(function* () { + yield* logger.effect.warn( + `Rejecting startup message from agent ${socket.data.agentId}: ${rejection.reason}`, + ); + yield* onEvent({ type: "agent.protocolRejected", payload: rejection }); + yield* Effect.sync(() => socket.close(1002, rejection.reason)); + yield* closeSession(); + }); + return { connectionId: socket.data.id, handleMessage: (data: string) => { return Effect.gen(function* () { + const currentState = yield* Ref.get(state); + + if (!currentState.isReady) { + const startupMessage = parseAgentStartupMessage(data); + if (!("success" in startupMessage)) { + yield* rejectStartupMessage(startupMessage); + return; + } + } + const parsed = parseAgentMessage(data); if (parsed === null) { @@ -221,6 +253,15 @@ export const createControllerAgentSession = ( } if (!parsed.success) { + if (!currentState.isReady) { + yield* rejectStartupMessage({ + reason: "invalid_agent_ready", + supportedProtocolMinVersion: SUPPORTED_AGENT_PROTOCOL_MIN_VERSION, + supportedProtocolMaxVersion: SUPPORTED_AGENT_PROTOCOL_MAX_VERSION, + }); + return; + } + yield* logger.effect.warn( `Invalid agent message from ${socket.data.agentId}: ${parsed.error.message}`, ); diff --git a/apps/agent/src/controller-session.ts b/apps/agent/src/controller-session.ts index 69d2a05d..bb1e9ca0 100644 --- a/apps/agent/src/controller-session.ts +++ b/apps/agent/src/controller-session.ts @@ -1,5 +1,6 @@ import { Effect, Fiber, Queue, Ref } from "effect"; import { + AGENT_PROTOCOL_VERSION, createAgentMessage, parseControllerMessage, type AgentWireMessage, @@ -141,7 +142,7 @@ export const createControllerSession = (ws: WebSocket): ControllerSession => { offerOutbound( createAgentMessage("agent.ready", { agentId: "", - protocolVersion: 1, + protocolVersion: AGENT_PROTOCOL_VERSION, hostname: resolveResticHostname(), platform: process.platform, capabilities: { backup: true, restore: true, volume: true, restic: true }, diff --git a/packages/contracts/src/agent-protocol.ts b/packages/contracts/src/agent-protocol.ts index 748661a8..295d2e5d 100644 --- a/packages/contracts/src/agent-protocol.ts +++ b/packages/contracts/src/agent-protocol.ts @@ -21,6 +21,27 @@ import { const compressionModeSchema = z.enum(["off", "auto", "max"]) satisfies z.ZodType; +export const AGENT_PROTOCOL_VERSION = 1; +export const SUPPORTED_AGENT_PROTOCOL_MIN_VERSION = 1; +export const SUPPORTED_AGENT_PROTOCOL_MAX_VERSION = 1; + +export type AgentProtocolRejectionReason = + | "agent_too_old" + | "agent_too_new" + | "invalid_agent_ready" + | "unexpected_startup_message" + | "invalid_startup_json"; + +export type AgentProtocolRejection = { + reason: AgentProtocolRejectionReason; + protocolVersion?: number; + supportedProtocolMinVersion: number; + supportedProtocolMaxVersion: number; + hostname?: string; + platform?: string; + messageType?: string; +}; + const backupExecutionOptionsSchema = z.object({ oneFileSystem: z.boolean(), excludePatterns: z.array(z.string()).nullable(), @@ -179,6 +200,21 @@ const agentReadySchema = z.object({ }), }); +const agentStartupMessageSchema = z.object({ + type: z.string(), + payload: z.unknown().optional(), +}); + +const stableAgentReadySchema = z.object({ + type: z.literal("agent.ready"), + payload: z.object({ + protocolVersion: z.number(), + hostname: z.string().optional(), + platform: z.string().optional(), + capabilities: z.record(z.string(), z.unknown()).optional(), + }), +}); + const backupStartedSchema = z.object({ type: z.literal("backup.started"), payload: z.object({ jobId: z.string(), scheduleId: z.string() }), @@ -303,6 +339,69 @@ export const parseAgentMessage = (data: string) => { return agentMessageSchema.safeParse(parsed); }; +export const parseAgentStartupMessage = (data: string): AgentProtocolRejection | { success: true; data: unknown } => { + const parsed = safeJsonParse(data); + if (parsed === null) { + return { + reason: "invalid_startup_json", + supportedProtocolMinVersion: SUPPORTED_AGENT_PROTOCOL_MIN_VERSION, + supportedProtocolMaxVersion: SUPPORTED_AGENT_PROTOCOL_MAX_VERSION, + }; + } + + const startupMessage = agentStartupMessageSchema.safeParse(parsed); + if (!startupMessage.success) { + return { + reason: "unexpected_startup_message", + supportedProtocolMinVersion: SUPPORTED_AGENT_PROTOCOL_MIN_VERSION, + supportedProtocolMaxVersion: SUPPORTED_AGENT_PROTOCOL_MAX_VERSION, + }; + } + + if (startupMessage.data.type !== "agent.ready") { + return { + reason: "unexpected_startup_message", + messageType: startupMessage.data.type, + supportedProtocolMinVersion: SUPPORTED_AGENT_PROTOCOL_MIN_VERSION, + supportedProtocolMaxVersion: SUPPORTED_AGENT_PROTOCOL_MAX_VERSION, + }; + } + + const readyMessage = stableAgentReadySchema.safeParse(parsed); + if (!readyMessage.success) { + return { + reason: "invalid_agent_ready", + supportedProtocolMinVersion: SUPPORTED_AGENT_PROTOCOL_MIN_VERSION, + supportedProtocolMaxVersion: SUPPORTED_AGENT_PROTOCOL_MAX_VERSION, + }; + } + + const { protocolVersion, hostname, platform } = readyMessage.data.payload; + if (protocolVersion < SUPPORTED_AGENT_PROTOCOL_MIN_VERSION) { + return { + reason: "agent_too_old", + protocolVersion, + hostname, + platform, + supportedProtocolMinVersion: SUPPORTED_AGENT_PROTOCOL_MIN_VERSION, + supportedProtocolMaxVersion: SUPPORTED_AGENT_PROTOCOL_MAX_VERSION, + }; + } + + if (protocolVersion > SUPPORTED_AGENT_PROTOCOL_MAX_VERSION) { + return { + reason: "agent_too_new", + protocolVersion, + hostname, + platform, + supportedProtocolMinVersion: SUPPORTED_AGENT_PROTOCOL_MIN_VERSION, + supportedProtocolMaxVersion: SUPPORTED_AGENT_PROTOCOL_MAX_VERSION, + }; + } + + return { success: true, data: parsed }; +}; + export const createControllerMessage = ( type: TType, payload: PayloadForMessage,