karlseguin.tngl.sh
A websocket implementation for zig
54 kB
1903 lines
1const std = @import("std");
2const builtin = @import("builtin");
3const proto = @import("../proto.zig");
4const buffer = @import("../buffer.zig");
5
6const net = std.net;
7const posix = std.posix;
8const Thread = std.Thread;
9const Allocator = std.mem.Allocator;
10const FixedBufferAllocator = std.heap.FixedBufferAllocator;
11
12const log = std.log.scoped(.websocket);
13
14const OpCode = proto.OpCode;
15const Reader = proto.Reader;
16const Message = proto.Message;
17pub const Handshake = @import("handshake.zig").Handshake;
18const FallbackAllocator = @import("fallback_allocator.zig").FallbackAllocator;
19
20const DEFAULT_BUFFER_SIZE = 2048;
21
22const EMPTY_PONG = ([2]u8{ @intFromEnum(OpCode.pong), 0 })[0..];
23// CLOSE, 2 length, code
24const CLOSE_NORMAL = ([_]u8{ @intFromEnum(OpCode.close), 2, 3, 232 })[0..]; // code: 1000
25const CLOSE_PROTOCOL_ERROR = ([_]u8{ @intFromEnum(OpCode.close), 2, 3, 234 })[0..]; //code: 1002
26
27const force_blocking: bool = blk: {
28 const build = @import("build");
29 if (@hasDecl(build, "websocket_blocking")) {
30 break :blk build.websocket_blocking;
31 }
32 break :blk false;
33};
34
35pub fn blockingMode() bool {
36 if (force_blocking) {
37 return true;
38 }
39 return switch (builtin.os.tag) {
40 .linux, .macos, .ios, .tvos, .watchos, .freebsd, .netbsd, .dragonfly, .openbsd => false,
41 else => true,
42 };
43}
44
45pub const Config = struct {
46 port: u16 = 9882,
47 address: []const u8 = "127.0.0.1",
48 unix_path: ?[]const u8 = null,
49
50 worker_count: ?u8 = null,
51
52 max_conn: usize = 16_384,
53 max_message_size: usize = 65_536,
54
55 handshake: Config.Handshake = .{},
56 thread_pool: ThreadPool = .{},
57 buffers: Config.Buffers = .{},
58
59 pub const ThreadPool = struct {
60 count: ?u16 = null,
61 backlog: ?u32 = null,
62 buffer_size: ?usize = null,
63 };
64
65 const Handshake = struct {
66 timeout: u32 = 10,
67 max_size: ?u16 = null,
68 max_headers: ?u16 = null,
69 count: ?u16 = null,
70 };
71
72 const Buffers = struct {
73 size: ?usize = null,
74 pool: ?usize = null,
75 large_size: ?usize = null,
76 large_count: ?u16 = null,
77 };
78
79 fn workerCount(self: *const Config) usize {
80 if (comptime blockingMode()) {
81 return 1;
82 }
83 return self.worker_count orelse 1;
84 }
85};
86
87pub fn Server(comptime H: type) type {
88 return struct {
89 config: Config,
90 allocator: Allocator,
91
92 _state: WorkerState,
93 _signals: []posix.fd_t,
94 _mut: Thread.Mutex,
95 _cond: Thread.Condition,
96
97 const Self = @This();
98
99 pub fn init(allocator: Allocator, config: Config) !Self {
100 if (blockingMode()) {
101 if (config.buffers.pool) |p| {
102 if (p > 1) {
103 log.warn("blockingMode() cannot utilize a small buffer pool, using per-connection buffer instead", .{});
104 }
105 }
106 }
107
108 const signals = try allocator.alloc(posix.fd_t, config.workerCount());
109 errdefer allocator.free(signals);
110
111 var state = try WorkerState.init(allocator, config);
112 errdefer state.deinit();
113
114 return .{
115 ._mut = .{},
116 ._cond = .{},
117 ._state = state,
118 ._signals = signals,
119 .config = config,
120 .allocator = allocator,
121 };
122 }
123
124 pub fn deinit(self: *Self) void {
125 self._state.deinit();
126 self.allocator.free(self._signals);
127 }
128
129 pub fn listenInNewThread(self: *Self, ctx: anytype) !Thread {
130 self._mut.lock();
131 defer self._mut.unlock();
132 const thrd = try Thread.spawn(.{}, Self.listen, .{self, ctx});
133
134 // we don't return until listen() signals us that the server is up
135 self._cond.wait(&self._mut);
136 return thrd;
137 }
138
139 pub fn listen(self: *Self, ctx: anytype) !void {
140 self._mut.lock();
141 errdefer self._mut.unlock();
142
143 const config = &self.config;
144
145 var no_delay = true;
146 const address = blk: {
147 if (config.unix_path) |unix_path| {
148 if (comptime builtin.os.tag == .windows) {
149 return error.UnixPathNotSupported;
150 }
151 no_delay = false;
152 std.fs.deleteFileAbsolute(unix_path) catch {};
153 break :blk try net.Address.initUnix(unix_path);
154 } else {
155 const listen_port = config.port;
156 const listen_address = config.address;
157 break :blk try net.Address.parseIp(listen_address, listen_port);
158 }
159 };
160
161 const socket = blk: {
162 var sock_flags: u32 = posix.SOCK.STREAM | posix.SOCK.CLOEXEC;
163 if (blockingMode() == false) sock_flags |= posix.SOCK.NONBLOCK;
164
165 const socket_proto = if (address.any.family == posix.AF.UNIX) @as(u32, 0) else posix.IPPROTO.TCP;
166 break :blk try posix.socket(address.any.family, sock_flags, socket_proto);
167 };
168
169 if (no_delay) {
170 // TODO: Broken on darwin:
171 // https://github.com/ziglang/zig/issues/17260
172 // if (@hasDecl(os.TCP, "NODELAY")) {
173 // try os.setsockopt(socket.sockfd.?, os.IPPROTO.TCP, os.TCP.NODELAY, &std.mem.toBytes(@as(c_int, 1)));
174 // }
175 try posix.setsockopt(socket, posix.IPPROTO.TCP, 1, &std.mem.toBytes(@as(c_int, 1)));
176 }
177
178 if (@hasDecl(posix.SO, "REUSEPORT_LB")) {
179 try posix.setsockopt(socket, posix.SOL.SOCKET, posix.SO.REUSEPORT_LB, &std.mem.toBytes(@as(c_int, 1)));
180 } else if (@hasDecl(posix.SO, "REUSEPORT")) {
181 try posix.setsockopt(socket, posix.SOL.SOCKET, posix.SO.REUSEPORT, &std.mem.toBytes(@as(c_int, 1)));
182 } else {
183 try posix.setsockopt(socket, posix.SOL.SOCKET, posix.SO.REUSEADDR, &std.mem.toBytes(@as(c_int, 1)));
184 }
185
186 {
187 const socklen = address.getOsSockLen();
188 try posix.bind(socket, &address.any, socklen);
189 try posix.listen(socket, 1024); // kernel backlog
190 }
191
192 const C = @TypeOf(ctx);
193
194 if (comptime blockingMode()) {
195 errdefer posix.close(socket);
196 var w = try Blocking(H).init(self.allocator, &self._state);
197 defer w.deinit();
198
199 const thrd = try std.Thread.spawn(.{}, Blocking(H).run, .{&w, socket, ctx});
200 log.info("starting blocking worker to listen on {}", .{address});
201
202 // incase listenInNewThread was used and is waiting for us to start
203 self._cond.signal();
204
205 // this is what we'll shutdown when stop() is called
206 self._signals[0] = socket;
207 self._mut.unlock();
208 thrd.join();
209 } else {
210 defer posix.close(socket);
211 const W = NonBlocking(H, C);
212
213 const allocator = self.allocator;
214
215 var signals = self._signals;
216 const worker_count = signals.len;
217 const threads = try allocator.alloc(Thread, worker_count);
218 const workers = try allocator.alloc(W, worker_count);
219
220 var started: usize = 0;
221
222 errdefer for (0..started) |i| {
223 // on success, these will be closed by a call to stop();
224 posix.close(signals[i]);
225 };
226
227 defer {
228 for (0..started) |i| {
229 workers[i].deinit();
230 }
231 allocator.free(threads);
232 allocator.free(workers);
233 }
234
235 for (0..worker_count) |i| {
236 const pipe = try posix.pipe2(.{.NONBLOCK = true});
237 errdefer posix.close(pipe[1]);
238
239 workers[i] = try W.init(self.allocator, &self._state, ctx);
240 errdefer workers[i].deinit();
241
242 threads[i] = try Thread.spawn(.{}, W.run, .{ &workers[i], socket, pipe[0] });
243
244 signals[i] = pipe[1];
245 started += 1;
246 }
247
248 log.info("starting nonblocking worker to listen on {}", .{address});
249
250 // in case startInNewThread is waiting
251 self._cond.signal();
252
253 self._mut.unlock();
254
255 for (threads) |thrd |{
256 thrd.join();
257 }
258 }
259 }
260
261 pub fn stop(self: *Self) void {
262 self._mut.lock();
263 defer self._mut.unlock();
264 for (self._signals) |s| {
265 posix.close(s);
266 }
267 }
268 };
269}
270
271// This is our Blocking worker. It's very different than NonBlocking and much simpler.
272pub fn Blocking(comptime H: type) type {
273 return struct {
274 allocator: Allocator,
275 handshake_timeout: Timeout,
276 connection_buffer_size: usize,
277 conn_manager: ConnManager(H, false),
278 handshake_pool: *Handshake.Pool,
279 buffer_provider: *buffer.Provider,
280
281 const Timeout = struct {
282 sec: u32,
283 timeval: [@sizeOf(std.posix.timeval)]u8,
284
285 // if sec is null, it means we want to cancel the timeout.
286 fn init(sec: u32) Timeout {
287 return .{
288 .sec = sec,
289 .timeval = std.mem.toBytes(posix.timeval{ .sec = @intCast(sec), .usec = 0 }),
290 };
291 }
292
293 pub const none = std.mem.toBytes(posix.timeval{.sec = 0, .usec = 0});
294 };
295
296 const Self = @This();
297
298 pub fn init(allocator: Allocator, state: *WorkerState) !Self {
299 const config = &state.config;
300
301 var conn_manager = try ConnManager(H, false).init(allocator);
302 errdefer conn_manager.deinit();
303
304 return .{
305 .conn_manager = conn_manager,
306 .allocator = allocator,
307 .handshake_pool = state.handshake_pool,
308 .buffer_provider = &state.buffer_provider,
309 .handshake_timeout = Timeout.init(config.handshake.timeout),
310 .connection_buffer_size = config.buffers.size orelse DEFAULT_BUFFER_SIZE,
311 };
312 }
313
314 pub fn deinit(self: *Self) void {
315 self.conn_manager.deinit();
316 }
317
318 pub fn run(self: *Self, listener: posix.socket_t, ctx: anytype) void {
319 defer self.shutdown();
320 while (true) {
321 var address: net.Address = undefined;
322 var address_len: posix.socklen_t = @sizeOf(net.Address);
323 const socket = posix.accept(listener, &address.any, &address_len, posix.SOCK.CLOEXEC) catch |err| {
324 if (err == error.ConnectionAborted) {
325 log.info("received shutdown signal", .{});
326 return;
327 }
328 log.err("failed to accept socket: {}", .{err});
329 continue;
330 };
331 log.debug("({}) connected", .{address});
332
333 const thread = std.Thread.spawn(.{}, Self.handleConnection, .{self, socket, address, ctx}) catch |err| {
334 posix.close(socket);
335 log.err("({}) failed to spawn connection thread: {}", .{address, err});
336 continue;
337 };
338 thread.detach();
339 }
340 }
341
342 // Called in a thread started above in listen.
343 // Wrapper around _handleConnection so that we can handle erros
344 fn handleConnection(self: *Self, socket: posix.socket_t, address: net.Address, ctx: anytype) void {
345 self._handleConnection(socket, address, ctx) catch |err| {
346 log.err("({}) uncaught error in connection handler: {}", .{address, err});
347 };
348 }
349
350 fn _handleConnection(self: *Self, socket: posix.socket_t, address: net.Address, ctx: anytype) !void {
351 const conn_manager = &self.conn_manager;
352 const hc = try conn_manager.create(socket, address, timestamp());
353
354 {
355 // Do our handshake
356 errdefer self.cleanupConn(hc);
357 const timeout = self.handshake_timeout;
358 const deadline = timestamp() + timeout.sec;
359 try posix.setsockopt(socket, posix.SOL.SOCKET, posix.SO.RCVTIMEO, &timeout.timeval);
360
361 while (true) {
362 if (handleHandshake(H, self, hc, ctx) == false) {
363 self.cleanupConn(hc);
364 return;
365 }
366 if (hc.handler != null) {
367 // if we have a handler, the our handshake completed
368 break;
369 }
370 if (timestamp() > deadline) {
371 self.cleanupConn(hc);
372 return;
373 }
374 }
375 }
376
377 return self.readLoop(hc);
378 }
379
380 // The readloop is extracted from _handleConnection so that it can be called
381 // directly when integrating with an http server
382 pub fn readLoop(self: *Self, hc: *HandlerConn(H)) !void {
383 defer self.cleanupConn(hc);
384
385 try posix.setsockopt(hc.socket, posix.SOL.SOCKET, posix.SO.RCVTIMEO, &Timeout.none);
386
387 // In BlockingMode, we always assign a reader for the duration of the connection
388 // In scenarios where client rarely send data, this is going to use up an unecessary amount
389 // of memory, but unlike nonblocking mode, we can't initialize the buffer just-in-time.
390 const reader_buf = try self.allocator.alloc(u8, self.connection_buffer_size);
391 defer self.allocator.free(reader_buf);
392
393 hc.reader = Reader.init(reader_buf, self.buffer_provider);
394 while (true) {
395 if (handleClientData(H, hc, self.allocator, undefined) == false) {
396 break;
397 }
398 }
399 }
400
401 fn cleanupConn(self: *Self, hc: *HandlerConn(H)) void {
402 hc.conn.close();
403 self.conn_manager.cleanup(hc);
404 }
405
406 fn shutdown(self: *Self) void {
407 var conn_manager = &self.conn_manager;
408 conn_manager.shutdown(self);
409
410 // wait up to 1 second for every connection to cleanly shutdown
411 var i: usize = 0;
412 while (i < 10) : (i += 1) {
413 if (conn_manager.count() == 0) {
414 return;
415 }
416 std.time.sleep(std.time.ns_per_ms * 100);
417 }
418 }
419
420 // called for each hc when shutting down
421 fn shutdownCleanup(_: *Self, _: *HandlerConn(H)) void {
422
423 }
424 };
425}
426
427fn NonBlocking(comptime H: type, comptime C: type) type {
428 return struct {
429 ctx: C,
430
431 // KQueue or Epoll, depending on the platform
432 loop: Loop,
433 thread_pool: *ThreadPool,
434
435 handshake_timeout: u32,
436 handshake_pool: *Handshake.Pool,
437 base: NonBlockingBase(H, true),
438
439 const Self = @This();
440 const ThreadPool = @import("thread_pool.zig").ThreadPool(Self.dataAvailable);
441
442 pub fn init(allocator: Allocator, state: *WorkerState, ctx: C) !Self {
443 var base = try NonBlockingBase(H, true).init(allocator, state);
444 errdefer base.deinit();
445
446 const loop = try Loop.init();
447 errdefer loop.deinit();
448
449 const config = &state.config;
450 var thread_pool = try ThreadPool.init(allocator, .{
451 .count = config.thread_pool.count orelse 4,
452 .backlog = config.thread_pool.backlog orelse 500,
453 .buffer_size = config.thread_pool.buffer_size orelse 32_768,
454 });
455 errdefer thread_pool.deinit();
456
457 return .{
458 .ctx = ctx,
459 .loop = loop,
460 .base = base,
461 .thread_pool = thread_pool,
462 .handshake_pool = state.handshake_pool,
463 .handshake_timeout = state.config.handshake.timeout,
464 };
465 }
466
467 pub fn deinit(self: *Self) void {
468 self.loop.deinit();
469 self.base.deinit();
470 self.thread_pool.deinit();
471 }
472
473 fn run(self: *Self, listener: posix.socket_t, signal: posix.fd_t) void {
474 self.loop.monitorAccept(listener) catch |err| {
475 log.err("failed to add monitor to listening socket: {}", .{err});
476 return;
477 };
478
479 self.loop.monitorSignal(signal) catch |err| {
480 log.err("failed to add monitor to signal pipe: {}", .{err});
481 return;
482 };
483
484 const thread_pool = self.thread_pool;
485 const conn_manager = &self.base.conn_manager;
486 const handshake_timeout = self.handshake_timeout;
487
488 var now = timestamp();
489 var oldest_pending_start_time: ?u32 = null;
490 while (true) {
491 const handshake_cutoff = now - handshake_timeout;
492 const timeout = self.prepareToWait(handshake_cutoff) orelse blk: {
493 if (oldest_pending_start_time) |started| {
494 break :blk if (started < handshake_cutoff) 1 else @as(i32, @intCast(started - handshake_cutoff));
495 }
496 break :blk null;
497 };
498
499 var it = self.loop.wait(timeout) catch |err| {
500 log.err("failed to wait on events: {}", .{err});
501 std.time.sleep(std.time.ns_per_s);
502 continue;
503 };
504
505 now = timestamp();
506 oldest_pending_start_time = null;
507 while (it.next()) |data| {
508 if (data == 0) {
509 self.accept(listener, now) catch |err| {
510 log.err("accept error: {}", .{err});
511 std.time.sleep(std.time.ns_per_ms);
512 };
513 continue;
514 }
515
516 if (data == 1) {
517 self.base.shutdown();
518 return;
519 }
520
521 const hc: *HandlerConn(H) = @ptrFromInt(data);
522 if (hc.state == .handshake) {
523 // we need to get this out of the pending list, so that it doesn't
524 // cause a timeout while we're processing it.
525 // But we stll need to care about its timeout.
526 if (oldest_pending_start_time) |s| {
527 oldest_pending_start_time = @min(hc.conn.started, s);
528 } else {
529 oldest_pending_start_time = hc.conn.started;
530 }
531 conn_manager.activate(hc);
532 }
533 thread_pool.spawn(.{self, hc});
534 }
535 }
536 }
537
538 // Enforces timeouts, and returns when the next timeout should be checked.
539 fn prepareToWait(self: *Self, cutoff: u32) ?i32 {
540 const cm = &self.base.conn_manager;
541
542 // always ordered from oldest to newest, so once we find a conneciton
543 // that isn't timed out, we can stop
544 cm.lock.lock();
545 defer cm.lock.unlock();
546
547 var next_conn = cm.pending.head;
548 while (next_conn) |hc| {
549 const conn = &hc.conn;
550 const started = conn.started;
551 if (started > cutoff) {
552 // This is the first connection which hasn't timed out
553 // return the time until it times out.
554 return @intCast(started - cutoff);
555 }
556
557 next_conn = hc.next;
558
559 // this connection has timed out. Don't use self.cleanup since there's
560 // a bunch of stuff we can assume here..like there's no handler or reader
561 conn.close();
562 log.debug("({}) handshake timeout", .{conn.address});
563 if (hc.handshake) |h| {
564 h.release();
565 }
566 cm.pending.remove(hc);
567 cm.pool.destroy(hc);
568 }
569 return null;
570 }
571
572 fn accept(self: *Self, listener: posix.fd_t, now: u32) !void {
573 const max_conn = self.base.max_conn;
574 const conn_manager = &self.base.conn_manager;
575
576 while (conn_manager.count() < max_conn) {
577 var address: net.Address = undefined;
578 var address_len: posix.socklen_t = @sizeOf(net.Address);
579
580 const socket = posix.accept(listener, &address.any, &address_len, posix.SOCK.CLOEXEC) catch |err| {
581 // When available, we use SO_REUSEPORT_LB or SO_REUSEPORT, so WouldBlock
582 // should not be possible in those cases, but if it isn't available
583 // this error should be ignored as it means another thread picked it up.
584 return if (err == error.WouldBlock) {} else err;
585 };
586
587 log.debug("({}) connected", .{address});
588
589 {
590 errdefer posix.close(socket);
591 // socket is _probably_ in NONBLOCKING mode (it inherits
592 // the flag from the listening socket).
593 const flags = try posix.fcntl(socket, posix.F.GETFL, 0);
594 const nonblocking = @as(u32, @bitCast(posix.O{ .NONBLOCK = true }));
595 if (flags & nonblocking == nonblocking) {
596 // Yup, it's in nonblocking mode. Disable that flag to
597 // put it in blocking mode.
598 _ = try posix.fcntl(socket, posix.F.SETFL, flags & ~nonblocking);
599 }
600 }
601 const hc = try self.base.newConn(socket, address, now);
602 self.loop.monitorRead(hc, false) catch |err| {
603 self.base.cleanupConn(hc);
604 return err;
605 };
606 }
607 }
608
609 // Called in a thread-pool thread/
610 // !! Access to self has to be synchronized !!
611 // There can only be 1 dataAvailable executing per HC at any given time.
612 // Access to HC *should not* need to be synchronized. Access to hc.conn
613 // needs to be synchronized once &hc.conn is passed to the handler during
614 // handshake (Conn is self-synchronized).
615 // Shutdown throws a wrench in our synchronization model, since we could
616 // be shutting down while we're processing data...but hopefully the way we
617 // shutdown (by waiting for all the thread pools threads to end) solves this.
618 // Else, we'll need to throw a bunch of locking around HC just to handle shutdown.
619 fn dataAvailable(self: *Self, hc: *HandlerConn(H), thread_buf: []u8) void {
620 var success = false;
621 if (hc.handler == null) {
622 success = self.dataForHandshake(hc) catch |err| blk: {
623 log.err("({}) error processing handshake: {}", .{hc.conn.address, err});
624 break :blk false;
625 };
626 } else {
627 success = self.base.dataAvailable(hc, thread_buf);
628 }
629
630 var conn = &hc.conn;
631 var closed: bool = undefined;
632 if (success == false) {
633 conn.close();
634 closed = true;
635 } else {
636 closed = conn.isClosed();
637 }
638
639 if (closed) {
640 self.base.cleanupConn(hc);
641 } else {
642 self.loop.monitorRead(hc, true) catch |err| {
643 log.debug("({}) failed to add read event monitor: {}", .{conn.address, err});
644 conn.close();
645 self.base.cleanupConn(hc);
646 };
647 }
648 }
649
650 fn dataForHandshake(self: *Self, hc: *HandlerConn(H)) !bool {
651 if (handleHandshake(H, self, hc, self.ctx) == false) {
652 return false;
653 }
654
655 if (hc.handler == null) {
656 // still don't have a handshake
657 self.base.conn_manager.inactive(hc);
658 }
659
660 return true;
661 }
662 };
663}
664
665fn NonBlockingBase(comptime H: type, comptime MANAGE_HS: bool) type {
666 return struct {
667 allocator: Allocator,
668
669 buffer_provider: *buffer.Provider,
670
671 // App can configure the use of a "small" buffer pool. This is the difference
672 // between assigning a buffer to the connection just-in-time per message, or always
673 small_buffer_pool: ?buffer.Pool,
674 connection_buffer_size: usize,
675
676 max_conn: usize,
677 conn_manager: ConnManager(H, MANAGE_HS),
678
679 const Self = @This();
680
681 fn init(allocator: Allocator, state: *WorkerState) !Self {
682 const config = &state.config;
683
684 var conn_manager = try ConnManager(H, MANAGE_HS).init(allocator);
685 errdefer conn_manager.deinit();
686
687 const connection_buffer_size = config.buffers.size orelse DEFAULT_BUFFER_SIZE;
688 var small_buffer_pool: ?buffer.Pool = null;
689 if (config.buffers.pool) |pool_count| {
690 small_buffer_pool = try buffer.Pool.init(allocator, pool_count, connection_buffer_size);
691 }
692
693 errdefer if (small_buffer_pool) |sbp| {
694 sbp.deinit();
695 };
696
697 return .{
698 .allocator = allocator,
699 .max_conn = config.max_conn,
700 .conn_manager = conn_manager,
701 .small_buffer_pool = small_buffer_pool,
702 .buffer_provider = &state.buffer_provider,
703 .connection_buffer_size = connection_buffer_size,
704 };
705 }
706
707 fn deinit(self: *Self) void {
708 self.conn_manager.deinit();
709 if (self.small_buffer_pool) |*sbp| {
710 sbp.deinit();
711 }
712 }
713
714 fn newConn(self: *Self, socket: posix.socket_t, address: net.Address, time: u32) !*HandlerConn(H) {
715 return self.conn_manager.create(socket, address, time);
716 }
717
718 pub fn dataAvailable(self: *Self, hc: *HandlerConn(H), thread_buf: []u8) bool {
719 return self._dataAvailable(hc, thread_buf) catch |err| {
720 log.err("({}) error processing client message: {}", .{hc.conn.address, err});
721 return false;
722 };
723 }
724
725 fn _dataAvailable(self: *Self, hc: *HandlerConn(H), thread_buf: []u8) !bool {
726 if (hc.reader == null) {
727 const reader_buf = if (self.small_buffer_pool) |*sbp| try sbp.acquireOrCreate() else try self.allocator.alloc(u8, self.connection_buffer_size);
728 hc.reader = Reader.init(reader_buf, self.buffer_provider);
729 }
730 const reader = &hc.reader.?;
731
732 var fba = FixedBufferAllocator.init(thread_buf);
733 const ok = handleClientData(H, hc, self.allocator, &fba);
734
735 if (self.small_buffer_pool) |*sbp| {
736 if (reader.isEmpty()) {
737 sbp.release(reader.static);
738 hc.reader = null;
739 }
740 }
741
742 return ok;
743 }
744
745 fn cleanupConn(self: *Self, hc: *HandlerConn(H)) void {
746 if (hc.reader) |*reader| {
747 if (self.small_buffer_pool) |*sbp| {
748 sbp.release(reader.static);
749 } else {
750 self.allocator.free(reader.static);
751 }
752 }
753 self.conn_manager.cleanup(hc);
754 }
755
756 fn shutdown(self: *Self) void {
757 log.info("received shutdown signal", .{});
758 self.conn_manager.shutdown(self);
759 }
760
761 // called for each hc when shutting down
762 fn shutdownCleanup(self: *Self, hc: *HandlerConn(H)) void {
763 if (hc.reader) |*reader| {
764 if (self.small_buffer_pool == null) {
765 self.allocator.free(reader.static);
766 }
767 }
768 }
769 };
770}
771
772const Loop = switch (@import("builtin").os.tag) {
773 .macos, .ios, .tvos, .watchos, .freebsd, .netbsd, .dragonfly, .openbsd => KQueue,
774 .linux => EPoll,
775 else => unreachable,
776};
777
778const KQueue = struct {
779 q: i32,
780 change_count: usize,
781 change_buffer: [16]Kevent,
782 event_list: [64]Kevent,
783
784 const Kevent = posix.Kevent;
785
786 fn init() !KQueue {
787 return .{
788 .q = try posix.kqueue(),
789 .change_count = 0,
790 .change_buffer = undefined,
791 .event_list = undefined,
792 };
793 }
794
795 fn deinit(self: KQueue) void {
796 posix.close(self.q);
797 }
798
799 fn monitorAccept(self: *KQueue, fd: c_int) !void {
800 try self.change(fd, 0, posix.system.EVFILT.READ, posix.system.EV.ADD);
801 }
802
803 fn monitorSignal(self: *KQueue, fd: c_int) !void {
804 try self.change(fd, 1, posix.system.EVFILT.READ, posix.system.EV.ADD);
805 }
806
807 // Normally, we add the socket in the worker thread with rearm == false.
808 // Because this is a DISPATCH, it'll only fire once until we rearm it which
809 // we do by re-enabling it with rearm == true.
810 // However, notice that the rearm path also has the EV.ADD flag. From the above
811 // description, this should not be necessary.
812 // But monitorRead where rearm == true is also used by our generic ServerLoop when
813 // taking over a connection (say from httpz). Hence, we need the EV.ADD flag too.
814 fn monitorRead(self: *KQueue, hc: anytype, comptime rearm: bool) !void {
815 if (rearm == false) {
816 return self.change(hc.socket, @intFromPtr(hc), posix.system.EVFILT.READ, posix.system.EV.ADD | posix.system.EV.ENABLE | posix.system.EV.DISPATCH);
817 }
818 const event = Kevent{
819 .ident = @intCast(hc.socket),
820 .filter = posix.system.EVFILT.READ,
821 .flags = posix.system.EV.ADD | posix.system.EV.ENABLE | posix.system.EV.DISPATCH,
822 .fflags = 0,
823 .data = 0,
824 .udata = @intFromPtr(hc),
825 };
826 _ = try posix.kevent(self.q, &.{event}, &[_]Kevent{}, null);
827 }
828
829 fn change(self: *KQueue, fd: posix.fd_t, data: usize, filter: i16, flags: u16) !void {
830 var change_count = self.change_count;
831 var change_buffer = &self.change_buffer;
832
833 if (change_count == change_buffer.len) {
834 // calling this with an empty event_list will return immediate
835 _ = try posix.kevent(self.q, change_buffer, &[_]Kevent{}, null);
836 change_count = 0;
837 }
838 change_buffer[change_count] = .{
839 .ident = @intCast(fd),
840 .filter = filter,
841 .flags = flags,
842 .fflags = 0,
843 .data = 0,
844 .udata = data,
845 };
846 self.change_count = change_count + 1;
847 }
848
849 fn wait(self: *KQueue, timeout_sec: ?i32) !Iterator {
850 const event_list = &self.event_list;
851 const timeout: ?posix.timespec = if (timeout_sec) |ts| posix.timespec{ .sec = ts, .nsec = 0 } else null;
852 const event_count = try posix.kevent(self.q, self.change_buffer[0..self.change_count], event_list, if (timeout) |ts| &ts else null);
853 self.change_count = 0;
854
855 return .{
856 .index = 0,
857 .events = event_list[0..event_count],
858 };
859 }
860
861 const Iterator = struct {
862 index: usize,
863 events: []Kevent,
864
865 fn next(self: *Iterator) ?usize {
866 const index = self.index;
867 const events = self.events;
868 if (index == events.len) {
869 return null;
870 }
871 self.index = index + 1;
872 return self.events[index].udata;
873 }
874 };
875};
876
877const EPoll = struct {
878 q: i32,
879 event_list: [64]EpollEvent,
880
881 const linux = std.os.linux;
882 const EpollEvent = linux.epoll_event;
883
884 fn init() !EPoll {
885 return .{
886 .event_list = undefined,
887 .q = try posix.epoll_create1(0),
888 };
889 }
890
891 fn deinit(self: EPoll) void {
892 posix.close(self.q);
893 }
894
895 fn monitorAccept(self: *EPoll, fd: c_int) !void {
896 var event = linux.epoll_event{ .events = linux.EPOLL.IN, .data = .{ .ptr = 0 } };
897 return std.posix.epoll_ctl(self.q, linux.EPOLL.CTL_ADD, fd, &event);
898 }
899
900 fn monitorSignal(self: *EPoll, fd: c_int) !void {
901 var event = linux.epoll_event{ .events = linux.EPOLL.IN, .data = .{ .ptr = 1 } };
902 return std.posix.epoll_ctl(self.q, linux.EPOLL.CTL_ADD, fd, &event);
903 }
904
905 fn monitorRead(self: *EPoll, hc: anytype, comptime rearm: bool) !void {
906 const op = if (rearm) linux.EPOLL.CTL_MOD else linux.EPOLL.CTL_ADD;
907 var event = linux.epoll_event{ .events = linux.EPOLL.IN | linux.EPOLL.ONESHOT, .data = .{ .ptr = @intFromPtr(hc) } };
908 return posix.epoll_ctl(self.q, op, hc.socket, &event);
909 }
910
911 fn wait(self: *EPoll, timeout_sec: ?i32) !Iterator {
912 const event_list = &self.event_list;
913 var timeout: i32 = -1;
914 if (timeout_sec) |sec| {
915 if (sec > 2147483) {
916 // max supported timeout by epoll_wait.
917 timeout = 2147483647;
918 } else {
919 timeout = sec * 1000;
920 }
921 }
922
923 const event_count = posix.epoll_wait(self.q, event_list, timeout);
924 return .{
925 .index = 0,
926 .events = event_list[0..event_count],
927 };
928 }
929
930 const Iterator = struct {
931 index: usize,
932 events: []EpollEvent,
933
934 fn next(self: *Iterator) ?usize {
935 const index = self.index;
936 const events = self.events;
937 if (index == events.len) {
938 return null;
939 }
940 self.index = index + 1;
941 return self.events[index].data.ptr;
942 }
943 };
944};
945
946// We want to extract as much common logic as possible from the Blocking and
947// NonBlockign workers. The code from this point on is meant to be used with
948// both workers, independently from the blocking/nonblocking nonsense.
949
950
951// Abstraction ontop of NonBlocking and Blocking. Exists solely for integration
952// with httpz (or any other http library I guess.). Serves a similar purpose
953// as Server, but doesn't accept/listen.
954pub fn Worker(comptime H: type) type {
955 return struct {
956 worker: W,
957
958 const Self = @This();
959 const W = if (blockingMode()) Blocking(H) else NonBlockingBase(H, false);
960
961 pub fn init(allocator: Allocator, state: *WorkerState) !Self {
962 return .{
963 .worker = try W.init(allocator, state),
964 };
965 }
966
967 pub fn deinit(self: *Self) void {
968 self.worker.deinit();
969 }
970
971 pub fn createConn(self: *Self, socket: posix.socket_t, address: net.Address, now: u32) !*HandlerConn(H) {
972 return self.worker.conn_manager.create(socket, address, now);
973 }
974
975 pub fn cleanupConn(self: *Self, hc: *HandlerConn(H)) void {
976 self.worker.cleanupConn(hc);
977 }
978
979 pub fn shutdown(self: *Self) void {
980 self.worker.shutdown();
981 }
982 };
983}
984
985// These are things that both the Blocking and NonBlocking workers need. Just cleaner
986// to have a single place for it. This could used to be directly in Server(H), with
987// Server(H) having handshake_pool and buffer_provider fields, but it was extracted
988// into its own struct for integration with webservers (i.e. httpz). The goal is
989// that a webserver can have websocket support without starting a full Server.
990pub const WorkerState = struct {
991 config: Config,
992 handshake_pool: *Handshake.Pool,
993 buffer_provider: buffer.Provider,
994
995 pub fn init(allocator: Allocator, config: Config) !WorkerState {
996 const handshake_pool_count = config.handshake.count orelse 32;
997 const handshake_max_size = config.handshake.max_size orelse 1024;
998 const handshake_max_headers = config.handshake.max_headers orelse 10;
999
1000 var handshake_pool = try Handshake.Pool.init(allocator, handshake_pool_count, handshake_max_size, handshake_max_headers);
1001 errdefer handshake_pool.deinit();
1002
1003 const max_message_size = config.max_message_size;
1004 const large_buffer_count = config.buffers.large_count orelse 8;
1005 const large_buffer_size = config.buffers.large_size orelse @min((config.buffers.size orelse DEFAULT_BUFFER_SIZE) * 2, max_message_size);
1006
1007 var buffer_provider = try buffer.Provider.init(allocator, .{
1008 .max = max_message_size,
1009 .size = large_buffer_size,
1010 .count = large_buffer_count,
1011 });
1012 errdefer buffer_provider.deinit();
1013
1014 return .{
1015 .config = config,
1016 .handshake_pool = handshake_pool,
1017 .buffer_provider = buffer_provider,
1018 };
1019 }
1020
1021 pub fn deinit(self: *WorkerState) void {
1022 self.handshake_pool.deinit();
1023 self.buffer_provider.deinit();
1024 }
1025};
1026
1027// In the Blocking worker, all the state could be stored on the spawn'd threads
1028// stack. The only reason we use a HandlerConn(H) in there is to be able to re-use
1029// code with the NonBlocking worker.
1030//
1031// For the NonBlocking worker, HandlerConn(H) is critical as it contains all the
1032// state for a connection. It lives on the heap, a pointer is registered into
1033// the event loop, and passed back when data is ready to read.
1034// * If handler is null, it means we haven't done our handshake yet.
1035// * If handler is null AND handshake is null, it means we havent' received
1036// any data yet (we delay creating the handshake state until we at least have
1037// some data ready).
1038pub fn HandlerConn(comptime H: type) type {
1039 return struct {
1040 state: State,
1041 conn: Conn,
1042 handler: ?H,
1043 reader: ?Reader,
1044 socket: posix.socket_t, // denormalization from conn.stream.handle
1045 handshake: ?*Handshake.State,
1046 next: ?*HandlerConn(H) = null,
1047 prev: ?*HandlerConn(H) = null,
1048
1049 const State = enum {
1050 handshake,
1051 active,
1052 };
1053 };
1054}
1055
1056pub fn ConnManager(comptime H: type, comptime MANAGE_HS: bool) type {
1057 return struct {
1058 lock: Thread.Mutex,
1059 active: List(HandlerConn(H)),
1060 pending: List(HandlerConn(H)),
1061 pool: std.heap.MemoryPool(HandlerConn(H)),
1062
1063 const Self = @This();
1064
1065 pub fn init(allocator: Allocator) !Self {
1066 var pool = std.heap.MemoryPool(HandlerConn(H)).init(allocator);
1067 errdefer pool.deinit();
1068
1069 return .{
1070 .lock = .{},
1071 .pool = pool,
1072 .active = .{},
1073 .pending = .{},
1074 };
1075 }
1076
1077 pub fn deinit(self: *Self) void {
1078 self.pool.deinit();
1079 }
1080
1081 pub fn count(self: *Self) usize {
1082 self.lock.lock();
1083 defer self.lock.unlock();
1084 if (MANAGE_HS == false) {
1085 return self.active.len;
1086 }
1087 return self.active.len + self.pending.len;
1088 }
1089
1090 pub fn create(self: *Self, socket: posix.socket_t, address: net.Address, now: u32) !*HandlerConn(H) {
1091 errdefer posix.close(socket);
1092
1093 self.lock.lock();
1094 defer self.lock.unlock();
1095
1096 const hc = try self.pool.create();
1097 hc.* = .{
1098 .state = if (MANAGE_HS) .handshake else .active,
1099 .socket = socket,
1100 .handler = null,
1101 .handshake = null,
1102 .reader = null,
1103 .conn = .{
1104 ._closed = false,
1105 .started = now,
1106 .address = address,
1107 .stream = .{.handle = socket},
1108 },
1109 };
1110
1111 if (comptime MANAGE_HS) {
1112 // Still waiting for a handshake. Only care about this with the full
1113 // NonBlocking worker
1114 self.pending.insert(hc);
1115 } else {
1116 self.active.insert(hc);
1117 }
1118 return hc;
1119 }
1120
1121 pub fn activate(self: *Self, hc: *HandlerConn(H)) void {
1122 std.debug.assert(MANAGE_HS == true);
1123
1124 // our caller made sute this was the case
1125 std.debug.assert(hc.state == .handshake);
1126 self.lock.lock();
1127 defer self.lock.unlock();
1128 self.pending.remove(hc);
1129 self.active.insert(hc);
1130 hc.state = .active;
1131 }
1132
1133 pub fn inactive(self: *Self, hc: *HandlerConn(H)) void {
1134 std.debug.assert(MANAGE_HS == false);
1135
1136 // this should only be called when we need more data to complete the handshake
1137 // which should only happen on an active connection
1138 std.debug.assert(hc.state == .active);
1139 hc.state = .handshake;
1140
1141 self.lock.lock();
1142 defer self.lock.unlock();
1143 self.active.remove(hc);
1144 self.pending.insert(hc);
1145 }
1146
1147 pub fn cleanup(self: *Self, hc: *HandlerConn(H)) void {
1148 if (hc.handshake) |h| {
1149 h.release();
1150 }
1151
1152 if (hc.reader) |*r| {
1153 r.deinit();
1154 }
1155
1156 if (hc.handler) |*h| {
1157 if (comptime std.meta.hasFn(H, "close")) {
1158 h.close();
1159 }
1160 hc.handler = null;
1161 }
1162
1163 self.lock.lock();
1164 if (hc.state == .active) {
1165 self.active.remove(hc);
1166 } else {
1167 self.pending.remove(hc);
1168 }
1169
1170 self.pool.destroy(hc);
1171 self.lock.unlock();
1172 }
1173
1174 pub fn shutdown(self: *Self, worker: anytype) void {
1175 self.lock.lock();
1176 defer self.lock.unlock();
1177
1178 shutdownList(self.active.head, worker);
1179 shutdownList(self.pending.head, worker);
1180 }
1181
1182 // This is sloppy and leaves things in an unrecoverable state. To keep
1183 // things clean, we should call self.cleanup(hc) on each entry in the list
1184 // but that does a bunch of things we don't need if we know that we're
1185 // shutting down - like returning data to the pools, nd popping items
1186 // out of the list.
1187 fn shutdownList(head: ?*HandlerConn(H), worker: anytype) void {
1188 var next_node = head;
1189 while (next_node) |hc| {
1190
1191 if (comptime std.meta.hasFn(H, "close")) {
1192 if (hc.handler) |*h| {
1193 h.close();
1194 hc.handler = null;
1195 }
1196 }
1197
1198 worker.shutdownCleanup(hc);
1199
1200 const conn = &hc.conn;
1201 if (conn.isClosed() == false) {
1202 conn.writeClose();
1203 conn.close();
1204 }
1205 next_node = hc.next;
1206 }
1207 }
1208 };
1209}
1210
1211// This is what actually gets exposed to the app
1212pub const Conn = struct {
1213 _closed: bool,
1214 started: u32,
1215 stream: net.Stream,
1216 address: net.Address,
1217 lock: Thread.Mutex = .{},
1218
1219 pub fn isClosed(self: *Conn) bool {
1220 // don't use lock to protect _closed. `isClosed` is called from
1221 // the worker thread and we don't want that potentially blocked while
1222 // a write is going on.
1223 return @atomicLoad(bool, &self._closed, .monotonic);
1224 }
1225
1226 pub fn close(self: *Conn) void {
1227 if (@atomicRmw(bool, &self._closed, .Xchg, true, .monotonic) == false) {
1228 posix.close(self.stream.handle);
1229 }
1230 }
1231
1232 pub fn writeBin(self: *Conn, data: []const u8) !void {
1233 return self.writeFrame(.binary, data);
1234 }
1235
1236 pub fn writeText(self: *Conn, data: []const u8) !void {
1237 return self.writeFrame(.text, data);
1238 }
1239
1240 pub fn write(self: *Conn, data: []const u8) !void {
1241 return self.writeFrame(.text, data);
1242 }
1243
1244 pub fn writePing(self: *Conn, data: []u8) !void {
1245 return self.writeFrame(.ping, data);
1246 }
1247
1248 pub fn writePong(self: *Conn, data: []u8) !void {
1249 return self.writeFrame(.pong, data);
1250 }
1251
1252 pub fn writeClose(self: *Conn) void {
1253 self.writeFramed(CLOSE_NORMAL) catch {};
1254 }
1255
1256 pub fn writeCloseWithCode(self: *Conn, code: u16) void {
1257 var buf: [2]u8 = undefined;
1258 std.mem.writeInt(u16, &buf, code, .big);
1259 self.writeFrame(.close, &buf) catch {};
1260 }
1261
1262 pub fn writeCloseWithReason(self: *Conn, code: u16, reason: []const u8) !void {
1263 if (reason.len > 123) {
1264 return error.ReasonTooLong;
1265 }
1266 var buf: [4]u8 = undefined;
1267 buf[0] = @intFromEnum(OpCode.close);
1268 buf[1] = @intCast(reason.len + 2);
1269 std.mem.writeInt(u16, buf[2..], code, .big);
1270
1271 var vec = [2]std.posix.iovec_const{
1272 .{ .len = buf.len, .base = &buf },
1273 .{ .len = reason.len, .base = reason.ptr },
1274 };
1275
1276 return self.writeAllIOVec(&vec);
1277 }
1278
1279 pub fn writeFrame(self: *Conn, op_code: OpCode, data: []const u8) !void {
1280 const l = data.len;
1281 const stream = self.stream;
1282
1283 // maximum possible prefix length. op_code + length_type + 8byte length
1284 var buf: [10]u8 = undefined;
1285 var header: []const u8 = undefined;
1286 buf[0] = @intFromEnum(op_code);
1287
1288 if (l <= 125) {
1289 buf[1] = @intCast(l);
1290 header = buf[0..2];
1291 } else if (l < 65536) {
1292 buf[1] = 126;
1293 buf[2] = @intCast((l >> 8) & 0xFF);
1294 buf[3] = @intCast(l & 0xFF);
1295 header = buf[0..4];
1296 } else {
1297 buf[1] = 127;
1298 buf[2] = @intCast((l >> 56) & 0xFF);
1299 buf[3] = @intCast((l >> 48) & 0xFF);
1300 buf[4] = @intCast((l >> 40) & 0xFF);
1301 buf[5] = @intCast((l >> 32) & 0xFF);
1302 buf[6] = @intCast((l >> 24) & 0xFF);
1303 buf[7] = @intCast((l >> 16) & 0xFF);
1304 buf[8] = @intCast((l >> 8) & 0xFF);
1305 buf[9] = @intCast(l & 0xFF);
1306 header = buf[0..];
1307 }
1308
1309 if (l == 0) {
1310 // no body, just write the header
1311 self.lock.lock();
1312 defer self.lock.unlock();
1313 return stream.writeAll(header);
1314 }
1315
1316 var vec = [2]std.posix.iovec_const{
1317 .{ .len = header.len, .base = header.ptr },
1318 .{ .len = data.len, .base = data.ptr },
1319 };
1320
1321 return self.writeAllIOVec(&vec);
1322 }
1323
1324 pub fn writeFramed(self: *Conn, data: []const u8) !void {
1325 self.lock.lock();
1326 defer self.lock.unlock();
1327 try self.stream.writeAll(data);
1328 }
1329
1330 fn writeAllIOVec(self: *Conn, vec: []std.posix.iovec_const) !void {
1331 const socket = self.stream.handle;
1332
1333 self.lock.lock();
1334 defer self.lock.unlock();
1335
1336 var i: usize = 0;
1337 while (true) {
1338 var n = try std.posix.writev(socket, vec[i..]);
1339 while (n >= vec[i].len) {
1340 n -= vec[i].len;
1341 i += 1;
1342 if (i >= vec.len) return;
1343 }
1344 vec[i].base += n;
1345 vec[i].len -= n;
1346 }
1347 }
1348};
1349
1350fn handleHandshake(comptime H: type, worker: anytype, hc: *HandlerConn(H), ctx: anytype) bool {
1351 return _handleHandshake(H, worker, hc, ctx) catch |err| {
1352 log.warn("({}) uncaugh error processing handshake: {}", .{hc.conn.address, err});
1353 return false;
1354 };
1355}
1356
1357fn _handleHandshake(comptime H: type, worker: anytype, hc: *HandlerConn(H), ctx: anytype) !bool {
1358 std.debug.assert(hc.handler == null);
1359
1360 var state = hc.handshake orelse blk: {
1361 const s = try worker.handshake_pool.acquire();
1362 hc.handshake = s;
1363 break :blk s;
1364 };
1365
1366 var buf = state.buf;
1367 var conn = &hc.conn;
1368 const len = state.len;
1369
1370 if (len == buf.len) {
1371 log.warn("({}) handshake request exceeded maximum configured size ({d})", .{conn.address, buf.len});
1372 return false;
1373 }
1374
1375 const n = posix.read(hc.socket, buf[len..]) catch |err| {
1376 switch (err) {
1377 error.BrokenPipe, error.ConnectionResetByPeer => log.debug("({}) handshake connection closed: {}", .{conn.address, err}),
1378 error.WouldBlock => {
1379 std.debug.assert(blockingMode());
1380 log.debug("({}) handshake timeout", .{conn.address});
1381 },
1382 else => log.warn("({}) handshake error reading from socket: {}", .{conn.address, err}),
1383 }
1384 return false;
1385 };
1386
1387 if (n == 0) {
1388 log.debug("({}) handshake connection closed", .{conn.address});
1389 return false;
1390 }
1391
1392 state.len = len + n;
1393 const handshake = Handshake.parse(state) catch |err| {
1394 log.debug("({}) error parsing handshake: {}", .{conn.address, err});
1395 respondToHandshakeError(conn, err);
1396 return false;
1397 } orelse {
1398 // we need more data
1399 return true;
1400 };
1401
1402 state.release();
1403 hc.handshake = null;
1404
1405 // After this, the app has access to &hc.conn, so any access to the
1406 // conn has to be synchronized (which the conn does internally).
1407
1408 const handler = H.init(handshake, conn, ctx) catch |err| {
1409 if (comptime std.meta.hasFn(H, "handshakeErrorResponse")) {
1410 preHandOffWrite(H.handshakeErrorResponse(err));
1411 } else {
1412 respondToHandshakeError(conn, err);
1413 }
1414 log.debug("({}) " ++ @typeName(H) ++ ".init rejected request {}", .{conn.address, err});
1415 return false;
1416 };
1417
1418
1419 hc.handler = handler;
1420 try conn.writeFramed(&Handshake.createReply(handshake.key));
1421
1422 if (comptime std.meta.hasFn(H, "afterInit")) {
1423 handler.afterInit(ctx) catch |err| {
1424 log.debug("({}) " ++ @typeName(H) ++ ".afterInit error: {}", .{conn.address, err});
1425 return false;
1426 };
1427 }
1428
1429 log.debug("({}) connection successfully upgraded", .{conn.address});
1430 return true;
1431}
1432
1433fn handleClientData(comptime H: type, hc: *HandlerConn(H), allocator: Allocator, fba: *FixedBufferAllocator) bool {
1434 std.debug.assert(hc.handshake == null);
1435 return _handleClientData(H, hc, allocator, fba) catch |err| {
1436 log.warn("({}) uncaugh error handling incoming data: {}", .{hc.conn.address, err});
1437 return false;
1438 };
1439}
1440
1441fn _handleClientData(comptime H: type, hc: *HandlerConn(H), allocator: Allocator, fba: *FixedBufferAllocator) !bool {
1442 var conn = &hc.conn;
1443 var reader = &hc.reader.?;
1444 reader.fill(conn.stream) catch |err| {
1445 switch (err) {
1446 error.BrokenPipe, error.Closed, error.ConnectionResetByPeer => log.debug("({}) connection closed: {}", .{conn.address, err}),
1447 else => log.warn("({}) error reading from connection: {}", .{conn.address, err}),
1448 }
1449 return false;
1450 };
1451
1452 const handler = &hc.handler.?;
1453 while (true) {
1454 const has_more, const message = reader.read() catch |err| {
1455 switch (err) {
1456 error.LargeControl => conn.writeFramed(CLOSE_PROTOCOL_ERROR) catch {},
1457 error.ReservedFlags => conn.writeFramed(CLOSE_PROTOCOL_ERROR) catch {},
1458 else => {},
1459 }
1460 log.debug("({}) invalid websocket packet: {}", .{conn.address, err});
1461 return false;
1462 } orelse {
1463 // everything is fine, we just need more data
1464 return true;
1465 };
1466
1467 const message_type = message.type;
1468 defer reader.done(message_type);
1469
1470 log.debug("({}) received {s} message", .{hc.conn.address, @tagName(message_type)});
1471 switch (message_type) {
1472 .text, .binary => {
1473 const params = @typeInfo(@TypeOf(H.clientMessage)).Fn.params;
1474 const needs_allocator = comptime params[1].type == Allocator;
1475
1476 var arena: std.heap.ArenaAllocator = undefined;
1477 var fallback_allocator: FallbackAllocator = undefined;
1478 var aa: Allocator = undefined;
1479
1480 if (comptime needs_allocator) {
1481 arena = std.heap.ArenaAllocator.init(allocator);
1482 if (comptime blockingMode()) {
1483 aa = arena.allocator();
1484 } else {
1485 fallback_allocator = FallbackAllocator{
1486 .fba = fba,
1487 .fallback = arena.allocator(),
1488 .fixed = fba.allocator(),
1489 };
1490 aa = fallback_allocator.allocator();
1491 }
1492 }
1493
1494 defer if (comptime needs_allocator) {
1495 arena.deinit();
1496 };
1497
1498 switch (comptime params.len) {
1499 2 => handler.clientMessage(message.data) catch return false,
1500 3 => if (needs_allocator) {
1501 handler.clientMessage(aa, message.data) catch return false;
1502 } else {
1503 handler.clientMessage(message.data, if (message_type == .text) .text else .binary) catch return false;
1504 },
1505 4 => handler.clientMessage(aa, message.data, if (message_type == .text) .text else .binary) catch return false,
1506 else => @compileError(@typeName(H) ++ ".clientMessage has invalid parameter count"),
1507 }
1508 },
1509 .pong => if (comptime std.meta.hasFn(H, "clientPong")) {
1510 try handler.clientPong();
1511 },
1512 .ping => {
1513 const data = message.data;
1514 if (comptime std.meta.hasFn(H, "clientPing")) {
1515 try handler.clientPing(data);
1516 } else if (data.len == 0) {
1517 try hc.conn.writeFramed(EMPTY_PONG);
1518 } else {
1519 try hc.conn.writeFrame(.pong, data);
1520 }
1521 },
1522 .close => {
1523 defer conn.close();
1524 const data = message.data;
1525 if (comptime std.meta.hasFn(H, "clientClose")) {
1526 return handler.clientClose(data);
1527 }
1528
1529 const l = data.len;
1530 if (l == 0) {
1531 conn.writeClose();
1532 return false;
1533 }
1534
1535 if (l == 1) {
1536 // close with a payload always has to have at least a 2-byte payload,
1537 // since a 2-byte code is required
1538 try conn.writeFramed(CLOSE_PROTOCOL_ERROR);
1539 return false;
1540 }
1541
1542 const code = @as(u16, @intCast(data[1])) | (@as(u16, @intCast(data[0])) << 8);
1543 if (code < 1000 or code == 1004 or code == 1005 or code == 1006 or (code > 1013 and code < 3000)) {
1544 try conn.writeFramed(CLOSE_PROTOCOL_ERROR);
1545 return false;
1546 }
1547
1548 if (l == 2) {
1549 try conn.writeFramed(CLOSE_NORMAL);
1550 return false;
1551 }
1552
1553 const payload = data[2..];
1554 if (!std.unicode.utf8ValidateSlice(payload)) {
1555 // if we have a payload, it must be UTF8 (why?!)
1556 try conn.writeFramed(CLOSE_PROTOCOL_ERROR);
1557 } else {
1558 conn.writeClose();
1559 }
1560 return false;
1561 },
1562 }
1563
1564 if (conn.isClosed()) {
1565 return false;
1566 }
1567
1568 if (has_more == false) {
1569 // we don't have more data ready to be processed in our buffer
1570 // back to our caller for more data
1571 return true;
1572 }
1573 }
1574}
1575
1576fn respondToHandshakeError(conn: *Conn, err: anyerror) void {
1577 const response = switch (err) {
1578 error.Close => return,
1579 error.RequestTooLarge => buildError(400, "too large"),
1580 error.Timeout, error.WouldBlock => buildError(400, "timeout"),
1581 error.InvalidProtocol => buildError(400, "invalid protocol"),
1582 error.InvalidRequestLine => buildError(400, "invalid requestline"),
1583 error.InvalidHeader => buildError(400, "invalid header"),
1584 error.InvalidUpgrade => buildError(400, "invalid upgrade"),
1585 error.InvalidVersion => buildError(400, "invalid version"),
1586 error.InvalidConnection => buildError(400, "invalid connection"),
1587 error.MissingHeaders => buildError(400, "missingheaders"),
1588 error.Empty => buildError(400, "invalid request"),
1589 else => buildError(400, "unknown"),
1590 };
1591 preHandOffWrite(conn, response);
1592}
1593
1594fn buildError(comptime status: u16, comptime err: []const u8) []const u8 {
1595 return std.fmt.comptimePrint("HTTP/1.1 {d} \r\nConnection: Close\r\nError: {s}\r\nContent-Length: 0\r\n\r\n", .{ status, err });
1596}
1597
1598fn preHandOffWrite(conn: *Conn, response: []const u8) void {
1599 // "preHandOff" means we haven't given the application handler a reference
1600 // to *Conn yet. In theory, this means we don't need to worry about thread-safety
1601 // However, it is possible for the worker to be stopped while we're doing this
1602 // which causes issues unless we lock
1603 conn.lock.lock();
1604 defer conn.lock.unlock();
1605
1606 if (conn.isClosed()) {
1607 return;
1608 }
1609
1610 const socket = conn.stream.handle;
1611 const timeout = std.mem.toBytes(posix.timeval{ .sec = 5, .usec = 0 });
1612 posix.setsockopt(socket, posix.SOL.SOCKET, posix.SO.SNDTIMEO, &timeout) catch return;
1613
1614 var pos: usize = 0;
1615 while (pos < response.len) {
1616 const n = posix.write(socket, response[pos..]) catch return;
1617 if (n == 0) {
1618 // closed
1619 return;
1620 }
1621 pos += n;
1622 }
1623}
1624
1625fn timestamp() u32 {
1626 if (comptime @hasDecl(std.c, "CLOCK") == false) {
1627 return @intCast(std.time.timestamp());
1628 }
1629 var ts: posix.timespec = undefined;
1630 posix.clock_gettime(posix.CLOCK.REALTIME, &ts) catch unreachable;
1631 return @intCast(ts.sec);
1632}
1633
1634// intrusive doubly-linked list with count, not thread safe
1635fn List(comptime T: type) type {
1636 return struct {
1637 len: usize = 0,
1638 head: ?*T = null,
1639 tail: ?*T = null,
1640
1641 const Self = @This();
1642
1643 pub fn insert(self: *Self, node: *T) void {
1644 if (self.tail) |tail| {
1645 tail.next = node;
1646 node.prev = tail;
1647 self.tail = node;
1648 } else {
1649 self.head = node;
1650 self.tail = node;
1651 }
1652 self.len += 1;
1653 node.next = null;
1654 }
1655
1656 pub fn remove(self: *Self, node: *T) void {
1657 if (node.prev) |prev| {
1658 prev.next = node.next;
1659 } else {
1660 self.head = node.next;
1661 }
1662
1663 if (node.next) |next| {
1664 next.prev = node.prev;
1665 } else {
1666 self.tail = node.prev;
1667 }
1668 node.prev = null;
1669 node.next = null;
1670 self.len -= 1;
1671 }
1672 };
1673}
1674
1675const t = @import("../t.zig");
1676
1677var test_thread: Thread = undefined;
1678var test_server: Server(TestHandler) = undefined;
1679var global_test_allocator = std.heap.GeneralPurposeAllocator(.{}){};
1680
1681test "tests:beforeAll" {
1682 test_server = try Server(TestHandler).init(global_test_allocator.allocator(), .{
1683 .port = 9292,
1684 .address = "127.0.0.1",
1685 });
1686 test_thread = try test_server.listenInNewThread({});
1687}
1688
1689test "tests:afterAll" {
1690 test_server.stop();
1691 test_thread.join();
1692 test_server.deinit();
1693 try t.expectEqual(false, global_test_allocator.detectLeaks());
1694}
1695
1696test "Server: invalid handshake" {
1697 const stream = try testStream(false);
1698 defer stream.close();
1699
1700 try stream.writeAll("GET / HTTP/1.1\r\n\r\n");
1701 var buf: [1024]u8 = undefined;
1702 var pos: usize = 0;
1703 while (pos < buf.len) {
1704 const n = try stream.read(buf[0..]);
1705 if (n == 0) {
1706 break;
1707 }
1708 pos += n;
1709 } else {
1710 unreachable;
1711 }
1712
1713 try t.expectString("HTTP/1.1 400 \r\nConnection: Close\r\nError: missingheaders\r\nContent-Length: 0\r\n\r\n", buf[0..pos]);
1714}
1715
1716test "Server: read and write" {
1717 const stream = try testStream(true);
1718 defer stream.close();
1719
1720 try stream.writeAll(&proto.frame(.text, "over"));
1721 var buf: [12]u8 = undefined;
1722 _ = try stream.readAtLeast(&buf, 6);
1723 try t.expectSlice(u8, &.{129, 4, '9', '0', '0', '0'}, buf[0..6]);
1724}
1725
1726test "Server: handleMessage allocator" {
1727 const stream = try testStream(true);
1728 defer stream.close();
1729
1730 try stream.writeAll(&proto.frame(.text, "dyn"));
1731 var buf: [12]u8 = undefined;
1732 _ = try stream.readAtLeast(&buf, 12);
1733 try t.expectSlice(u8, &.{129, 10, 'o', 'v', 'e', 'r', ' ', '9', '0', '0', '0', '!'}, buf[0..12]);
1734}
1735
1736// Same as above, but client doesn't shutdown the connection
1737// When afterAll is runs and things are shutdown, this should still be properly cleaned up
1738test "Server: dirty handleMessage allocator" {
1739 const stream = try testStream(true);
1740
1741 try stream.writeAll(&proto.frame(.text, "dyn"));
1742 var buf: [12]u8 = undefined;
1743 _ = try stream.readAtLeast(&buf, 12);
1744 try t.expectSlice(u8, &.{129, 10, 'o', 'v', 'e', 'r', ' ', '9', '0', '0', '0', '!'}, buf[0..12]);
1745}
1746
1747test "Conn: close" {
1748 {
1749 // plain close
1750 const stream = try testStream(true);
1751 defer stream.close();
1752 try stream.writeAll(&proto.frame(.text, "close1"));
1753 var buf: [4]u8 = undefined;
1754 _ = try stream.readAtLeast(&buf, 4);
1755 try t.expectSlice(u8, &.{136, 2, 3, 232}, buf[0..4]);
1756 }
1757
1758 {
1759 // close with code
1760 const stream = try testStream(true);
1761 defer stream.close();
1762 try stream.writeAll(&proto.frame(.text, "close2"));
1763 var buf: [4]u8 = undefined;
1764 _ = try stream.readAtLeast(&buf, 4);
1765 try t.expectSlice(u8, &.{136, 2, 0, 0x7b}, buf[0..4]);
1766 }
1767
1768 {
1769 // close with reason
1770 const stream = try testStream(true);
1771 defer stream.close();
1772 try stream.writeAll(&proto.frame(.text, "close3"));
1773 var buf: [7]u8 = undefined;
1774 _ = try stream.readAtLeast(&buf, 7);
1775 try t.expectSlice(u8, &.{136, 5, 0, 0xea, 'b', 'y', 'e'}, buf[0..7]);
1776 }
1777}
1778
1779fn testStream(handshake: bool) !net.Stream {
1780 const timeout = std.mem.toBytes(std.posix.timeval{.sec = 0, .usec = 20_000});
1781 const address = try std.net.Address.parseIp("127.0.0.1", 9292);
1782 const stream = try std.net.tcpConnectToAddress(address);
1783 try std.posix.setsockopt(stream.handle, std.posix.SOL.SOCKET, std.posix.SO.RCVTIMEO, &timeout);
1784 try std.posix.setsockopt(stream.handle, std.posix.SOL.SOCKET, std.posix.SO.SNDTIMEO, &timeout);
1785
1786 if (handshake == false) {
1787 return stream;
1788 }
1789
1790 try stream.writeAll("GET / HTTP/1.1\r\ncontent-length: 0\r\nupgrade: websocket\r\nsec-websocket-version: 13\r\nconnection: upgrade\r\nsec-websocket-key: my-key\r\n\r\n");
1791 var buf: [1024]u8 = undefined;
1792 var pos: usize = 0;
1793 while (pos < buf.len) {
1794 const n = try stream.read(buf[0..]);
1795 if (n == 0) break;
1796
1797 pos += n;
1798 if (std.mem.endsWith(u8, buf[0..pos], "\r\n\r\n")) {
1799 break;
1800 }
1801 } else {
1802 unreachable;
1803 }
1804
1805 try t.expectString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: upgrade\r\nSec-Websocket-Accept: L8KGBs4w2MNLLzhfzlVoM0scCIE=\r\n\r\n", buf[0..pos]);
1806
1807 return stream;
1808}
1809
1810const TestHandler = struct {
1811 conn: *Conn,
1812
1813 pub fn init(_: Handshake, conn: *Conn, _: void) !TestHandler {
1814 return .{
1815 .conn = conn,
1816 };
1817 }
1818 pub fn clientMessage(self: *TestHandler, allocator: Allocator, data: []const u8,) !void {
1819 if (std.mem.eql(u8, data, "over")) {
1820 return self.conn.writeText("9000");
1821 }
1822 if (std.mem.eql(u8, data, "dyn")) {
1823 return self.conn.writeText(try std.fmt.allocPrint(allocator, "over {d}!", .{9000}));
1824 }
1825 if (std.mem.eql(u8, data, "close1")) {
1826 return self.conn.writeClose();
1827 }
1828 if (std.mem.eql(u8, data, "close2")) {
1829 return self.conn.writeCloseWithCode(123);
1830 }
1831 if (std.mem.eql(u8, data, "close3")) {
1832 return self.conn.writeCloseWithReason(234, "bye");
1833 }
1834 }
1835};
1836
1837test "List" {
1838 var list = List(TestNode){};
1839 try expectList(&.{}, list);
1840
1841 var n1 = TestNode{ .id = 1 };
1842 list.insert(&n1);
1843 try expectList(&.{1}, list);
1844
1845 list.remove(&n1);
1846 try expectList(&.{}, list);
1847
1848 var n2 = TestNode{ .id = 2 };
1849 list.insert(&n2);
1850 list.insert(&n1);
1851 try expectList(&.{ 2, 1 }, list);
1852
1853 var n3 = TestNode{ .id = 3 };
1854 list.insert(&n3);
1855 try expectList(&.{ 2, 1, 3 }, list);
1856
1857 list.remove(&n1);
1858 try expectList(&.{ 2, 3 }, list);
1859
1860 list.insert(&n1);
1861 try expectList(&.{ 2, 3, 1 }, list);
1862
1863 list.remove(&n2);
1864 try expectList(&.{ 3, 1 }, list);
1865
1866 list.remove(&n1);
1867 try expectList(&.{3}, list);
1868
1869 list.remove(&n3);
1870 try expectList(&.{}, list);
1871}
1872
1873const TestNode = struct {
1874 id: i32,
1875 next: ?*TestNode = null,
1876 prev: ?*TestNode = null,
1877};
1878
1879fn expectList(expected: []const i32, list: List(TestNode)) !void {
1880 if (expected.len == 0) {
1881 try t.expectEqual(null, list.head);
1882 try t.expectEqual(null, list.tail);
1883 return;
1884 }
1885
1886 var i: usize = 0;
1887 var next = list.head;
1888 while (next) |node| {
1889 try t.expectEqual(expected[i], node.id);
1890 i += 1;
1891 next = node.next;
1892 }
1893 try t.expectEqual(expected.len, i);
1894
1895 i = expected.len;
1896 var prev = list.tail;
1897 while (prev) |node| {
1898 i -= 1;
1899 try t.expectEqual(expected[i], node.id);
1900 prev = node.prev;
1901 }
1902 try t.expectEqual(0, i);
1903}