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