danvics/zerobyte
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

feat: enforce protocol version between agent and controller (#938)

Browse source 
* feat: enforce protocol version between agent and controller

* chore: add logging for protocol rejected message
This commit is contained in:
Nico 2026-06-03 18:29:36 +02:00 committed by GitHub
parent be3182793d
commit 333c11986d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 270 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

37
app/server/modules/agents/__tests__/controller-runtime.test.ts
View file

@ -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", {

79
app/server/modules/agents/__tests__/session.test.ts
View file

@ -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<AgentMessage, { type: "backup.progress" }>;
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() });

4
app/server/modules/agents/agents-manager.ts
View file

@ -247,6 +247,10 @@ const handleAgentManagerEvent = (event: AgentManagerEvent) => {
);
break;
}
case "agent.protocolRejected": {
logger.warn(`Rejected agent protocol for ${event.agentName} (${event.agentId}): ${event.payload.reason}`);
break;
}
case "backup.started": {
getActiveBackupRun(event.payload.jobId, event.payload.scheduleId, event.type, event.agentId);
break;

8
app/server/modules/agents/controller/server.ts
View file

@ -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));

43
app/server/modules/agents/controller/session.ts
View file

@ -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<AgentConnectionData>;
type SessionState = {
isReady: boolean;
protocolVersion: number | null;
lastSeenAt: number | null;
lastPongAt: number | null;
};
@ -35,6 +40,7 @@ type PendingCommand = { deferred: Deferred.Deferred<VolumeCommandResponsePayload
export type ControllerAgentSessionEvent =
| Exclude<AgentMessage, { type: "volume.commandResult" }>
| { 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<string, PendingCommand>());
const state = yield* Ref.make<SessionState>({
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}`,
);

3
apps/agent/src/controller-session.ts
View file

@ -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 },

99
packages/contracts/src/agent-protocol.ts
View file

@ -21,6 +21,27 @@ import {
const compressionModeSchema = z.enum(["off", "auto", "max"]) satisfies z.ZodType<CompressionMode>;
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 = <TType extends ControllerMessage["type"]>(
type: TType,
payload: PayloadForMessage<ControllerMessage, TType>,