src/client/buffered_wire.zig

Ref: main worktree-web-client-spec   Size: 7.9 KiB   History

const std = @import("std");
const client = @import("client.zig");
const proto = @import("term").protocol;

/// Opaque transport bytes held by a forwarding connection, across fragmented
/// frames and bursts. Frame parsing remains forwarding-only in this wrapper.
pub const transport_queue_max: usize = 1024 * 1024;
// Incremental reads and queued writes keep partial stream frames and a
// peer which stops reading from blocking the mailbox or stop(). QUIC keeps
// its existing framing and outgoing queue in Link.
pub const Wire = struct {
    alloc: std.mem.Allocator,
    tr: *client.Transport,
    input: std.ArrayList(u8) = .empty,
    output: std.ArrayList(u8) = .empty,

    pub fn init(alloc: std.mem.Allocator, tr: *client.Transport) !Wire {
        switch (tr.link) {
            .fd => |fd| try nonblocking(fd),
            .pipe => |p| {
                try nonblocking(p.r);
                try nonblocking(p.w);
            },
            .quic => {},
        }
        return .{ .alloc = alloc, .tr = tr };
    }
    pub fn deinit(self: *Wire) void {
        self.input.deinit(self.alloc);
        self.output.deinit(self.alloc);
    }
    pub fn writeFd(self: *Wire) std.posix.fd_t {
        return switch (self.tr.link) {
            .fd => |fd| fd,
            .pipe => |p| p.w,
            .quic => self.tr.pollFd(),
        };
    }
    /// Stream backlog becomes writable through its output fd. QUIC backlog
    /// advances only when service processes packets or a transport timer, so
    /// polling its UDP socket for OUT would spin while stream credit is full.
    pub fn pendingWriteFd(self: *Wire) ?std.posix.fd_t {
        if (!self.pending()) return null;
        return switch (self.tr.link) {
            .fd => |fd| fd,
            .pipe => |p| p.w,
            .quic => null,
        };
    }
    pub fn pending(self: *Wire) bool {
        return if (self.tr.link == .quic) self.tr.link.quic.qout.items.len > 0 else self.output.items.len > 0;
    }
    pub fn pendingBytes(self: *const Wire) usize {
        return if (self.tr.link == .quic) self.tr.link.quic.qout.items.len else self.output.items.len;
    }
    pub fn send(self: *Wire, kind: proto.MsgType, payload: []const u8) !void {
        if (self.tr.link == .quic) return self.tr.writeFrame(kind, payload);
        try proto.appendFrame(&self.output, self.alloc, kind, payload);
        try self.flush();
    }
    pub fn flush(self: *Wire) !void {
        if (self.tr.link == .quic) return self.tr.flushQuic();
        if (self.output.items.len == 0) return;
        const n = std.posix.write(self.writeFd(), self.output.items) catch |err| switch (err) {
            error.WouldBlock => return,
            else => return err,
        };
        self.output.replaceRangeAssumeCapacity(0, n, &.{});
    }
    pub fn read(self: *Wire) !client.Incoming {
        return self.readLimited(proto.frame_header_len + proto.max_payload);
    }
    /// Forwarding peers accept much smaller frames than the terminal wire's
    /// paste-sized maximum. Bound the partial frame too, rather than letting a
    /// claimed terminal-sized payload quietly become a forwarding queue.
    pub fn readLimited(self: *Wire, max_buffered: usize) !client.Incoming {
        if (self.tr.link == .quic) {
            if (self.tr.link.quic.cl.in.items.len > max_buffered) return error.FrameTooLarge;
            return self.tr.readFrame(self.alloc);
        }
        var need: usize = proto.frame_header_len;
        if (self.input.items.len >= proto.frame_header_len) {
            const len = std.mem.readInt(u32, self.input.items[1..5], .little);
            if (len > proto.max_payload) return error.FrameTooLarge;
            need += len;
        }
        if (need > max_buffered) return error.FrameTooLarge;
        if (self.input.items.len >= need) {
            if (try proto.takeFrame(self.alloc, &self.input)) |frame| return .{ .frame = frame };
            unreachable;
        }
        if (self.input.items.len >= max_buffered) return error.FrameTooLarge;
        var buf: [64 * 1024]u8 = undefined;
        const room = max_buffered - self.input.items.len;
        const n = std.posix.read(self.tr.pollFd(), buf[0..@min(buf.len, @min(need - self.input.items.len, room))]) catch |err| switch (err) {
            error.WouldBlock => return .incomplete,
            else => return err,
        };
        if (n == 0) return .closed;
        try self.input.appendSlice(self.alloc, buf[0..n]);
        if (self.input.items.len >= proto.frame_header_len) {
            const staged_len = std.mem.readInt(u32, self.input.items[1..5], .little);
            if (staged_len > proto.max_payload or proto.frame_header_len + staged_len > max_buffered) return error.FrameTooLarge;
        }
        if (try proto.takeFrame(self.alloc, &self.input)) |frame| return .{ .frame = frame };
        return .incomplete;
    }
};

fn nonblocking(fd: std.posix.fd_t) !void {
    const flags = try std.posix.fcntl(fd, std.posix.F.GETFL, 0);
    const bits: u32 = @bitCast(std.posix.O{ .NONBLOCK = true });
    _ = try std.posix.fcntl(fd, std.posix.F.SETFL, flags | bits);
}

test "limited wire rejects a claimed forwarding queue before reading its body" {
    var pair: [2]std.posix.fd_t = undefined;
    try std.testing.expectEqual(@as(c_int, 0), std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair));
    defer std.posix.close(pair[1]);
    var tr: client.Transport = .{ .link = .{ .fd = pair[0] } };
    defer tr.close();
    var wire = try Wire.init(std.testing.allocator, &tr);
    defer wire.deinit();

    const header = proto.encodeHeader(.forward_data, 1024);
    try proto.writeAllFd(pair[1], &header);
    try std.testing.expectError(error.FrameTooLarge, wire.readLimited(512));
}

test "forwarding budget admits a maximum legal forwarding frame" {
    var pair: [2]std.posix.fd_t = undefined;
    try std.testing.expectEqual(@as(c_int, 0), std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &pair));
    defer std.posix.close(pair[1]);
    var tr: client.Transport = .{ .link = .{ .fd = pair[0] } };
    defer tr.close();
    var wire = try Wire.init(std.testing.allocator, &tr);
    defer wire.deinit();

    var payload: [proto.forward_id_len + proto.forward_data_max]u8 = undefined;
    try proto.writeFrame(pair[1], .forward_data, &payload);
    const frame = while (true) switch (try wire.readLimited(transport_queue_max)) {
        .frame => |frame| break frame,
        .incomplete => continue,
        .closed => return error.UnexpectedEof,
    };
    defer frame.deinit(std.testing.allocator);
    try std.testing.expectEqual(proto.MsgType.forward_data, frame.type);
    try std.testing.expectEqual(payload.len, frame.payload.len);
    try std.testing.expect(proto.frame_header_len + frame.payload.len <= transport_queue_max);
}

test "only stream backlog requests write readiness" {
    const alloc = std.testing.allocator;
    const pair = blk: {
        var fds: [2]std.posix.fd_t = undefined;
        try std.testing.expectEqual(@as(c_int, 0), std.c.socketpair(std.posix.AF.UNIX, std.posix.SOCK.STREAM, 0, &fds));
        break :blk fds;
    };
    defer std.posix.close(pair[0]);
    defer std.posix.close(pair[1]);

    var stream_tr: client.Transport = .{ .link = .{ .fd = pair[0] } };
    var stream_wire: Wire = .{ .alloc = alloc, .tr = &stream_tr };
    defer stream_wire.deinit();
    try std.testing.expectEqual(@as(?std.posix.fd_t, null), stream_wire.pendingWriteFd());
    try stream_wire.output.append(alloc, 1);
    try std.testing.expectEqual(@as(?std.posix.fd_t, pair[0]), stream_wire.pendingWriteFd());

    var quic_tr: client.Transport = .{ .link = .{ .quic = .{ .cl = undefined, .alloc = alloc } } };
    defer quic_tr.link.quic.qout.deinit(alloc);
    var quic_wire: Wire = .{ .alloc = alloc, .tr = &quic_tr };
    defer quic_wire.deinit();
    try quic_tr.link.quic.qout.append(alloc, 1);
    try std.testing.expect(quic_wire.pending());
    try std.testing.expectEqual(@as(?std.posix.fd_t, null), quic_wire.pendingWriteFd());
}