sscoble/pi-mono
1
0
Fork 0
mirror of https://github.com/badlogic/pi-mono.git synced 2026-11-05 23:37:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

fix(protocol): reject cyclic values and organize sources

Browse source
This commit is contained in:
Christian Klotz 2026-07-30 22:07:42 +03:00
parent bc031ae457
commit bd2cfabc5f
No known key found for this signature in database
10 changed files with 207 additions and 199 deletions
Download patch file Download diff file Expand all files Collapse all files

2
packages/protocol/src/cbor-decoder.ts → packages/protocol/src/cbor/decoder.ts
View file

@ -5,7 +5,7 @@ import {
resolveOptions,
textDecoder,
UINT32_BASE,
} from "./cbor-options.ts";
} from "./types.ts";
class CborReader {
private readonly bytes: Uint8Array;

2
packages/protocol/src/cbor-encoder.ts → packages/protocol/src/cbor/encoder.ts
View file

@ -7,7 +7,7 @@ import {
textDecoder,
textEncoder,
UINT32_BASE,
} from "./cbor-options.ts";
} from "./types.ts";
class CborWriter {
private buffer: Uint8Array;

6
packages/protocol/src/cbor.ts → packages/protocol/src/cbor/index.ts
View file

@ -1,9 +1,9 @@
export { decodeCbor } from "./cbor-decoder.ts";
export { encodeCbor } from "./cbor-encoder.ts";
export { decodeCbor } from "./decoder.ts";
export { encodeCbor } from "./encoder.ts";
export {
CborError,
type CborOptions,
DEFAULT_MAX_CBOR_BYTE_LENGTH,
DEFAULT_MAX_CBOR_CONTAINER_LENGTH,
DEFAULT_MAX_CBOR_DEPTH,
} from "./cbor-options.ts";
} from "./types.ts";

0
packages/protocol/src/cbor-options.ts → packages/protocol/src/cbor/types.ts
View file

29
packages/protocol/src/codec.ts
View file

@ -1,5 +1,5 @@
import { Check } from "typebox/value";
import { decodeCbor, encodeCbor } from "./cbor.ts";
import { decodeCbor, encodeCbor } from "./cbor/index.ts";
import {
assertCompleteFrame,
DEFAULT_MAX_FRAME_LENGTH,
@ -7,8 +7,13 @@ import {
FrameDecoder,
type FrameDecoderOptions,
} from "./framing.ts";
import { type ClientMessage, ClientMessageSchema, type ServerMessage, ServerMessageSchema } from "./messages.ts";
import { PROTOCOL_VERSION } from "./schemas.ts";
import {
type ClientMessage,
ClientMessageSchema,
PROTOCOL_VERSION,
type ServerMessage,
ServerMessageSchema,
} from "./schemas.ts";
export class ProtocolValidationError extends Error {
constructor(message: string, _value?: unknown) {
@ -17,25 +22,31 @@ export class ProtocolValidationError extends Error {
}
}
function isProtocolValue(value: unknown, optionalProperty = false): boolean {
function isProtocolValue(value: unknown, optionalProperty = false, ancestors = new Set<object>()): boolean {
if (value === undefined) return optionalProperty;
if (value === null || typeof value === "boolean" || typeof value === "number" || typeof value === "string") {
return true;
}
if (Array.isArray(value)) return value.every((item) => isProtocolValue(item));
if (typeof value !== "object" || Object.getPrototypeOf(value) !== Object.prototype) return false;
return Object.values(value).every((item) => isProtocolValue(item, true));
if (typeof value !== "object" || ancestors.has(value)) return false;
ancestors.add(value);
try {
if (Array.isArray(value)) return value.every((item) => isProtocolValue(item, false, ancestors));
if (Object.getPrototypeOf(value) !== Object.prototype) return false;
return Object.values(value).every((item) => isProtocolValue(item, true, ancestors));
} finally {
ancestors.delete(value);
}
}
export function parseClientMessage(value: unknown): ClientMessage {
if (!Check(ClientMessageSchema, value) || !isProtocolValue(value)) {
if (!isProtocolValue(value) || !Check(ClientMessageSchema, value)) {
throw new ProtocolValidationError("Invalid client protocol message");
}
return value;
}
export function parseServerMessage(value: unknown): ServerMessage {
if (!Check(ServerMessageSchema, value) || !isProtocolValue(value)) {
if (!isProtocolValue(value) || !Check(ServerMessageSchema, value)) {
throw new ProtocolValidationError("Invalid server protocol message");
}
return value;

3
packages/protocol/src/index.ts
View file

@ -1,5 +1,4 @@
export * from "./cbor.ts";
export * from "./cbor/index.ts";
export * from "./codec.ts";
export * from "./framing.ts";
export * from "./messages.ts";
export * from "./schemas.ts";

182
packages/protocol/src/messages.ts
View file

@ -1,182 +0,0 @@
import Type, { type Static } from "typebox";
import {
ModelRefSchema,
PROTOCOL_VERSION,
ProtocolErrorSchema,
ServerSnapshotSchema,
SessionSnapshotSchema,
SessionSummarySchema,
ThinkingLevelSchema,
TranscriptProgressSchema,
} from "./schemas.ts";
const IdSchema = Type.String({ minLength: 1 });
const StrictObject = <const T extends Parameters<typeof Type.Object>[0]>(properties: T) =>
Type.Object(properties, { additionalProperties: false });
const PromptPayloadProperties = {
sessionId: IdSchema,
text: Type.String(),
} as const;
export const ListCommandSchema = StrictObject({ command: Type.Literal("list") });
export const CreateCommandSchema = StrictObject({
command: Type.Literal("create"),
cwd: Type.Optional(Type.String({ minLength: 1 })),
name: Type.Optional(Type.String()),
model: Type.Optional(ModelRefSchema),
thinkingLevel: Type.Optional(ThinkingLevelSchema),
});
export const AttachCommandSchema = StrictObject({ command: Type.Literal("attach"), sessionId: IdSchema });
export const DetachCommandSchema = StrictObject({ command: Type.Literal("detach"), sessionId: IdSchema });
export const PromptCommandSchema = StrictObject({ command: Type.Literal("prompt"), ...PromptPayloadProperties });
export const SteerCommandSchema = StrictObject({ command: Type.Literal("steer"), ...PromptPayloadProperties });
export const AbortCommandSchema = StrictObject({ command: Type.Literal("abort"), sessionId: IdSchema });
export const SetModelCommandSchema = StrictObject({
command: Type.Literal("set_model"),
sessionId: IdSchema,
model: ModelRefSchema,
});
export const SetThinkingCommandSchema = StrictObject({
command: Type.Literal("set_thinking"),
sessionId: IdSchema,
thinkingLevel: ThinkingLevelSchema,
});
export const CommandSchema = Type.Union([
ListCommandSchema,
CreateCommandSchema,
AttachCommandSchema,
DetachCommandSchema,
PromptCommandSchema,
SteerCommandSchema,
AbortCommandSchema,
SetModelCommandSchema,
SetThinkingCommandSchema,
]);
export type Command = Static<typeof CommandSchema>;
export type CommandName = Command["command"];
export const CreateResultSchema = StrictObject({
command: Type.Literal("create"),
session: SessionSnapshotSchema,
});
export const AttachResultSchema = StrictObject({
command: Type.Literal("attach"),
session: SessionSnapshotSchema,
});
export const PromptResultSchema = StrictObject({
command: Type.Literal("prompt"),
session: SessionSnapshotSchema,
});
export const SteerResultSchema = StrictObject({
command: Type.Literal("steer"),
session: SessionSnapshotSchema,
});
export const AbortResultSchema = StrictObject({
command: Type.Literal("abort"),
session: SessionSnapshotSchema,
});
export const SetModelResultSchema = StrictObject({
command: Type.Literal("set_model"),
session: SessionSnapshotSchema,
});
export const SetThinkingResultSchema = StrictObject({
command: Type.Literal("set_thinking"),
session: SessionSnapshotSchema,
});
export const ListResultSchema = StrictObject({
command: Type.Literal("list"),
sessions: Type.Array(SessionSummarySchema),
});
export const DetachResultSchema = StrictObject({
command: Type.Literal("detach"),
sessionId: IdSchema,
});
export const CommandResultSchema = Type.Union([
ListResultSchema,
CreateResultSchema,
AttachResultSchema,
DetachResultSchema,
PromptResultSchema,
SteerResultSchema,
AbortResultSchema,
SetModelResultSchema,
SetThinkingResultSchema,
]);
export type CommandResult = Static<typeof CommandResultSchema>;
export type ResultForCommand<TCommand extends Command> = TCommand["command"] extends "list"
? Static<typeof ListResultSchema>
: TCommand["command"] extends "detach"
? Static<typeof DetachResultSchema>
: Extract<CommandResult, { command: TCommand["command"] }>;
/** Must be the first frame sent by a client. Version is intentionally an integer, not a coercible string. */
export const ClientHelloSchema = StrictObject({
type: Type.Literal("hello"),
version: Type.Integer({ minimum: 0 }),
token: Type.String({ minLength: 1 }),
});
export type ClientHello = Static<typeof ClientHelloSchema>;
export const RequestEnvelopeSchema = StrictObject({
type: Type.Literal("request"),
id: IdSchema,
request: CommandSchema,
});
export type RequestEnvelope = Static<typeof RequestEnvelopeSchema>;
export const ClientMessageSchema = Type.Union([ClientHelloSchema, RequestEnvelopeSchema]);
export type ClientMessage = Static<typeof ClientMessageSchema>;
export const ServerEventSchema = Type.Union([
StrictObject({ type: Type.Literal("server_snapshot"), snapshot: ServerSnapshotSchema }),
StrictObject({ type: Type.Literal("session_snapshot"), snapshot: SessionSnapshotSchema }),
StrictObject({
type: Type.Literal("session_progress"),
sessionId: IdSchema,
progress: TranscriptProgressSchema,
}),
StrictObject({ type: Type.Literal("session_removed"), sessionId: IdSchema }),
]);
export type ServerEvent = Static<typeof ServerEventSchema>;
export const ServerHelloSchema = StrictObject({
type: Type.Literal("hello"),
version: Type.Literal(PROTOCOL_VERSION),
connectionId: IdSchema,
snapshot: ServerSnapshotSchema,
});
export const ServerHelloErrorSchema = StrictObject({
type: Type.Literal("hello_error"),
error: ProtocolErrorSchema,
});
export const ResponseEnvelopeSchema = Type.Union([
StrictObject({
type: Type.Literal("response"),
id: IdSchema,
ok: Type.Literal(true),
result: CommandResultSchema,
}),
StrictObject({
type: Type.Literal("response"),
id: IdSchema,
ok: Type.Literal(false),
error: ProtocolErrorSchema,
}),
]);
export const EventEnvelopeSchema = StrictObject({
type: Type.Literal("event"),
event: ServerEventSchema,
});
export const ServerMessageSchema = Type.Union([
ServerHelloSchema,
ServerHelloErrorSchema,
ResponseEnvelopeSchema,
EventEnvelopeSchema,
]);
export type ServerHello = Static<typeof ServerHelloSchema>;
export type ServerHelloError = Static<typeof ServerHelloErrorSchema>;
export type ResponseEnvelope = Static<typeof ResponseEnvelopeSchema>;
export type EventEnvelope = Static<typeof EventEnvelopeSchema>;
export type ServerMessage = Static<typeof ServerMessageSchema>;

167
packages/protocol/src/schemas.ts
View file

@ -243,3 +243,170 @@ export const ProtocolErrorSchema = StrictObject({
});
export type ProtocolErrorCode = Static<typeof ProtocolErrorCodeSchema>;
export type ProtocolError = Static<typeof ProtocolErrorSchema>;
const PromptPayloadProperties = {
sessionId: IdSchema,
text: Type.String(),
} as const;
export const ListCommandSchema = StrictObject({ command: Type.Literal("list") });
export const CreateCommandSchema = StrictObject({
command: Type.Literal("create"),
cwd: Type.Optional(Type.String({ minLength: 1 })),
name: Type.Optional(Type.String()),
model: Type.Optional(ModelRefSchema),
thinkingLevel: Type.Optional(ThinkingLevelSchema),
});
export const AttachCommandSchema = StrictObject({ command: Type.Literal("attach"), sessionId: IdSchema });
export const DetachCommandSchema = StrictObject({ command: Type.Literal("detach"), sessionId: IdSchema });
export const PromptCommandSchema = StrictObject({ command: Type.Literal("prompt"), ...PromptPayloadProperties });
export const SteerCommandSchema = StrictObject({ command: Type.Literal("steer"), ...PromptPayloadProperties });
export const AbortCommandSchema = StrictObject({ command: Type.Literal("abort"), sessionId: IdSchema });
export const SetModelCommandSchema = StrictObject({
command: Type.Literal("set_model"),
sessionId: IdSchema,
model: ModelRefSchema,
});
export const SetThinkingCommandSchema = StrictObject({
command: Type.Literal("set_thinking"),
sessionId: IdSchema,
thinkingLevel: ThinkingLevelSchema,
});
export const CommandSchema = Type.Union([
ListCommandSchema,
CreateCommandSchema,
AttachCommandSchema,
DetachCommandSchema,
PromptCommandSchema,
SteerCommandSchema,
AbortCommandSchema,
SetModelCommandSchema,
SetThinkingCommandSchema,
]);
export type Command = Static<typeof CommandSchema>;
export type CommandName = Command["command"];
export const CreateResultSchema = StrictObject({
command: Type.Literal("create"),
session: SessionSnapshotSchema,
});
export const AttachResultSchema = StrictObject({
command: Type.Literal("attach"),
session: SessionSnapshotSchema,
});
export const PromptResultSchema = StrictObject({
command: Type.Literal("prompt"),
session: SessionSnapshotSchema,
});
export const SteerResultSchema = StrictObject({
command: Type.Literal("steer"),
session: SessionSnapshotSchema,
});
export const AbortResultSchema = StrictObject({
command: Type.Literal("abort"),
session: SessionSnapshotSchema,
});
export const SetModelResultSchema = StrictObject({
command: Type.Literal("set_model"),
session: SessionSnapshotSchema,
});
export const SetThinkingResultSchema = StrictObject({
command: Type.Literal("set_thinking"),
session: SessionSnapshotSchema,
});
export const ListResultSchema = StrictObject({
command: Type.Literal("list"),
sessions: Type.Array(SessionSummarySchema),
});
export const DetachResultSchema = StrictObject({
command: Type.Literal("detach"),
sessionId: IdSchema,
});
export const CommandResultSchema = Type.Union([
ListResultSchema,
CreateResultSchema,
AttachResultSchema,
DetachResultSchema,
PromptResultSchema,
SteerResultSchema,
AbortResultSchema,
SetModelResultSchema,
SetThinkingResultSchema,
]);
export type CommandResult = Static<typeof CommandResultSchema>;
export type ResultForCommand<TCommand extends Command> = TCommand["command"] extends "list"
? Static<typeof ListResultSchema>
: TCommand["command"] extends "detach"
? Static<typeof DetachResultSchema>
: Extract<CommandResult, { command: TCommand["command"] }>;
/** Must be the first frame sent by a client. Version is intentionally an integer, not a coercible string. */
export const ClientHelloSchema = StrictObject({
type: Type.Literal("hello"),
version: Type.Integer({ minimum: 0 }),
token: Type.String({ minLength: 1 }),
});
export type ClientHello = Static<typeof ClientHelloSchema>;
export const RequestEnvelopeSchema = StrictObject({
type: Type.Literal("request"),
id: IdSchema,
request: CommandSchema,
});
export type RequestEnvelope = Static<typeof RequestEnvelopeSchema>;
export const ClientMessageSchema = Type.Union([ClientHelloSchema, RequestEnvelopeSchema]);
export type ClientMessage = Static<typeof ClientMessageSchema>;
export const ServerEventSchema = Type.Union([
StrictObject({ type: Type.Literal("server_snapshot"), snapshot: ServerSnapshotSchema }),
StrictObject({ type: Type.Literal("session_snapshot"), snapshot: SessionSnapshotSchema }),
StrictObject({
type: Type.Literal("session_progress"),
sessionId: IdSchema,
progress: TranscriptProgressSchema,
}),
StrictObject({ type: Type.Literal("session_removed"), sessionId: IdSchema }),
]);
export type ServerEvent = Static<typeof ServerEventSchema>;
export const ServerHelloSchema = StrictObject({
type: Type.Literal("hello"),
version: Type.Literal(PROTOCOL_VERSION),
connectionId: IdSchema,
snapshot: ServerSnapshotSchema,
});
export const ServerHelloErrorSchema = StrictObject({
type: Type.Literal("hello_error"),
error: ProtocolErrorSchema,
});
export const ResponseEnvelopeSchema = Type.Union([
StrictObject({
type: Type.Literal("response"),
id: IdSchema,
ok: Type.Literal(true),
result: CommandResultSchema,
}),
StrictObject({
type: Type.Literal("response"),
id: IdSchema,
ok: Type.Literal(false),
error: ProtocolErrorSchema,
}),
]);
export const EventEnvelopeSchema = StrictObject({
type: Type.Literal("event"),
event: ServerEventSchema,
});
export const ServerMessageSchema = Type.Union([
ServerHelloSchema,
ServerHelloErrorSchema,
ResponseEnvelopeSchema,
EventEnvelopeSchema,
]);
export type ServerHello = Static<typeof ServerHelloSchema>;
export type ServerHelloError = Static<typeof ServerHelloErrorSchema>;
export type ResponseEnvelope = Static<typeof ResponseEnvelopeSchema>;
export type EventEnvelope = Static<typeof EventEnvelopeSchema>;
export type ServerMessage = Static<typeof ServerMessageSchema>;

2
packages/protocol/test/cbor.test.ts → packages/protocol/test/cbor/cbor.test.ts
View file

@ -6,7 +6,7 @@ import {
DEFAULT_MAX_CBOR_DEPTH,
decodeCbor,
encodeCbor,
} from "../src/index.ts";
} from "../../src/index.ts";
function fromHex(hex: string): Uint8Array {
if (hex.length % 2 !== 0) throw new Error("Hex fixture must contain whole bytes");

13
packages/protocol/test/protocol.test.ts
View file

@ -121,6 +121,19 @@ describe("protocol validation", () => {
expect(parseServerMessage(message)).toEqual(message);
});
test("rejects cyclic protocol values with a protocol validation error", () => {
const details: Record<string, unknown> = {};
details.self = details;
const message = {
type: "response",
id: "request-1",
ok: false,
error: { code: "invalid_request", message: "invalid", details },
};
expect(() => parseServerMessage(message)).toThrow(ProtocolValidationError);
});
test("validation errors do not retain rejected payloads", () => {
let thrown: unknown;
try {