tiny.http.ClientWebSocket

1pub const ClientWebSocket = struct {2    allocator: Allocator,3    connection: PersistentConnection,4    read_storage: []u8,5    write_storage: []u8,6    frame_payload_bytes: usize,7    read_length: usize = 0,8    pending_consumed: usize = 0,9    open: bool = true,1011    fn init(12        allocator: Allocator,13        connection: PersistentConnection,14        scratch: WebSocketScratch,15    ) !ClientWebSocket {16        const expected = try alloc_phase.capacity.add(17            usize,18            client_websocket.frame_header_bytes,19            scratch.frame_payload_bytes,20        );21        if (scratch.read.len != expected or22            scratch.write.len != expected)23        {24            return error.InvalidWebSocketScratch;25        }26        return .{27            .allocator = allocator,28            .connection = connection,29            .read_storage = scratch.read,30            .write_storage = scratch.write,31            .frame_payload_bytes = scratch.frame_payload_bytes,32        };33    }3435    pub fn deinit(self: *ClientWebSocket) void {36        destroyConnection(37            self.allocator,38            self.connection,39            true,40        );41        std.crypto.secureZero(u8, self.read_storage);42        std.crypto.secureZero(u8, self.write_storage);43        self.* = undefined;44    }4546    pub fn pollReadable(47        self: *ClientWebSocket,48        timeout_ms: i32,49    ) !bool {50        if (timeout_ms < 0) return error.InvalidTimeout;51        if (self.read_length > self.pending_consumed or52            connectionBuffered(self.connection))53        {54            return true;55        }56        return try sys.pollReadable(57            connectionSocket(self.connection),58            timeout_ms,59        );60    }6162    pub fn sendBinary(63        self: *ClientWebSocket,64        payload: []const u8,65    ) !void {66        try self.sendFrame(.binary, payload);67    }6869    pub fn receive(70        self: *ClientWebSocket,71    ) !?WebSocketMessage {72        self.releaseBorrowed();73        while (self.open) {74            if (self.read_length == 0) {75                try self.readFrame();76            }77            const parsed = http.Frame.parse(78                self.wire()[0..self.read_length],79                self.frame_payload_bytes,80            ) catch |err| return err;81            if (parsed.frame.mask != null) {82                return error.MaskedServerFrame;83            }84            if (!parsed.frame.fin or85                parsed.frame.opcode == .continuation)86            {87                return error.FragmentedServerMessage;88            }89            switch (parsed.frame.opcode) {90                .binary, .text => {91                    self.pending_consumed = parsed.consumed;92                    return .{93                        .opcode = parsed.frame.opcode,94                        .payload = parsed.frame.payload,95                    };96                },97                .ping => {98                    try self.sendFrame(99                        .pong,100                        parsed.frame.payload,101                    );102                    self.consume(parsed.consumed);103                },104                .pong => self.consume(parsed.consumed),105                .close => {106                    self.sendFrame(107                        .close,108                        parsed.frame.payload,109                    ) catch {};110                    self.consume(parsed.consumed);111                    self.open = false;112                    return null;113                },114                .continuation => unreachable,115            }116        }117        return null;118    }119120    pub fn close(self: *ClientWebSocket) !void {121        if (!self.open) return;122        try self.sendFrame(.close, &.{});123        self.open = false;124    }125126    fn sendFrame(127        self: *ClientWebSocket,128        opcode: http.Opcode,129        payload: []const u8,130    ) !void {131        if (!self.open and opcode != .close) {132            return error.ConnectionClosed;133        }134        var mask: [4]u8 = undefined;135        try sys_root.random.secureBytes(&mask);136        const encoded = try (http.Frame{137            .fin = true,138            .opcode = opcode,139            .mask = mask,140            .payload = payload,141        }).serializeInto(142            self.output(),143            self.frame_payload_bytes,144        );145        defer std.crypto.secureZero(u8, encoded);146        const writer = connectionWriter(147            self.connection,148        );149        try writer.writeAll(encoded);150        try flushConnection(self.connection);151    }152153    fn readFrame(self: *ClientWebSocket) !void {154        const reader = connectionReader(self.connection);155        const target = self.wire();156        try readClientWebSocketExact(157            reader,158            target[0..2],159        );160        if ((target[1] & 0x80) != 0) {161            return error.MaskedServerFrame;162        }163        const short_length = target[1] & 0x7f;164        var header_bytes: usize = 2;165        const payload_bytes: usize = switch (short_length) {166            0...125 => short_length,167            126 => length: {168                try readClientWebSocketExact(169                    reader,170                    target[2..4],171                );172                header_bytes = 4;173                break :length std.mem.readInt(174                    u16,175                    target[2..4],176                    .big,177                );178            },179            127 => length: {180                try readClientWebSocketExact(181                    reader,182                    target[2..10],183                );184                header_bytes = 10;185                const wide = std.mem.readInt(186                    u64,187                    target[2..10],188                    .big,189                );190                if (wide > self.frame_payload_bytes or191                    wide > std.math.maxInt(usize))192                {193                    return error.PayloadTooLarge;194                }195                break :length @intCast(wide);196            },197            else => unreachable,198        };199        if (payload_bytes > self.frame_payload_bytes) {200            return error.PayloadTooLarge;201        }202        const total = std.math.add(203            usize,204            header_bytes,205            payload_bytes,206        ) catch return error.PayloadTooLarge;207        if (total > target.len) {208            return error.FrameCapacity;209        }210        try readClientWebSocketExact(211            reader,212            target[header_bytes..total],213        );214        self.read_length = total;215    }216217    fn releaseBorrowed(self: *ClientWebSocket) void {218        if (self.pending_consumed == 0) return;219        self.consume(self.pending_consumed);220        self.pending_consumed = 0;221    }222223    fn consume(224        self: *ClientWebSocket,225        consumed: usize,226    ) void {227        const remaining = self.read_length - consumed;228        const previous_length = self.read_length;229        std.mem.copyForwards(230            u8,231            self.wire()[0..remaining],232            self.wire()[consumed..self.read_length],233        );234        std.crypto.secureZero(235            u8,236            self.wire()[remaining..previous_length],237        );238        self.read_length = remaining;239    }240241    fn wire(self: *ClientWebSocket) []u8 {242        return self.read_storage;243    }244245    fn output(self: *ClientWebSocket) []u8 {246        return self.write_storage;247    }248};

1pub const ClientWebSocket = client.ClientWebSocket;