tiny.http.WebSocket

1pub const WebSocket = struct {2    connection: *Connection,3    state: WebSocketState,4    storage: SessionStorage,5    read_length: usize = 0,6    pending_consumed: usize = 0,7    fragment_opcode: ?Opcode,8    fragment_length: usize = 0,9    fragment_returned: bool = false,1011    pub const Limits: type = SessionLimits;12    pub const Capacity: type = SessionCapacity;13    pub const Storage: type = SessionStorage;1415    pub fn init(16        allocator: std.mem.Allocator,17        limits: SessionLimits,18        connection: *Connection,19    ) SessionStorage.InitError!WebSocket {20        return .{21            .connection = connection,22            .state = .open,23            .storage = try SessionStorage.init(allocator, limits),24            .fragment_opcode = null,25        };26    }2728    pub fn activate(self: *WebSocket) void {29        self.storage.activate();30    }3132    pub fn deinit(self: *WebSocket, allocator: std.mem.Allocator) void {33        self.storage.deinit(allocator);34        self.* = undefined;35    }3637    pub fn capacity(self: *const WebSocket) SessionCapacity {38        return self.storage.capacity;39    }4041    pub fn validateClientHandshake(req: *const Request) HandshakeError!WebSocketClientHandshake {42        if (req.method != .GET) return HandshakeError.InvalidMethod;43        if (req.version != .http_1_1) return HandshakeError.InvalidHttpVersion;44        if (!req.isWebSocketUpgrade()) return HandshakeError.InvalidUpgrade;4546        const version = req.headers.get("Sec-WebSocket-Version") orelse return HandshakeError.MissingVersion;47        if (!std.mem.eql(u8, version, "13")) return HandshakeError.UnsupportedVersion;4849        const key = req.getWebSocketKey() orelse return HandshakeError.MissingKey;50        try validateClientKey(key);5152        return .{ .key = key };53    }5455    pub fn validateClientKey(key: []const u8) HandshakeError!void {56        const decoded_len = std.base64.standard.Decoder.calcSizeForSlice(key) catch return HandshakeError.InvalidKey;57        if (decoded_len != 16) return HandshakeError.InvalidKey;5859        var decoded: [16]u8 = undefined;60        std.base64.standard.Decoder.decode(&decoded, key) catch return HandshakeError.InvalidKey;61    }6263    pub fn handshake(key: []const u8) HandshakeError![28]u8 {64        try validateClientKey(key);65        return accept_key(key);66    }6768    pub fn sendText(self: *WebSocket, data: []const u8) !void {69        try self.sendFrame(.{70            .fin = true,71            .opcode = .text,72            .mask = null,73            .payload = data,74        });75    }7677    pub fn sendBinary(self: *WebSocket, data: []const u8) !void {78        try self.sendFrame(.{79            .fin = true,80            .opcode = .binary,81            .mask = null,82            .payload = data,83        });84    }8586    pub fn sendPing(self: *WebSocket, data: []const u8) !void {87        try self.sendFrame(.{88            .fin = true,89            .opcode = .ping,90            .mask = null,91            .payload = data,92        });93    }9495    pub fn sendPong(self: *WebSocket, data: []const u8) !void {96        try self.sendFrame(.{97            .fin = true,98            .opcode = .pong,99            .mask = null,100            .payload = data,101        });102    }103104    pub fn close(self: *WebSocket, code: u16, reason: []const u8) !void {105        if (self.state == .closed) return;106        const payload_len = std.math.add(usize, 2, reason.len) catch {107            return FrameError.InvalidFrame;108        };109        if (payload_len > maximum_control_payload_bytes) return FrameError.InvalidFrame;110111        var payload: [maximum_control_payload_bytes]u8 = undefined;112        std.mem.writeInt(u16, payload[0..2], code, .big);113        @memcpy(payload[2..][0..reason.len], reason);114115        try self.sendFrame(.{116            .fin = true,117            .opcode = .close,118            .mask = null,119            .payload = payload[0..payload_len],120        });121122        self.state = .closing;123    }124125    pub fn receive(self: *WebSocket) !?WebSocketMessage {126        std.debug.assert(self.storage.phase == .steady);127        self.releaseBorrowedMessage();128        while (self.state != .closed) {129            const buffered_length = self.read_length;130            if (try self.processBufferedFrame()) |message| return message;131            if (self.state == .closed) return null;132            if (self.read_length != buffered_length) continue;133            try self.readFrameBytes();134        }135        return null;136    }137138    fn processBufferedFrame(self: *WebSocket) !?WebSocketMessage {139        if (self.read_length == 0) return null;140        const frame_region = self.frameRegion();141        const parsed = Frame.parseServer(142            frame_region[0..self.read_length],143            self.storage.capacity.frame_payload_bytes,144        ) catch |err| switch (err) {145            error.IncompleteFrame => return null,146            error.MaskRequired, error.InvalidOpcode, error.InvalidFrame => {147                self.rejectProtocol();148                return null;149            },150            error.PayloadTooLarge => {151                self.rejectMessageTooBig();152                return null;153            },154            else => return err,155        };156        if (parsed.frame.opcode.isControl()) {157            try self.processControlFrame(parsed.frame, parsed.consumed);158            return null;159        }160        return self.processDataFrame(parsed.frame, parsed.consumed);161    }162163    fn processControlFrame(self: *WebSocket, frame: Frame, consumed: usize) !void {164        defer self.consumeFrame(consumed);165        switch (frame.opcode) {166            .ping => try self.sendPong(frame.payload),167            .pong => {},168            .close => {169                if (self.state != .closing) self.replyToClose(frame.payload);170                self.state = .closed;171                self.connection.markClosing();172            },173            else => unreachable,174        }175    }176177    fn replyToClose(self: *WebSocket, payload: []const u8) void {178        if (payload.len >= 2) {179            const code = std.mem.readInt(u16, payload[0..2], .big);180            self.close(code, payload[2..]) catch |err| {181                log.warn("failed to send close response: {s}", .{@errorName(err)});182            };183        } else {184            self.close(WebSocketCloseCode.normal, "") catch |err| {185                log.warn("failed to send close response: {s}", .{@errorName(err)});186            };187        }188    }189190    fn processDataFrame(self: *WebSocket, frame: Frame, consumed: usize) ?WebSocketMessage {191        if (frame.fin) return self.finishDataFrame(frame, consumed);192        if (frame.opcode == .continuation) {193            if (self.fragment_opcode == null) {194                self.rejectProtocol();195                return null;196            }197        } else if (self.fragment_opcode == null) {198            self.fragment_opcode = frame.opcode;199            self.fragment_length = 0;200        } else {201            self.rejectProtocol();202            return null;203        }204        if (!self.retainFragment(frame.payload)) return null;205        self.consumeFrame(consumed);206        return null;207    }208209    fn finishDataFrame(self: *WebSocket, frame: Frame, consumed: usize) ?WebSocketMessage {210        if (self.fragment_opcode) |opcode| {211            if (frame.opcode != .continuation) {212                self.rejectProtocol();213                return null;214            }215            if (!self.retainFragment(frame.payload)) return null;216            self.consumeFrame(consumed);217            self.fragment_opcode = null;218            self.fragment_returned = true;219            return .{220                .opcode = opcode,221                .payload = self.messageRegion()[0..self.fragment_length],222            };223        }224        if (frame.opcode == .continuation) {225            self.rejectProtocol();226            return null;227        }228        if (frame.payload.len > self.storage.capacity.message_payload_bytes) {229            self.rejectMessageTooBig();230            return null;231        }232        self.pending_consumed = consumed;233        return .{ .opcode = frame.opcode, .payload = frame.payload };234    }235236    fn retainFragment(self: *WebSocket, payload: []const u8) bool {237        if (!fragmentLengthAllowed(238            self.fragment_length,239            payload.len,240            self.storage.capacity.message_payload_bytes,241        )) {242            self.rejectMessageTooBig();243            return false;244        }245        self.appendFragment(payload);246        return true;247    }248249    fn readFrameBytes(self: *WebSocket) !void {250        const frame_region = self.frameRegion();251        if (self.read_length == frame_region.len) {252            self.rejectMessageTooBig();253            return;254        }255        const read = self.connection.read(256            frame_region[self.read_length..],257        ) catch |err| switch (err) {258            error.WouldBlock => return error.WouldBlock,259            error.ConnectionClosed => {260                self.state = .closed;261                self.connection.markClosing();262                return;263            },264            else => return err,265        };266        if (read == 0) {267            self.state = .closed;268            self.connection.markClosing();269            return;270        }271        self.read_length += read;272    }273274    fn sendFrame(self: *WebSocket, frame: Frame) !void {275        std.debug.assert(self.storage.phase == .steady);276        std.debug.assert(frame.mask == null);277        var header: [maximum_frame_header_bytes]u8 = undefined;278        const encoded = try frame.writeHeader(&header, std.math.maxInt(usize));279        try self.connection.write(encoded);280        if (frame.payload.len != 0) try self.connection.write(frame.payload);281    }282283    fn frameRegion(self: *WebSocket) []u8 {284        return self.storage.frame(self.storage.capacity.frame_payload_bytes) catch unreachable;285    }286287    fn messageRegion(self: *WebSocket) []u8 {288        return self.storage.message(self.storage.capacity.message_payload_bytes) catch unreachable;289    }290291    fn appendFragment(self: *WebSocket, payload: []const u8) void {292        const message = self.messageRegion();293        @memcpy(message[self.fragment_length..][0..payload.len], payload);294        self.fragment_length += payload.len;295    }296297    fn consumeFrame(self: *WebSocket, consumed: usize) void {298        std.debug.assert(consumed <= self.read_length);299        const remaining = self.read_length - consumed;300        const frame = self.frameRegion();301        std.mem.copyForwards(u8, frame[0..remaining], frame[consumed..self.read_length]);302        self.read_length = remaining;303    }304305    fn releaseBorrowedMessage(self: *WebSocket) void {306        if (self.pending_consumed != 0) {307            self.consumeFrame(self.pending_consumed);308            self.pending_consumed = 0;309        }310        if (self.fragment_returned) {311            self.fragment_length = 0;312            self.fragment_returned = false;313        }314    }315316    fn rejectProtocol(self: *WebSocket) void {317        self.close(WebSocketCloseCode.protocol_error, "Invalid WebSocket frame") catch |err| {318            log.warn("failed to send protocol error close: {s}", .{@errorName(err)});319        };320        self.state = .closed;321        self.connection.markClosing();322    }323324    fn rejectMessageTooBig(self: *WebSocket) void {325        self.close(WebSocketCloseCode.message_too_big, "Message too big") catch |close_err| {326            log.warn("failed to send message-too-big close: {s}", .{@errorName(close_err)});327        };328        self.state = .closed;329        self.connection.markClosing();330    }331};

1pub const WebSocket = websocket.WebSocket;