forked from
karlseguin.tngl.sh/websocket.zig
A websocket implementation for zig
add handshake.res_headers
Karl Seguin
(May 13, 2025, 10:01 AM +0800)
ba6068ff
c1c53b06
+112
-35
+15
-4
readme.md
+15
-4
readme.md
···
40
40
conn: *ws.Conn,
41
41
42
42
// You must define a public init function which takes
43
-
pub fn init(h: ws.Handshake, conn: *ws.Conn, app: *App) !Handler {
43
+
pub fn init(h: *ws.Handshake, conn: *ws.Conn, app: *App) !Handler {
44
44
// `h` contains the initial websocket "handshake" request
45
45
// It can be used to apply application-specific logic to verify / allow
46
46
// the connection (e.g. valid url, query string parameters, or headers)
···
71
71
When you create a `websocket.Server(Handler)`, the specified `Handler` is your structure which will receive messages. It must have a public `init` function and `clientMessage` method. Other methods, such as `close` can optionally be defined.
72
72
73
73
### init
74
-
The `init` method is called with a `websocket.Handshake`, a `*websocket.Conn` and whatever app-specific value was passed into `Server(H).init`.
74
+
The `init` method is called with a `*websocket.Handshake`, a `*websocket.Conn` and whatever app-specific value was passed into `Server(H).init`.
75
75
76
76
When `init` is called, the handshake response has not yet been sent to the client (this allows your `init` method to return an error which will cause websocket.zig to send an error response and close the connection). As such, you should not use/write to the `*websocket.Conn` at this point. Instead, use the `afterInit` method, described next.
77
77
78
78
The websocket specification requires the initial "handshake" to contain certain headers and values. The library validates these headers. However applications may have additional requirements before allowing the connection to be "upgraded" to a websocket connection. For example, a one-time-use token could be required in the querystring. Applications should use the provided `websocket.Handshake` to apply any application-specific verification and optionally return an error to terminate the connection.
79
79
80
-
The `websocket.Handshake` exposes the following fields:
80
+
The `*websocket.Handshake` exposes the following fields:
81
81
82
82
* `url: []const u8` - URL of the request in its original casing
83
83
* `method: []const u8` - Method of the request in its original casing
···
85
85
86
86
If you set the `max_headers` configuration value to > 0, then you can use `req.headers.get("HEADER_NAME")` to extract a header value from the given name:
87
87
88
+
If you set the `max_res_headers` configuration value to > 0, then you can set headers to be sent in the handshake response:
89
+
90
+
```zig
91
+
pub fn init(h: *ws.Handshake, conn: *ws.Conn, app: *App) !Handler {
92
+
h.res_headers.add("set-cookie", "delicious")
93
+
//...
94
+
}
95
+
```
96
+
97
+
Note that, currently, the total length of the headers added to `res_headers` should not exceed 1024 characters, else you will exeperience an out of bounds segfault.
98
+
88
99
```zig
89
100
// the last parameter, an *App in this case, is an application-specific
90
101
// value that you passed into server.listen()
91
-
pub fn init(h: websocket.Handshake, conn: websocket.Conn, app: *App) !Handler {
102
+
pub fn init(h: *websocket.Handshake, conn: websocket.Conn, app: *App) !Handler {
92
103
// get returns a ?[]const u8
93
104
// the name is lowercase
94
105
// the value is in its original case
+88
-25
src/server/handshake.zig
+88
-25
src/server/handshake.zig
···
21
21
key: []const u8,
22
22
method: []const u8,
23
23
headers: *KeyValue,
24
+
res_headers: KeyValue,
24
25
raw_header: []const u8,
25
26
compression: ?Compression,
26
27
···
47
48
return error.InvalidProtocol;
48
49
}
49
50
50
-
var headers = &state.headers;
51
+
var headers = &state.req_headers;
51
52
52
53
var key: []const u8 = "";
53
54
var required_headers: u8 = 0;
···
116
117
.method = method,
117
118
.headers = headers,
118
119
.compression = compression,
120
+
.res_headers = state.res_headers,
119
121
.raw_header = request[request_line_end + 2 .. request_length + 2],
120
122
};
121
123
}
122
124
123
-
pub fn createReply(key: []const u8, compression: ?websocket.Compression, buf: []u8) []const u8 {
125
+
pub fn createReply(key: []const u8, headers: *const KeyValue, compression: ?websocket.Compression, buf: []u8) ![]const u8 {
124
126
const HEADER =
125
127
"HTTP/1.1 101 Switching Protocols\r\n" ++
126
128
"Upgrade: websocket\r\n" ++
···
165
167
}
166
168
}
167
169
170
+
for (headers.keys[0..headers.len], headers.values[0..headers.len]) |k, v| {
171
+
pos += (try std.fmt.bufPrint(buf[pos..], "\r\n{s}: {s}", .{k, v})).len;
172
+
}
173
+
168
174
const end = pos + 4;
169
175
@memcpy(buf[pos..end], "\r\n\r\n");
170
176
return buf[0..end];
···
224
230
// a buffer to read data into
225
231
buf: []u8,
226
232
227
-
// Headers
228
-
headers: KeyValue,
233
+
// Headers from the request
234
+
req_headers: KeyValue,
235
+
236
+
// Headers that we want to send in the response
237
+
res_headers: KeyValue,
229
238
230
239
pool: *M.Pool,
231
240
···
234
243
const buf = try allocator.alloc(u8, pool.buffer_size);
235
244
errdefer allocator.free(buf);
236
245
237
-
const headers = try KeyValue.init(allocator, pool.max_headers);
238
-
errdefer headers.deinit(allocator);
246
+
const req_headers = try KeyValue.init(allocator, pool.max_req_headers);
247
+
errdefer req_headers.deinit(allocator);
248
+
249
+
const res_headers = try KeyValue.init(allocator, pool.max_res_headers);
250
+
errdefer res_headers.deinit(allocator);
239
251
240
252
return .{
241
253
.buf = buf,
242
254
.pool = pool,
243
-
.headers = headers,
255
+
.req_headers = req_headers,
256
+
.res_headers = res_headers,
244
257
};
245
258
}
246
259
247
260
fn deinit(self: *State) void {
248
261
const allocator = self.pool.allocator;
249
262
allocator.free(self.buf);
250
-
self.headers.deinit(allocator);
263
+
self.req_headers.deinit(allocator);
264
+
self.res_headers.deinit(allocator);
251
265
}
252
266
253
267
pub fn release(self: *State) void {
254
268
self.len = 0;
255
-
self.headers.len = 0;
269
+
self.req_headers.len = 0;
270
+
self.res_headers.len = 0;
256
271
self.pool.release(self);
257
272
}
258
273
};
···
282
297
allocator.free(self.values);
283
298
}
284
299
285
-
fn add(self: *KeyValue, key: []const u8, value: []const u8) void {
300
+
pub fn add(self: *KeyValue, key: []const u8, value: []const u8) void {
286
301
const len = self.len;
287
302
var keys = self.keys;
288
303
if (len == keys.len) {
···
354
369
available: usize,
355
370
allocator: Allocator,
356
371
buffer_size: usize,
357
-
max_headers: usize,
372
+
max_req_headers: usize,
373
+
max_res_headers: usize,
358
374
states: []*Handshake.State,
359
375
360
-
pub fn init(allocator: Allocator, count: usize, buffer_size: usize, max_headers: usize) !*Pool {
376
+
pub fn init(allocator: Allocator, count: usize, buffer_size: usize, max_req_headers: usize, max_res_headers: usize) !*Pool {
361
377
const states = try allocator.alloc(*Handshake.State, count);
362
378
errdefer allocator.free(states);
363
379
···
369
385
.states = states,
370
386
.allocator = allocator,
371
387
.available = count,
372
-
.max_headers = max_headers,
373
388
.buffer_size = buffer_size,
389
+
.max_req_headers = max_req_headers,
390
+
.max_res_headers = max_res_headers,
374
391
};
375
392
376
393
for (0..count) |i| {
···
442
459
443
460
const t = @import("../t.zig");
444
461
test "handshake: parse" {
445
-
var pool = try Pool.init(t.allocator, 1, 512, 10);
462
+
var pool = try Pool.init(t.allocator, 1, 512, 10, 1);
446
463
defer pool.deinit();
447
464
448
465
{
···
509
526
510
527
test "handshake: reply" {
511
528
var buf: [512]u8 = undefined;
529
+
var res_headers = try KeyValue.init(t.allocator, 2);
530
+
defer res_headers.deinit(t.allocator);
531
+
512
532
{
513
533
// no compression
514
534
const expected =
···
517
537
"Connection: upgrade\r\n" ++
518
538
"Sec-Websocket-Accept: flzHu2DevQ2dSCSVqKSii5e9C2o=\r\n\r\n"
519
539
;
520
-
try t.expectString(expected, Handshake.createReply("this is my key", null, &buf));
540
+
try t.expectString(expected, try Handshake.createReply("this is my key", &res_headers, null, &buf));
521
541
}
522
542
523
543
{
···
529
549
"Sec-Websocket-Accept: flzHu2DevQ2dSCSVqKSii5e9C2o=\r\n" ++
530
550
"Sec-WebSocket-Extensions: permessage-deflate\r\n\r\n"
531
551
;
532
-
try t.expectString(expected, Handshake.createReply("this is my key", .{}, &buf));
552
+
try t.expectString(expected, try Handshake.createReply("this is my key", &res_headers, .{}, &buf));
533
553
}
534
554
535
555
{
···
541
561
"Sec-Websocket-Accept: flzHu2DevQ2dSCSVqKSii5e9C2o=\r\n" ++
542
562
"Sec-WebSocket-Extensions: permessage-deflate; server_no_context_takeover; client_no_context_takeover\r\n\r\n"
543
563
;
544
-
try t.expectString(expected, Handshake.createReply("this is my key", .{
564
+
try t.expectString(expected, try Handshake.createReply("this is my key", &res_headers, .{
565
+
.client_no_context_takeover = true,
566
+
.server_no_context_takeover = true,
567
+
}, &buf));
568
+
}
569
+
570
+
// With custom headers
571
+
res_headers.add("Set-Cookie", "Yummy!");
572
+
{
573
+
// no compression
574
+
const expected =
575
+
"HTTP/1.1 101 Switching Protocols\r\n" ++
576
+
"Upgrade: websocket\r\n" ++
577
+
"Connection: upgrade\r\n" ++
578
+
"Sec-Websocket-Accept: flzHu2DevQ2dSCSVqKSii5e9C2o=\r\n" ++
579
+
"Set-Cookie: Yummy!\r\n\r\n"
580
+
;
581
+
try t.expectString(expected, try Handshake.createReply("this is my key", &res_headers, null, &buf));
582
+
}
583
+
584
+
{
585
+
// compression
586
+
const expected =
587
+
"HTTP/1.1 101 Switching Protocols\r\n" ++
588
+
"Upgrade: websocket\r\n" ++
589
+
"Connection: upgrade\r\n" ++
590
+
"Sec-Websocket-Accept: flzHu2DevQ2dSCSVqKSii5e9C2o=\r\n" ++
591
+
"Sec-WebSocket-Extensions: permessage-deflate\r\n" ++
592
+
"Set-Cookie: Yummy!\r\n\r\n"
593
+
;
594
+
try t.expectString(expected, try Handshake.createReply("this is my key", &res_headers, .{}, &buf));
595
+
}
596
+
597
+
{
598
+
// compression
599
+
const expected =
600
+
"HTTP/1.1 101 Switching Protocols\r\n" ++
601
+
"Upgrade: websocket\r\n" ++
602
+
"Connection: upgrade\r\n" ++
603
+
"Sec-Websocket-Accept: flzHu2DevQ2dSCSVqKSii5e9C2o=\r\n" ++
604
+
"Sec-WebSocket-Extensions: permessage-deflate; server_no_context_takeover; client_no_context_takeover\r\n" ++
605
+
"Set-Cookie: Yummy!\r\n\r\n"
606
+
;
607
+
try t.expectString(expected, try Handshake.createReply("this is my key", &res_headers, .{
545
608
.client_no_context_takeover = true,
546
609
.server_no_context_takeover = true,
547
610
}, &buf));
···
584
647
585
648
test "pool: acquire and release" {
586
649
// not 100% sure this is testing exactly what I want, but it's ....something ?
587
-
var p = try Pool.init(t.allocator, 2, 10, 3);
650
+
var p = try Pool.init(t.allocator, 2, 10, 3, 1);
588
651
defer p.deinit();
589
652
590
653
var hs1a = p.acquire() catch unreachable;
···
596
659
try t.expectEqual(10, hs1a.buf.len);
597
660
try t.expectEqual(10, hs2a.buf.len);
598
661
try t.expectEqual(10, hs3a.buf.len);
599
-
try t.expectEqual(0, hs1a.headers.len);
600
-
try t.expectEqual(0, hs2a.headers.len);
601
-
try t.expectEqual(0, hs3a.headers.len);
602
-
try t.expectEqual(3, hs1a.headers.keys.len);
603
-
try t.expectEqual(3, hs2a.headers.keys.len);
604
-
try t.expectEqual(3, hs3a.headers.keys.len);
662
+
try t.expectEqual(0, hs1a.req_headers.len);
663
+
try t.expectEqual(0, hs2a.req_headers.len);
664
+
try t.expectEqual(0, hs3a.req_headers.len);
665
+
try t.expectEqual(3, hs1a.req_headers.keys.len);
666
+
try t.expectEqual(3, hs2a.req_headers.keys.len);
667
+
try t.expectEqual(3, hs3a.req_headers.keys.len);
605
668
606
669
p.release(hs1a);
607
670
···
614
677
}
615
678
616
679
test "Handshake.Pool: threadsafety" {
617
-
var p = try Pool.init(t.allocator, 4, 10, 2);
680
+
var p = try Pool.init(t.allocator, 4, 10, 2, 2);
618
681
defer p.deinit();
619
682
620
683
for (p.states) |hs| {
+8
-5
src/server/server.zig
+8
-5
src/server/server.zig
···
70
70
timeout: u32 = 10,
71
71
max_size: ?u16 = null,
72
72
max_headers: ?u16 = null,
73
+
max_res_headers: ?u16 = null,
73
74
count: ?u16 = null,
74
75
};
75
76
···
1022
1023
const handshake_pool_count = config.handshake.count orelse 32;
1023
1024
const handshake_max_size = config.handshake.max_size orelse 1024;
1024
1025
const handshake_max_headers = config.handshake.max_headers orelse 10;
1026
+
const handshake_max_res_headers = config.handshake.max_res_headers orelse 2;
1025
1027
1026
-
var handshake_pool = try Handshake.Pool.init(allocator, handshake_pool_count, handshake_max_size, handshake_max_headers);
1028
+
var handshake_pool = try Handshake.Pool.init(allocator, handshake_pool_count, handshake_max_size, handshake_max_headers, handshake_max_res_headers);
1027
1029
errdefer handshake_pool.deinit();
1028
1030
1029
1031
const max_message_size = config.max_message_size orelse DEFAULT_MAX_MESSAGE_SIZE;
···
1549
1551
// After this, the app has access to &hc.conn, so any access to the
1550
1552
// conn has to be synchronized (which the conn does internally).
1551
1553
1552
-
const handler = H.init(handshake, conn, ctx) catch |err| {
1554
+
const handler = H.init(&handshake, conn, ctx) catch |err| {
1553
1555
if (comptime std.meta.hasFn(H, "handshakeErrorResponse")) {
1554
1556
preHandOffWrite(H.handshakeErrorResponse(err));
1555
1557
} else {
···
1561
1563
1562
1564
hc.handler = handler;
1563
1565
1564
-
var reply_buf: [512]u8 = undefined;
1565
-
try conn.writeFramed(Handshake.createReply(handshake.key, agreed_compression, &reply_buf));
1566
+
var reply_buf: [2048]u8 = undefined;
1567
+
const handshake_reply = try Handshake.createReply(handshake.key, &handshake.res_headers, agreed_compression, &reply_buf);
1568
+
try conn.writeFramed(handshake_reply);
1566
1569
1567
1570
if (comptime std.meta.hasFn(H, "afterInit")) {
1568
1571
const params = @typeInfo(@TypeOf(H.afterInit)).@"fn".params;
···
1973
1976
const TestHandler = struct {
1974
1977
conn: *Conn,
1975
1978
1976
-
pub fn init(h: Handshake, conn: *Conn, _: void) !TestHandler {
1979
+
pub fn init(h: *const Handshake, conn: *Conn, _: void) !TestHandler {
1977
1980
try t.expectString("upgrade", h.headers.get("connection").?);
1978
1981
return .{
1979
1982
.conn = conn,