544 lines · 17.8 KB · Zig
RawDownload
const std = @import("std");const contracts = @import("contracts.zig"); const Allocator = std.mem.Allocator; pub const header_len: usize = 28;const magic = "FXTH"; pub const DecodeError = error{    TruncatedFrame,    InvalidFrameMagic,    InvalidFrameKind,    InvalidFrameSubject,    InvalidHostFrame,    HostFrameTooLarge,    InvalidPayload,} || contracts.HostMessageValidationError || Allocator.Error; pub const Header = struct {    envelope: contracts.MessageEnvelope,     pub fn decode(bytes: []const u8) DecodeError!Header {        if (bytes.len < header_len) return error.TruncatedFrame;        if (!std.mem.eql(u8, bytes[0..magic.len], magic)) {            return error.InvalidFrameMagic;        }         const revision = std.mem.readInt(u16, bytes[4..6], .little);        const kind = try decodeKind(bytes[6], bytes[7]);        const required_capabilities = std.mem.readInt(u64, bytes[8..16], .little);        const correlation_value = std.mem.readInt(u64, bytes[16..24], .little);        const payload_len = std.mem.readInt(u32, bytes[24..28], .little);        const correlation_id: ?contracts.CorrelationId =            if (correlation_value == 0) null else .{ .value = correlation_value };        const envelope = contracts.MessageEnvelope{            .revision = revision,            .kind = kind,            .required_capabilities = required_capabilities,            .correlation_id = correlation_id,            .payload_len = payload_len,        };        try envelope.validate(payload_len);        return .{ .envelope = envelope };    }     pub fn encode(self: Header) DecodeError![header_len]u8 {        try self.envelope.validate(self.envelope.payload_len);        const wire_kind = encodeKind(self.envelope.kind);        var bytes: [header_len]u8 = @splat(0);        @memcpy(bytes[0..magic.len], magic);        std.mem.writeInt(u16, bytes[4..6], self.envelope.revision, .little);        bytes[6] = wire_kind.kind;        bytes[7] = wire_kind.subject;        std.mem.writeInt(            u64,            bytes[8..16],            self.envelope.required_capabilities,            .little,        );        std.mem.writeInt(            u64,            bytes[16..24],            if (self.envelope.correlation_id) |id| id.value else 0,            .little,        );        std.mem.writeInt(u32, bytes[24..28], self.envelope.payload_len, .little);        return bytes;    }}; pub const EncodedFrame = struct {    bytes: []u8,     pub fn deinit(self: *EncodedFrame, alloc: Allocator) void {        alloc.free(self.bytes);        self.* = undefined;    }}; pub const DecodedFrame = struct {    envelope: contracts.MessageEnvelope,    parsed: std.json.Parsed(contracts.MessagePayload),     pub fn message(self: *const DecodedFrame) contracts.HostMessage {        return .{            .envelope = self.envelope,            .payload = self.parsed.value,        };    }     pub fn deinit(self: *DecodedFrame) void {        self.parsed.deinit();        self.* = undefined;    }}; pub fn encodeFrame(    alloc: Allocator,    revision: u16,    required_capabilities: u64,    correlation_id: ?contracts.CorrelationId,    payload: contracts.MessagePayload,) DecodeError!EncodedFrame {    var payload_out: std.Io.Writer.Allocating = .init(alloc);    defer payload_out.deinit();    try stringifyPayload(revision, payload, &payload_out.writer);    const payload_bytes = payload_out.written();    if (payload_bytes.len > contracts.max_host_frame_bytes) {        return error.HostFrameTooLarge;    }    const payload_len = std.math.cast(u32, payload_bytes.len) orelse        return error.HostFrameTooLarge;    const kind = kindForPayload(payload);    const envelope = contracts.MessageEnvelope{        .revision = revision,        .kind = kind,        .required_capabilities = required_capabilities,        .correlation_id = correlation_id,        .payload_len = payload_len,    };    try (contracts.HostMessage{        .envelope = envelope,        .payload = payload,    }).validate(payload_bytes.len);     const total_len = std.math.add(usize, header_len, payload_bytes.len) catch        return error.HostFrameTooLarge;    var bytes = try alloc.alloc(u8, total_len);    errdefer alloc.free(bytes);    const header = try (Header{ .envelope = envelope }).encode();    @memcpy(bytes[0..header_len], &header);    @memcpy(bytes[header_len..], payload_bytes);    return .{ .bytes = bytes };} fn stringifyPayload(    revision: u16,    payload: contracts.MessagePayload,    writer: *std.Io.Writer,) DecodeError!void {    if (revision == contracts.previous_protocol_revision) {        const projected = projectPayloadV4(payload);        std.json.Stringify.value(projected, .{}, writer) catch            return error.OutOfMemory;        return;    }    std.json.Stringify.value(payload, .{}, writer) catch        return error.OutOfMemory;} fn projectPayloadV4(payload: contracts.MessagePayload) contracts.MessagePayload {    return switch (payload) {        .hello => |hello| .{ .hello = hello },        .request => |request| .{ .request = request },        .response => |response| .{ .response = projectResultV4(response) },        .cancel => .cancel,    };} fn projectResultV4(result: contracts.Result) contracts.Result {    const inspect = switch (result) {        .success => |success| switch (success) {            .inspect => |value| value,            else => return result,        },        .failure => return result,    };    for (inspect.monitors) |monitor| {        if (monitor.state == .degraded) return .{ .failure = .{            .action = .inspect,            .code = .protocol_incompatible,            .session_id = inspect.session.session_id,        } };    }    return result;} pub fn projectResultForCapabilities(    result: contracts.Result,    capabilities: u64,) contracts.Result {    const failure = switch (result) {        .success => return result,        .failure => |value| value,    };    if (failure.code != .path_outside_workspace or        capabilities & contracts.protocol_capability_path_outside_workspace_error != 0)    {        return result;    }    var compatible = failure;    compatible.code = .invalid_request;    return .{ .failure = compatible };} pub fn decodeFrame(alloc: Allocator, bytes: []const u8) DecodeError!DecodedFrame {    const header = try Header.decode(bytes);    const payload_len: usize = header.envelope.payload_len;    const expected_len = std.math.add(usize, header_len, payload_len) catch        return error.HostFrameTooLarge;    if (bytes.len < expected_len) return error.TruncatedFrame;    if (bytes.len != expected_len) return error.InvalidHostFrame;     var parsed = std.json.parseFromSlice(        contracts.MessagePayload,        alloc,        bytes[header_len..],        .{ .allocate = .alloc_always },    ) catch |err| switch (err) {        error.OutOfMemory => return error.OutOfMemory,        else => return error.InvalidPayload,    };    errdefer parsed.deinit();    try (contracts.HostMessage{        .envelope = header.envelope,        .payload = parsed.value,    }).validate(payload_len);    return .{        .envelope = header.envelope,        .parsed = parsed,    };} pub fn readFrame(    alloc: Allocator,    reader: *std.Io.Reader,) (DecodeError || std.Io.Reader.Error)!DecodedFrame {    var header_bytes: [header_len]u8 = undefined;    reader.readSliceAll(&header_bytes) catch |err| switch (err) {        error.EndOfStream => return error.TruncatedFrame,        else => return err,    };    const header = try Header.decode(&header_bytes);    const payload_len: usize = header.envelope.payload_len;    const total_len = std.math.add(usize, header_len, payload_len) catch        return error.HostFrameTooLarge;    var bytes = try alloc.alloc(u8, total_len);    defer alloc.free(bytes);    @memcpy(bytes[0..header_len], &header_bytes);    reader.readSliceAll(bytes[header_len..]) catch |err| switch (err) {        error.EndOfStream => return error.TruncatedFrame,        else => return err,    };    return decodeFrame(alloc, bytes);} pub fn writeFrame(    writer: *std.Io.Writer,    frame: EncodedFrame,) std.Io.Writer.Error!void {    try writer.writeAll(frame.bytes);    try writer.flush();} fn kindForPayload(payload: contracts.MessagePayload) contracts.MessageKind {    return switch (payload) {        .hello => .hello,        .request => |request| .{ .request = request.action() },        .response => |response| .{ .response = response.action() },        .cancel => .cancel,    };} const WireKind = struct {    kind: u8,    subject: u8,}; fn encodeKind(kind: contracts.MessageKind) WireKind {    return switch (kind) {        .hello => .{ .kind = 0, .subject = 0 },        .request => |action| .{ .kind = 1, .subject = @intFromEnum(action) },        .response => |action| .{ .kind = 2, .subject = @intFromEnum(action) },        .cancel => .{ .kind = 4, .subject = 0 },    };} fn decodeKind(kind: u8, subject: u8) DecodeError!contracts.MessageKind {    return switch (kind) {        0 => if (subject == 0) .hello else error.InvalidFrameSubject,        1 => if (subject <= @intFromEnum(contracts.Action.close))            .{ .request = @enumFromInt(subject) }        else            error.InvalidFrameSubject,        2 => if (subject <= @intFromEnum(contracts.Action.close))            .{ .response = @enumFromInt(subject) }        else            error.InvalidFrameSubject,        3 => error.InvalidFrameKind,        4 => if (subject == 0) .cancel else error.InvalidFrameSubject,        else => error.InvalidFrameKind,    };} fn testRequestFrame(alloc: Allocator) !void {    var encoded = try encodeFrame(        alloc,        contracts.current_protocol_revision,        0,        .{ .value = 41 },        .{ .request = .{ .screen = .{ .session_id = "terminal-1" } } },    );    defer encoded.deinit(alloc);    var decoded = try decodeFrame(alloc, encoded.bytes);    defer decoded.deinit();    const message = decoded.message();    try std.testing.expectEqual(@as(u64, 41), message.envelope.correlation_id.?.value);    try std.testing.expectEqual(contracts.Action.screen, message.payload.request.action());} test "host protocol round trips owned correlated frames" {    try std.testing.checkAllAllocationFailures(        std.testing.allocator,        testRequestFrame,        .{},    );} test "host protocol reads fragmented frames and preserves cancel framing" {    var encoded = try encodeFrame(        std.testing.allocator,        contracts.current_protocol_revision,        0,        .{ .value = 7 },        .cancel,    );    defer encoded.deinit(std.testing.allocator);    var reader_buffer: [64]u8 = undefined;    var fragmented: std.testing.Reader = .init(&reader_buffer, &.{        .{ .buffer = encoded.bytes },    });    fragmented.artificial_limit = .limited(1);    var decoded = try readFrame(        std.testing.allocator,        &fragmented.interface,    );    defer decoded.deinit();    try std.testing.expectEqual(contracts.MessageKind.cancel, decoded.message().envelope.kind);    try std.testing.expectEqual(@as(u8, 4), encoded.bytes[6]);} test "host protocol reserves removed event kind" {    var encoded = try encodeFrame(        std.testing.allocator,        contracts.current_protocol_revision,        0,        .{ .value = 7 },        .cancel,    );    defer encoded.deinit(std.testing.allocator);    encoded.bytes[6] = 3;    try std.testing.expectError(        error.InvalidFrameKind,        decodeFrame(std.testing.allocator, encoded.bytes),    );} test "host protocol rejects partial truncated oversized and invalid correlation frames" {    var encoded = try encodeFrame(        std.testing.allocator,        contracts.current_protocol_revision,        0,        .{ .value = 9 },        .{ .request = .{ .screen = .{ .session_id = "terminal-1" } } },    );    defer encoded.deinit(std.testing.allocator);     try std.testing.expectError(        error.TruncatedFrame,        decodeFrame(std.testing.allocator, encoded.bytes[0 .. header_len - 1]),    );    try std.testing.expectError(        error.TruncatedFrame,        decodeFrame(std.testing.allocator, encoded.bytes[0 .. encoded.bytes.len - 1]),    );     var oversized = encoded.bytes[0..header_len].*;    std.mem.writeInt(        u32,        oversized[24..28],        contracts.max_host_frame_bytes + 1,        .little,    );    try std.testing.expectError(        error.HostFrameTooLarge,        Header.decode(&oversized),    );     var invalid_correlation = encoded.bytes[0..header_len].*;    @memset(invalid_correlation[16..24], 0);    try std.testing.expectError(        error.MissingCorrelationId,        Header.decode(&invalid_correlation),    );} test "host protocol rejects required capabilities before payload allocation" {    var encoded = try encodeFrame(        std.testing.allocator,        contracts.current_protocol_revision,        0,        null,        .{ .hello = .{            .range = contracts.local_protocol_range,            .capabilities = contracts.known_protocol_capabilities,        } },    );    defer encoded.deinit(std.testing.allocator);    std.mem.writeInt(u64, encoded.bytes[8..16], 1 << 63, .little);    try std.testing.expectError(        error.UnknownRequiredCapability,        decodeFrame(std.testing.allocator, encoded.bytes),    );} test "host protocol preserves opaque terminal output bytes" {    const output = [_]u8{ 0, 0x80, 0xff };    var encoded = try encodeFrame(        std.testing.allocator,        contracts.current_protocol_revision,        0,        .{ .value = 1 },        .{ .response = .{ .success = .{ .read = .{            .session = .{                .session_id = "terminal-opaque",                .lifecycle = .running,                .attention = .{},                .backend = .native,                .output_cursor = .{ .segment = 1, .offset = output.len },                .screen_recovery = .{ .unavailable = .missing },            },            .output = &output,            .raw_range = .{                .start = .{ .segment = 1, .offset = 0 },                .end = .{ .segment = 1, .offset = output.len },            },        } } } },    );    defer encoded.deinit(std.testing.allocator);    var decoded = try decodeFrame(std.testing.allocator, encoded.bytes);    defer decoded.deinit();    const result = switch (decoded.message().payload) {        .response => |response| response,        else => return error.TestExpectedResponse,    };    const read = switch (result) {        .success => |success| switch (success) {            .read => |value| value,            else => return error.TestExpectedRead,        },        .failure => return error.TestExpectedSuccess,    };    try std.testing.expectEqualSlices(u8, &output, read.output);} test "revision four projection never exposes degraded monitor state" {    var response = try encodeFrame(        std.testing.allocator,        contracts.previous_protocol_revision,        0,        .{ .value = 1 },        .{ .response = .{ .success = .{ .inspect = .{            .session = .{                .session_id = "terminal-1",                .lifecycle = .running,                .attention = .{},                .backend = .tmux,                .output_cursor = .{ .segment = 1, .offset = 0 },                .screen_recovery = .{ .unavailable = .raw_gap },            },            .shell = "/bin/sh",            .cwd = "/workspace",            .monitors = &.{.{                .monitor_id = "monitor-1",                .state = .degraded,            }},        } } } },    );    defer response.deinit(std.testing.allocator);    const response_json = response.bytes[header_len..];    try std.testing.expect(std.mem.find(u8, response_json, "degraded") == null);    try std.testing.expect(std.mem.find(u8, response_json, "protocol_incompatible") != null);} fn expectProjectedErrorCode(    capabilities: u64,    expected: contracts.StructuredErrorCode,) !void {    const projected = projectResultForCapabilities(.{ .failure = .{        .action = .start,        .code = .path_outside_workspace,    } }, capabilities);    var encoded = try encodeFrame(        std.testing.allocator,        contracts.current_protocol_revision,        0,        .{ .value = 1 },        .{ .response = projected },    );    defer encoded.deinit(std.testing.allocator);

Lines 1–500 of 544.