WebSocket

pub const WebSocket = struct {
    key: []const u8,
    input: *Reader,
    output: *Writer,

    pub const Header0 = packed struct(u8) {
        opcode: Opcode,
        rsv3: u1 = 0,
        rsv2: u1 = 0,
        rsv1: u1 = 0,
        fin: bool,
    };

    pub const Header1 = packed struct(u8) {
        payload_len: enum(u7) {
            len16 = 126,
            len64 = 127,
            _,
        },
        mask: bool,
    };

    pub const Opcode = enum(u4) {
        continuation = 0,
        text = 1,
        binary = 2,
        connection_close = 8,
        ping = 9,
        /// "A Pong frame MAY be sent unsolicited. This serves as a unidirectional
        /// heartbeat. A response to an unsolicited Pong frame is not expected."
        pong = 10,
        _,
    };

    pub const ReadSmallTextMessageError = error{
        ConnectionClose,
        UnexpectedOpCode,
        MessageOversize,
        MissingMaskBit,
        ReadFailed,
        EndOfStream,
    };

    pub const SmallMessage = struct {
        /// Can be text, binary, or ping.
        opcode: Opcode,
        data: []u8,
    };

    /// Reads the next message from the WebSocket stream, failing if the
    /// message does not fit into the input buffer. The returned memory points
    /// into the input buffer and is invalidated on the next read.
    pub fn readSmallMessage(ws: *WebSocket) ReadSmallTextMessageError!SmallMessage {
        const in = ws.input;
        while (true) {
            const header = try in.takeArray(2);
            const h0: Header0 = @bitCast(header[0]);
            const h1: Header1 = @bitCast(header[1]);

            switch (h0.opcode) {
                .text, .binary, .pong, .ping => {},
                .connection_close => return error.ConnectionClose,
                .continuation => return error.UnexpectedOpCode,
                _ => return error.UnexpectedOpCode,
            }

            if (!h0.fin) return error.MessageOversize;
            if (!h1.mask) return error.MissingMaskBit;

            const len: usize = switch (h1.payload_len) {
                .len16 => try in.takeInt(u16, .big),
                .len64 => std.math.cast(usize, try in.takeInt(u64, .big)) orelse return error.MessageOversize,
                else => @backingInt(h1.payload_len),
            };
            if (len > in.buffer.len) return error.MessageOversize;
            const mask: [4]u8 = (try in.takeArray(4)).*;
            const payload = try in.take(len);

            // Skip pongs.
            if (h0.opcode == .pong) continue;

            // The last item may contain a partial word of unused data.
            const floored_len = (payload.len / 4) * 4;

            const payload_chunks: [][4]u8 = @ptrCast(payload[0..floored_len]);
            for (payload_chunks) |*chunk| {
                const mask_i: u32 = @bitCast(mask);
                const chunk_i: u32 = @bitCast(chunk.*);
                chunk.* = @bitCast(chunk_i ^ mask_i);
            }
            for (payload[floored_len..], mask[0 .. payload.len - floored_len]) |*leftover, m| {
                leftover.* ^= m;
            }

            return .{
                .opcode = h0.opcode,
                .data = payload,
            };
        }
    }

    pub fn writeMessage(ws: *WebSocket, data: []const u8, op: Opcode) Writer.Error!void {
        var bufs: [1][]const u8 = .{data};
        try writeMessageVecUnflushed(ws, &bufs, op);
        try ws.output.flush();
    }

    pub fn writeMessageUnflushed(ws: *WebSocket, data: []const u8, op: Opcode) Writer.Error!void {
        var bufs: [1][]const u8 = .{data};
        try writeMessageVecUnflushed(ws, &bufs, op);
    }

    pub fn writeMessageVec(ws: *WebSocket, data: [][]const u8, op: Opcode) Writer.Error!void {
        try writeMessageVecUnflushed(ws, data, op);
        try ws.output.flush();
    }

    pub fn writeMessageVecUnflushed(ws: *WebSocket, data: [][]const u8, op: Opcode) Writer.Error!void {
        const total_len = l: {
            var total_len: u64 = 0;
            for (data) |iovec| total_len += iovec.len;
            break :l total_len;
        };
        const out = ws.output;
        try out.writeByte(@bitCast(@as(Header0, .{
            .opcode = op,
            .fin = true,
        })));
        switch (total_len) {
            0...125 => try out.writeByte(@bitCast(@as(Header1, .{
                .payload_len = @fromBackingInt(@intCast(total_len)),
                .mask = false,
            }))),
            126...0xffff => {
                try out.writeByte(@bitCast(@as(Header1, .{
                    .payload_len = .len16,
                    .mask = false,
                })));
                try out.writeInt(u16, @intCast(total_len), .big);
            },
            else => {
                try out.writeByte(@bitCast(@as(Header1, .{
                    .payload_len = .len64,
                    .mask = false,
                })));
                try out.writeInt(u64, total_len, .big);
            },
        }
        try out.writeVecAll(data);
    }

    pub fn flush(ws: *WebSocket) Writer.Error!void {
        try ws.output.flush();
    }
}