Validate RPC protocol frames with Schema

This commit is contained in:
elpresidank 2026-06-04 07:04:15 -05:00
parent 67b5e0dd5b
commit 5a945af345
3 changed files with 241 additions and 19 deletions

View file

@ -1,9 +1,41 @@
import { Effect, Queue, Scope } from "effect";
import * as S from "effect/Schema";
import * as RpcMessage from "effect/unstable/rpc/RpcMessage";
import * as RpcSerialization from "effect/unstable/rpc/RpcSerialization";
import * as RpcServer from "effect/unstable/rpc/RpcServer";
import * as Socket from "effect/unstable/socket/Socket";
export class RpcProtocolDecodeError extends S.TaggedErrorClass<RpcProtocolDecodeError>()(
"RpcProtocolDecodeError",
{
message: S.String,
cause: S.Unknown,
},
) {}
const HeaderSchema = S.mutable(S.Tuple([S.String, S.String]));
const FromClientEncodedSchema = S.Union([
S.TaggedStruct("Request", {
id: S.String,
tag: S.String,
payload: S.Unknown,
headers: S.Array(HeaderSchema),
traceId: S.optionalKey(S.String),
spanId: S.optionalKey(S.String),
sampled: S.optionalKey(S.Boolean),
}),
S.TaggedStruct("Ack", {
requestId: S.String,
}),
S.TaggedStruct("Interrupt", {
requestId: S.String,
}),
S.TaggedStruct("Ping", {}),
S.TaggedStruct("Eof", {}),
]).pipe(S.toTaggedUnion("_tag"));
const FromClientEncodedMessagesSchema = S.Array(FromClientEncodedSchema);
const decodeFromClientMessages = S.decodeUnknownEffect(FromClientEncodedMessagesSchema);
export const makeSocketRpcProtocol = Effect.gen(function* () {
const serialization = yield* RpcSerialization.RpcSerialization;
const disconnects = yield* Queue.make<number>();
@ -34,18 +66,17 @@ export const makeSocketRpcProtocol = Effect.gen(function* () {
});
const writeRaw = yield* socket.writer;
const encodeDefect = (cause: unknown) =>
Effect.sync(() => parser.encode(RpcMessage.ResponseDefectEncoded(cause))!);
const writeDefect = (cause: unknown) =>
Effect.sync(() => parser.encode(RpcMessage.ResponseDefectEncoded(cause))).pipe(
Effect.flatMap((encoded) => encoded === undefined ? Effect.void : writeRaw(encoded)),
);
const write = (response: RpcMessage.FromServerEncoded) =>
Effect.sync(() => parser.encode(response)).pipe(
Effect.flatMap((encoded) =>
encoded === undefined ? Effect.void : Effect.orDie(writeRaw(encoded)),
),
Effect.catchDefect((cause: unknown) =>
encodeDefect(cause).pipe(
Effect.flatMap((encoded) => Effect.orDie(writeRaw(encoded))),
Effect.orDie,
),
writeDefect(cause).pipe(Effect.orDie),
),
);
@ -53,7 +84,24 @@ export const makeSocketRpcProtocol = Effect.gen(function* () {
clientIds.add(clientId);
yield* socket.runRaw((data) =>
Effect.sync(() => parser.decode(data) as ReadonlyArray<RpcMessage.FromClientEncoded>).pipe(
Effect.try({
try: () => parser.decode(data),
catch: (cause) =>
RpcProtocolDecodeError.make({
message: "Failed to decode RPC socket frame",
cause,
}),
}).pipe(
Effect.flatMap((raw) =>
decodeFromClientMessages(raw).pipe(
Effect.mapError((cause) =>
RpcProtocolDecodeError.make({
message: "RPC socket frame did not contain valid client messages",
cause,
})
),
)
),
Effect.flatMap((decoded) =>
Effect.forEach(decoded, (message) => {
if (message._tag === "Request" && headers !== undefined) {
@ -65,11 +113,8 @@ export const makeSocketRpcProtocol = Effect.gen(function* () {
return writeRequest(clientId, message);
}, { discard: true }),
),
Effect.catchDefect((cause: unknown) =>
encodeDefect(cause).pipe(
Effect.flatMap((encoded) => writeRaw(encoded)),
),
),
Effect.catch((cause) => writeDefect(cause)),
Effect.catchDefect((cause: unknown) => writeDefect(cause)),
)
).pipe(
Effect.catchReason("SocketError", "SocketCloseError", () => Effect.void),