zzstoatzz.io
websocket
48 kB
1239 lines
1const std = @import("std");
2
3const libc = std.c;
4const posix = std.posix;
5
6// 0.16: std.Thread.sleep removed, use libc nanosleep
7fn sleep(ns: u64) void {
8 const secs = ns / std.time.ns_per_s;
9 const nsecs = ns % std.time.ns_per_s;
10 var ts: libc.timespec = .{
11 .sec = @intCast(secs),
12 .nsec = @intCast(nsecs),
13 };
14 _ = libc.nanosleep(&ts, null);
15}
16const ascii = std.ascii;
17const Allocator = std.mem.Allocator;
18const websocket = @import("../websocket.zig");
19
20const M = @This();
21
22const SpecialHeader = enum {
23 none,
24 upgrade,
25 connection,
26 @"sec-websocket-key",
27 @"sec-websocket-version",
28 @"sec-websocket-extensions",
29};
30
31pub const Handshake = struct {
32 url: []const u8,
33 key: []const u8,
34 method: []const u8,
35 headers: *KeyValue,
36 res_headers: *KeyValue,
37 raw_header: []const u8,
38 compression: ?Compression,
39
40 pub const Pool = M.Pool;
41
42 pub const Compression = struct {
43 client_no_context_takeover: bool,
44 server_no_context_takeover: bool,
45 };
46
47 /// Incrementally parse an HTTP/1.1 request from `state.buf[0..state.len]`.
48 ///
49 /// Return values:
50 /// - `null` — more data needed to make a decision
51 /// - `Some(Handshake)` — valid RFC 6455 upgrade request
52 /// - `error.MissingHeaders` — well-formed HTTP/1.1 request that is NOT a
53 /// websocket upgrade. Held back until the request body (per
54 /// Content-Length) has fully arrived, so the worker's `httpFallback`
55 /// sees a complete request rather than a truncated one.
56 /// - other named errors — request is malformed beyond recovery; the
57 /// worker responds 400.
58 ///
59 /// Streaming: end-of-headers is detected by `std.http.HeadParser`, a
60 /// stdlib SIMD state machine that handles every TCP-fragmentation
61 /// boundary correctly (including splits *inside* CRLF and CRLFCRLF
62 /// sequences). The parser feeds only new bytes per call, so repeated
63 /// calls as the buffer grows are O(new_bytes), not O(buf).
64 ///
65 /// RFC compliance:
66 /// - whitespace between header name and ':' is rejected per
67 /// RFC 7230 §3.2 (request smuggling vector)
68 /// - simultaneous Transfer-Encoding and Content-Length is rejected
69 /// per RFC 7230 §3.3 (request smuggling vector)
70 /// - obsolete line folding (RFC 7230 §3.2.4) is not supported;
71 /// a folded continuation line will surface as InvalidHeader
72 /// - request bodies with Transfer-Encoding only (no Content-Length)
73 /// are passed through to httpFallback without the worker
74 /// waiting for body completeness — chunked decoding is the
75 /// fallback handler's responsibility, not this parser's.
76 ///
77 /// Crash discipline (per Zig zen "runtime crashes are better than bugs"):
78 /// four invariant-protected paths use `unreachable` so a contract
79 /// violation surfaces as a debuggable panic rather than a silent
80 /// misclassification of valid input as malformed. all other error
81 /// paths correspond to genuinely malformed external input.
82 pub fn parse(state: *State) !?Handshake {
83 // Phase 1: streaming detection of end-of-headers via std.http.HeadParser.
84 // feed only NEW bytes since last call so the SIMD scan is amortized.
85 const newly_arrived = state.buf[state.head_fed..state.len];
86 state.head_fed += state.head_parser.feed(newly_arrived);
87 if (state.head_parser.state != .finished) return null;
88
89 // Phase 2: parse the bounded request.
90 // body_start = head_fed = position of first body byte (just past CRLFCRLF).
91 // Structure:
92 // REQUEST_LINE \r\n
93 // HEADER_LINE \r\n (zero or more)
94 // \r\n (blank line)
95 // [body]
96 const body_start = state.head_fed;
97 // HeadParser .finished means it consumed at least the 4-byte CRLFCRLF.
98 if (body_start < 4) unreachable;
99
100 // Any buffer of length ≥4 ending in CRLFCRLF contains at least one CRLF.
101 const first_crlf = std.mem.indexOf(u8, state.buf[0..body_start], "\r\n") orelse unreachable;
102 const request_line = state.buf[0..first_crlf];
103
104 if (!ascii.endsWithIgnoreCase(request_line, "http/1.1")) {
105 return error.InvalidProtocol;
106 }
107
108 // headers accumulator. parse may be called multiple times on the
109 // same state (e.g. while waiting for body bytes), so reset to keep
110 // re-runs idempotent.
111 var headers = &state.req_headers;
112 headers.len = 0;
113
114 var key: []const u8 = "";
115 var required_headers: u8 = 0;
116 var compression: ?Handshake.Compression = null;
117
118 // Errors that signal "this is not a websocket upgrade, dispatch via
119 // httpFallback" are *deferred* rather than returned inline. We need
120 // to finish parsing the headers regardless so the body-completeness
121 // check below sees `Content-Length`. Otherwise a fragmented POST
122 // (headers in one TCP read, body in the next) whose request also
123 // contains `Connection: keep-alive` or `Upgrade: ...` (non-websocket)
124 // would dispatch with a truncated body — handlers would observe
125 // body.len == 0 even though Content-Length was non-zero.
126 const DeferredFallbackErr = error{ InvalidUpgrade, InvalidConnection };
127 var deferred_fallback_err: ?DeferredFallbackErr = null;
128
129 // header lines occupy [first_crlf + 2 .. body_start - 2], excluding
130 // the blank line's CRLF. each line ends in CRLF (RFC 7230 §3.2.4
131 // line folding deprecated; we don't support it).
132 var rest = state.buf[first_crlf + 2 .. body_start - 2];
133 while (rest.len > 0) {
134 // rest is bounded by HeadParser's CRLFCRLF detection: every
135 // line in it ends in CRLF, so indexOf must find one.
136 const crlf = std.mem.indexOf(u8, rest, "\r\n") orelse unreachable;
137 const line = rest[0..crlf];
138 rest = rest[crlf + 2 ..];
139
140 // HeadParser stops at the FIRST empty line, so an empty line
141 // mid-iteration would mean rest was constructed wrong.
142 if (line.len == 0) unreachable;
143
144 // Leading whitespace = obsolete line folding (RFC 7230 §3.2.4)
145 // OR malformed header. Both are rejected.
146 if (ascii.isWhitespace(line[0])) return error.InvalidHeader;
147
148 const colon = std.mem.indexOfScalar(u8, line, ':') orelse return error.InvalidHeader;
149 if (colon == 0) return error.InvalidHeader; // empty name
150
151 // RFC 7230 §3.2: "No whitespace is allowed between the header
152 // field-name and colon. Servers MUST reject ... with a 400."
153 // This is a request-smuggling defense, not stylistic.
154 if (ascii.isWhitespace(line[colon - 1])) return error.WhitespaceBeforeColon;
155
156 const name = toLower(line[0..colon]);
157 const value = std.mem.trim(u8, line[colon + 1 ..], &ascii.whitespace);
158
159 headers.add(name, value);
160 switch (std.meta.stringToEnum(SpecialHeader, name) orelse .none) {
161 .upgrade => {
162 if (!ascii.eqlIgnoreCase("websocket", value)) {
163 // Defer — keep parsing to find Content-Length.
164 if (deferred_fallback_err == null) deferred_fallback_err = error.InvalidUpgrade;
165 } else {
166 required_headers |= 1;
167 }
168 },
169 .connection => {
170 // Connection: keep-alive, Upgrade — the spec allows multiple tokens
171 if (std.ascii.indexOfIgnoreCase(value, "upgrade") == null) {
172 // Defer — keep parsing to find Content-Length.
173 if (deferred_fallback_err == null) deferred_fallback_err = error.InvalidConnection;
174 } else {
175 required_headers |= 4;
176 }
177 },
178 .@"sec-websocket-key" => {
179 key = value;
180 required_headers |= 8;
181 },
182 .@"sec-websocket-version" => {
183 if (value.len != 2 or value[0] != '1' or value[1] != '3') {
184 return error.InvalidVersion;
185 }
186 required_headers |= 2;
187 },
188 .@"sec-websocket-extensions" => compression = try parseExtension(value),
189 .none => {},
190 }
191 }
192
193 // RFC 7230 §3.3: simultaneous Transfer-Encoding and Content-Length is
194 // a request-smuggling vector. The two headers can disagree about
195 // body length, letting an attacker hide a second request inside the
196 // first. Reject before any body-completeness logic runs.
197 if (headers.get("transfer-encoding") != null and headers.get("content-length") != null) {
198 return error.AmbiguousBodyLength;
199 }
200
201 // Single body-completeness gate covering BOTH `MissingHeaders` (no
202 // websocket headers at all) and `InvalidConnection`/`InvalidUpgrade`
203 // (looks like a normal HTTP/1.1 request with `Connection: keep-alive`
204 // or some non-websocket Upgrade token). Both paths feed httpFallback
205 // and both need a complete body in the buffer.
206 if (deferred_fallback_err != null or required_headers != 15) {
207 if (headers.get("content-length")) |cl_str| {
208 const content_length = std.fmt.parseInt(usize, cl_str, 10) catch {
209 return deferred_fallback_err orelse error.MissingHeaders;
210 };
211 const body_available = state.len - body_start;
212 if (body_available < content_length) {
213 // body not yet complete — ask caller for more data
214 return null;
215 }
216 }
217 return deferred_fallback_err orelse error.MissingHeaders;
218 }
219
220 // request line: METHOD <space> URL <space> HTTP/1.1
221 const sp = std.mem.indexOfScalar(u8, request_line, ' ') orelse return error.InvalidRequestLine;
222 const method = request_line[0..sp];
223 // url is between method-end and " HTTP/1.1" suffix (9 trailing chars)
224 if (request_line.len < sp + 1 + 9) return error.InvalidRequestLine;
225 const url = std.mem.trim(u8, request_line[sp + 1 .. request_line.len - 9], &ascii.whitespace);
226
227 return .{
228 .key = key,
229 .url = url,
230 .method = method,
231 .headers = headers,
232 .compression = compression,
233 .res_headers = &state.res_headers,
234 // raw_header preserves original byte-exact slice. matches the
235 // pre-rewrite contract: starts after the request line's CRLF
236 // and ends just past the last header's CRLF (i.e. excludes
237 // the blank line's CRLF). callers re-emit this verbatim.
238 .raw_header = state.buf[first_crlf + 2 .. body_start - 2],
239 };
240 }
241
242 pub fn createReply(key: []const u8, headers_: ?*KeyValue, compression: bool, buf: []u8) ![]const u8 {
243 return createReplyWithCompression(key, headers_, if (compression) .{
244 .server_no_context_takeover = true,
245 .client_no_context_takeover = true,
246 } else null, buf);
247 }
248
249 pub fn createReplyNegotiated(key: []const u8, headers_: ?*KeyValue, compression: ?Compression, buf: []u8) ![]const u8 {
250 return createReplyWithCompression(key, headers_, compression, buf);
251 }
252
253 fn createReplyWithCompression(key: []const u8, headers_: ?*KeyValue, compression: ?Compression, buf: []u8) ![]const u8 {
254 const HEADER =
255 "HTTP/1.1 101 Switching Protocols\r\n" ++
256 "Upgrade: websocket\r\n" ++
257 "Connection: upgrade\r\n" ++
258 "Sec-Websocket-Accept: ";
259
260 @memcpy(buf[0..HEADER.len], HEADER);
261 var pos = HEADER.len;
262
263 {
264 var h: [20]u8 = undefined;
265 var hasher = std.crypto.hash.Sha1.init(.{});
266 hasher.update(key);
267 // websocket spec always used this value
268 hasher.update("258EAFA5-E914-47DA-95CA-C5AB0DC85B11");
269 hasher.final(&h);
270
271 const end = pos + 28;
272 _ = std.base64.standard.Encoder.encode(buf[pos..end], h[0..]);
273 pos = end;
274 }
275
276 if (compression) |negotiated| {
277 const extension = "\r\nSec-WebSocket-Extensions: permessage-deflate";
278 @memcpy(buf[pos..][0..extension.len], extension);
279 pos += extension.len;
280 if (negotiated.server_no_context_takeover) {
281 const param = "; server_no_context_takeover";
282 @memcpy(buf[pos..][0..param.len], param);
283 pos += param.len;
284 }
285 if (negotiated.client_no_context_takeover) {
286 const param = "; client_no_context_takeover";
287 @memcpy(buf[pos..][0..param.len], param);
288 pos += param.len;
289 }
290 }
291
292 if (headers_) |headers| {
293 for (headers.keys[0..headers.len], headers.values[0..headers.len]) |k, v| {
294 pos += (try std.fmt.bufPrint(buf[pos..], "\r\n{s}: {s}", .{ k, v })).len;
295 }
296 }
297
298 const end = pos + 4;
299 @memcpy(buf[pos..end], "\r\n\r\n");
300 return buf[0..end];
301 }
302
303 pub fn parseExtension(value: []const u8) !?Handshake.Compression {
304 var offers = std.mem.splitScalar(u8, value, ',');
305 while (offers.next()) |offer| {
306 var params = std.mem.splitScalar(u8, offer, ';');
307 if (!std.mem.eql(u8, std.mem.trim(u8, params.first(), &ascii.whitespace), "permessage-deflate"))
308 continue;
309
310 var client_no_context_takeover = false;
311 var server_no_context_takeover = false;
312 var valid = true;
313 while (params.next()) |param_| {
314 const param = std.mem.trim(u8, param_, &ascii.whitespace);
315 if (std.mem.eql(u8, param, "client_no_context_takeover")) {
316 client_no_context_takeover = true;
317 } else if (std.mem.eql(u8, param, "server_no_context_takeover")) {
318 server_no_context_takeover = true;
319 } else if (std.mem.eql(u8, param, "client_max_window_bits") or
320 std.mem.startsWith(u8, param, "client_max_window_bits="))
321 {
322 // A larger decoder window can read a stream produced with
323 // any smaller client window.
324 } else if (!std.mem.eql(u8, param, "server_max_window_bits=15")) {
325 valid = false;
326 break;
327 }
328 }
329 if (!valid) continue;
330 return .{
331 .client_no_context_takeover = client_no_context_takeover,
332 .server_no_context_takeover = server_no_context_takeover,
333 };
334 }
335 return null;
336 }
337
338 // This is what we're pooling
339 pub const State = struct {
340 // length of data we have in buf
341 len: usize = 0,
342
343 // a buffer to read data into
344 buf: []u8,
345
346 // streaming end-of-headers detector. advances across multiple
347 // parse() calls as new bytes arrive. survives splits inside CRLF
348 // sequences (the recurring class of TCP-fragmentation bugs).
349 head_parser: std.http.HeadParser = .{},
350
351 // monotonically tracks how many bytes have been fed into
352 // head_parser. on each parse() call we feed only the new bytes
353 // [head_fed..len] so the SIMD scan is amortized.
354 head_fed: usize = 0,
355
356 // Headers from the request
357 req_headers: KeyValue,
358
359 // Headers that we want to send in the response
360 res_headers: KeyValue,
361
362 pool: *M.Pool,
363
364 fn init(pool: *M.Pool) !State {
365 const allocator = pool.allocator;
366 const buf = try allocator.alloc(u8, pool.buffer_size);
367 errdefer allocator.free(buf);
368
369 const req_headers = try Handshake.KeyValue.init(allocator, pool.max_req_headers);
370 errdefer req_headers.deinit(allocator);
371
372 const res_headers = try Handshake.KeyValue.init(allocator, pool.max_res_headers);
373 errdefer res_headers.deinit(allocator);
374
375 return .{
376 .buf = buf,
377 .pool = pool,
378 .req_headers = req_headers,
379 .res_headers = res_headers,
380 };
381 }
382
383 fn deinit(self: *State) void {
384 const allocator = self.pool.allocator;
385 allocator.free(self.buf);
386 self.req_headers.deinit(allocator);
387 self.res_headers.deinit(allocator);
388 }
389
390 pub fn release(self: *State) void {
391 self.len = 0;
392 self.head_parser = .{};
393 self.head_fed = 0;
394 self.req_headers.len = 0;
395 self.res_headers.len = 0;
396 self.pool.release(self);
397 }
398 };
399
400 pub const KeyValue = struct {
401 len: usize,
402 keys: [][]const u8,
403 values: [][]const u8,
404
405 fn init(allocator: Allocator, max: usize) !KeyValue {
406 const keys = try allocator.alloc([]const u8, max);
407 errdefer allocator.free(keys);
408
409 const values = try allocator.alloc([]const u8, max);
410 errdefer allocator.free(values);
411
412 return .{
413 .len = 0,
414 .keys = keys,
415 .values = values,
416 };
417 }
418
419 fn deinit(self: *const KeyValue, allocator: Allocator) void {
420 allocator.free(self.keys);
421 allocator.free(self.values);
422 }
423
424 pub fn add(self: *KeyValue, key: []const u8, value: []const u8) void {
425 const len = self.len;
426 var keys = self.keys;
427 if (len == keys.len) {
428 return;
429 }
430
431 keys[len] = key;
432 self.values[len] = value;
433 self.len = len + 1;
434 }
435
436 pub fn get(self: *const KeyValue, needle: []const u8) ?[]const u8 {
437 const keys = self.keys[0..self.len];
438 loop: for (keys, 0..) |key, i| {
439 // This is largely a reminder to myself that std.mem.eql isn't
440 // particularly fast. Here we at least avoid the 1 extra ptr
441 // equality check that std.mem.eql does, but we could do better
442 // TODO: monitor https://github.com/ziglang/zig/issues/8689
443 if (needle.len != key.len) {
444 continue;
445 }
446 for (needle, key) |n, k| {
447 if (n != k) {
448 continue :loop;
449 }
450 }
451 return self.values[i];
452 }
453
454 return null;
455 }
456
457 pub fn iterator(self: *const KeyValue) Iterator {
458 const len = self.len;
459 return .{
460 .pos = 0,
461 .keys = self.keys[0..len],
462 .values = self.values[0..len],
463 };
464 }
465
466 pub const Iterator = struct {
467 pos: usize,
468 keys: [][]const u8,
469 values: [][]const u8,
470
471 const KV = struct {
472 key: []const u8,
473 value: []const u8,
474 };
475
476 pub fn next(self: *Iterator) ?KV {
477 const pos = self.pos;
478 if (pos == self.keys.len) {
479 return null;
480 }
481
482 self.pos = pos + 1;
483 return .{
484 .key = self.keys[pos],
485 .value = self.values[pos],
486 };
487 }
488 };
489 };
490};
491
492pub const Pool = struct {
493 mutex: std.Io.Mutex,
494 io: std.Io,
495 available: usize,
496 allocator: Allocator,
497 buffer_size: usize,
498 max_req_headers: usize,
499 max_res_headers: usize,
500 states: []*Handshake.State,
501
502 pub fn init(allocator: Allocator, count: usize, buffer_size: usize, max_req_headers: usize, max_res_headers: usize) !*Pool {
503 const states = try allocator.alloc(*Handshake.State, count);
504 errdefer allocator.free(states);
505
506 const pool = try allocator.create(Pool);
507 errdefer allocator.destroy(pool);
508
509 pool.* = .{
510 .mutex = .init,
511 .io = std.Options.debug_io,
512 .states = states,
513 .allocator = allocator,
514 .available = count,
515 .buffer_size = buffer_size,
516 .max_req_headers = max_req_headers,
517 .max_res_headers = max_res_headers,
518 };
519
520 for (0..count) |i| {
521 const state = try allocator.create(Handshake.State);
522 errdefer allocator.destroy(state);
523
524 state.* = try Handshake.State.init(pool);
525 states[i] = state;
526 }
527
528 return pool;
529 }
530
531 pub fn deinit(self: *Pool) void {
532 const allocator = self.allocator;
533 for (self.states) |s| {
534 s.deinit();
535 allocator.destroy(s);
536 }
537 allocator.free(self.states);
538 allocator.destroy(self);
539 }
540
541 pub fn acquire(self: *Pool) !*Handshake.State {
542 const states = self.states;
543 const io = self.io;
544
545 self.mutex.lockUncancelable(io);
546 const available = self.available;
547 if (available == 0) {
548 // dont hold the lock over factory
549 self.mutex.unlock(io);
550
551 const allocator = self.allocator;
552 const state = try allocator.create(Handshake.State);
553 errdefer allocator.destroy(state);
554 state.* = try Handshake.State.init(self);
555 return state;
556 }
557 const index = available - 1;
558 const state = states[index];
559 self.available = index;
560 self.mutex.unlock(io);
561 return state;
562 }
563
564 fn release(self: *Pool, state: *Handshake.State) void {
565 var states = self.states;
566 const io = self.io;
567
568 self.mutex.lockUncancelable(io);
569 const available = self.available;
570 if (available == states.len) {
571 self.mutex.unlock(io);
572 state.deinit();
573 self.allocator.destroy(state);
574 return;
575 }
576 states[available] = state;
577 self.available = available + 1;
578 self.mutex.unlock(io);
579 }
580};
581
582fn toLower(str: []u8) []u8 {
583 for (str, 0..) |c, i| {
584 str[i] = ascii.toLower(c);
585 }
586 return str;
587}
588
589const t = @import("../t.zig");
590test "handshake: parse" {
591 var pool = try Pool.init(t.allocator, 1, 512, 10, 1);
592 defer pool.deinit();
593
594 {
595 var state = try pool.acquire();
596 defer state.release();
597
598 try t.expectEqual(null, try testHandshake("", state));
599 try t.expectEqual(null, try testHandshake("GET", state));
600 try t.expectEqual(null, try testHandshake("GET 1 HTTP/1.0\r", state));
601 try t.expectEqual(null, try testHandshake("GET 1 HTTP/1.0\r\n", state));
602
603 try t.expectError(error.InvalidProtocol, testHandshake("GET / HTTP/1.0\r\n\r\n", state));
604 try t.expectError(error.MissingHeaders, testHandshake("GET / HTTP/1.1\r\n\r\n", state));
605 try t.expectError(error.MissingHeaders, testHandshake("GET / HTTP/1.1\r\nConnection: upgrade\r\n\r\n", state));
606 try t.expectError(error.MissingHeaders, testHandshake("GET / HTTP/1.1\r\nConnection: upgrade\r\nUpgrade: websocket\r\n\r\n", state));
607 try t.expectError(error.MissingHeaders, testHandshake("GET / HTTP/1.1\r\nConnection: upgrade\r\nUpgrade: websocket\r\nsec-websocket-version:13\r\n\r\n", state));
608 }
609
610 {
611 var state = try pool.acquire();
612 defer state.release();
613
614 const h = (try testHandshake("GET /test?a=1 HTTP/1.1\r\nConnection: upgrade\r\nUpgrade: websocket\r\nsec-websocket-version:13\r\nsec-websocket-key: 9000!\r\nCustom: Header-Value\r\n\r\n", state)).?;
615 try t.expectString("9000!", h.key);
616 try t.expectString("GET", h.method);
617 try t.expectString("/test?a=1", h.url);
618 try t.expectString("connection: upgrade\r\nupgrade: websocket\r\nsec-websocket-version:13\r\nsec-websocket-key: 9000!\r\ncustom: Header-Value\r\n", h.raw_header);
619 try t.expectString("Header-Value", h.headers.get("custom").?);
620
621 var it = h.headers.iterator();
622 {
623 const kv = it.next().?;
624 try t.expectString("connection", kv.key);
625 try t.expectString("upgrade", kv.value);
626 }
627
628 {
629 const kv = it.next().?;
630 try t.expectString("upgrade", kv.key);
631 try t.expectString("websocket", kv.value);
632 }
633
634 {
635 const kv = it.next().?;
636 try t.expectString("sec-websocket-version", kv.key);
637 try t.expectString("13", kv.value);
638 }
639
640 {
641 const kv = it.next().?;
642 try t.expectString("sec-websocket-key", kv.key);
643 try t.expectString("9000!", kv.value);
644 }
645
646 {
647 const kv = it.next().?;
648 try t.expectString("custom", kv.key);
649 try t.expectString("Header-Value", kv.value);
650 }
651
652 try t.expectEqual(null, it.next());
653 }
654}
655
656test "handshake: parse — every byte-split of a valid upgrade is equivalent" {
657 // Regression net for the recurring class of TCP-fragmentation bugs.
658 // For a known-good request, feeding it 1-byte-at-a-time, 2-bytes-at-a-time,
659 // etc. must produce the same final outcome as feeding it whole.
660 // This catches any future parser change that loses idempotence under
661 // partial reads — including splits inside CRLF or CRLFCRLF sequences.
662 var pool = try Pool.init(t.allocator, 1, 512, 10, 1);
663 defer pool.deinit();
664
665 const req = "GET /chat HTTP/1.1\r\nConnection: upgrade\r\nUpgrade: websocket\r\nsec-websocket-version:13\r\nsec-websocket-key: abc\r\n\r\n";
666
667 var chunk: usize = 1;
668 while (chunk <= req.len) : (chunk += 1) {
669 var state = try pool.acquire();
670 defer state.release();
671
672 const h = (try testHandshakeIncremental(req, chunk, state)) orelse {
673 std.debug.print("chunk={d}: parse returned null instead of Handshake\n", .{chunk});
674 return error.UnexpectedNull;
675 };
676 try t.expectString("abc", h.key);
677 try t.expectString("GET", h.method);
678 try t.expectString("/chat", h.url);
679 }
680}
681
682test "handshake: parse — split inside CRLF and CRLFCRLF sequences" {
683 // Pin specific torture-cases that have caused production outages:
684 // - split between '\r' and '\n' of the request-line CRLF
685 // - split between request-line's CRLF and the next header's first byte
686 // - split inside the final CRLFCRLF (each of the 4 internal positions)
687 var pool = try Pool.init(t.allocator, 1, 512, 10, 1);
688 defer pool.deinit();
689
690 const req = "GET / HTTP/1.1\r\nConnection: upgrade\r\nUpgrade: websocket\r\nsec-websocket-version:13\r\nsec-websocket-key: k\r\n\r\n";
691
692 // Try every two-chunk split. Each split must produce the same handshake.
693 var split: usize = 1;
694 while (split < req.len) : (split += 1) {
695 var state = try pool.acquire();
696 defer state.release();
697
698 state.head_parser = .{};
699 state.head_fed = 0;
700 state.req_headers.len = 0;
701 state.len = 0;
702
703 // first chunk
704 @memcpy(state.buf[0..split], req[0..split]);
705 state.len = split;
706 const r1 = try Handshake.parse(state);
707 try t.expectEqual(null, r1);
708
709 // second chunk
710 @memcpy(state.buf[split..req.len], req[split..]);
711 state.len = req.len;
712 const h = (try Handshake.parse(state)).?;
713 try t.expectString("k", h.key);
714 }
715}
716
717test "handshake: parse — malformed inputs return errors, never panic" {
718 // Adversarial / malformed inputs should map to named errors. There
719 // must be no `unreachable` reachable from any input; every parse
720 // exits via either Some(handshake), null, or a named error.
721 var pool = try Pool.init(t.allocator, 1, 512, 10, 1);
722 defer pool.deinit();
723
724 const Case = struct { input: []const u8, expected: anyerror };
725 const cases = [_]Case{
726 // request line missing space → InvalidRequestLine after protocol check
727 .{ .input = "GETONLY\r\n\r\n", .expected = error.InvalidProtocol },
728 .{ .input = "GETONLY HTTP/1.1\r\n\r\n", .expected = error.MissingHeaders },
729 // request line ends mid-protocol
730 .{ .input = "GET / HTTP/1\r\n\r\n", .expected = error.InvalidProtocol },
731 // header line with no colon
732 .{ .input = "GET / HTTP/1.1\r\nNoColonHere\r\n\r\n", .expected = error.InvalidHeader },
733 // header with empty name
734 .{ .input = "GET / HTTP/1.1\r\n: value\r\n\r\n", .expected = error.InvalidHeader },
735 // upgrade present but not "websocket"
736 .{ .input = "GET / HTTP/1.1\r\nConnection: upgrade\r\nUpgrade: nope\r\n\r\n", .expected = error.InvalidUpgrade },
737 // version not 13
738 .{ .input = "GET / HTTP/1.1\r\nConnection: upgrade\r\nUpgrade: websocket\r\nsec-websocket-version: 12\r\n\r\n", .expected = error.InvalidVersion },
739 // connection without "upgrade" token
740 .{ .input = "GET / HTTP/1.1\r\nConnection: keep-alive\r\nUpgrade: websocket\r\n\r\n", .expected = error.InvalidConnection },
741 };
742
743 for (cases) |c| {
744 var state = try pool.acquire();
745 defer state.release();
746 try t.expectError(c.expected, testHandshake(c.input, state));
747 }
748}
749
750test "handshake: parse — POST body byte-split equivalence" {
751 // The headline bug from b45400c: POST with body split across reads
752 // must hold null until the full body has arrived, then return
753 // MissingHeaders so the worker dispatches httpFallback with the
754 // complete body. Testing every possible split position guards
755 // against any future regression in the body-completeness path.
756 var pool = try Pool.init(t.allocator, 1, 1024, 10, 1);
757 defer pool.deinit();
758
759 const req = "POST /xrpc/com.atproto.sync.requestCrawl HTTP/1.1\r\nHost: relay\r\nContent-Type: application/json\r\nContent-Length: 27\r\n\r\n{\"hostname\":\"pds.test.com\"}";
760
761 var split: usize = 1;
762 while (split < req.len) : (split += 1) {
763 var state = try pool.acquire();
764 defer state.release();
765
766 state.head_parser = .{};
767 state.head_fed = 0;
768 state.req_headers.len = 0;
769 state.len = 0;
770
771 @memcpy(state.buf[0..split], req[0..split]);
772 state.len = split;
773 // first call: must NOT commit to MissingHeaders before body arrives
774 const r1 = Handshake.parse(state) catch |e| blk: {
775 // legitimate early errors (bad request line) are fine
776 if (e != error.InvalidProtocol and e != error.InvalidRequestLine) return e;
777 break :blk @as(?Handshake, null);
778 };
779 try t.expectEqual(null, r1);
780
781 @memcpy(state.buf[split..req.len], req[split..]);
782 state.len = req.len;
783 try t.expectError(error.MissingHeaders, Handshake.parse(state));
784 }
785}
786
787test "handshake: parse — RFC 7230 §3.2 rejects whitespace before colon" {
788 // "No whitespace is allowed between the header field-name and colon.
789 // Servers must reject ... with a 400." This is a request-smuggling
790 // defense, not stylistic.
791 var pool = try Pool.init(t.allocator, 1, 512, 10, 1);
792 defer pool.deinit();
793
794 const cases = [_][]const u8{
795 // single space before colon
796 "GET / HTTP/1.1\r\nHost : x\r\n\r\n",
797 // tab before colon
798 "GET / HTTP/1.1\r\nHost\t: x\r\n\r\n",
799 // also catches it on a non-required header
800 "GET / HTTP/1.1\r\nConnection: upgrade\r\nUpgrade: websocket\r\nsec-websocket-version: 13\r\nsec-websocket-key: x\r\nFoo : bar\r\n\r\n",
801 };
802
803 for (cases) |c| {
804 var state = try pool.acquire();
805 defer state.release();
806 try t.expectError(error.WhitespaceBeforeColon, testHandshake(c, state));
807 }
808}
809
810test "handshake: parse — RFC 7230 §3.3 rejects Transfer-Encoding + Content-Length" {
811 // Combination is a request-smuggling vector. RFC says either reject or
812 // remove Content-Length before forwarding; we choose explicit rejection
813 // because we can't safely forward what we don't trust.
814 var pool = try Pool.init(t.allocator, 1, 512, 10, 1);
815 defer pool.deinit();
816
817 {
818 var state = try pool.acquire();
819 defer state.release();
820 try t.expectError(
821 error.AmbiguousBodyLength,
822 testHandshake(
823 "POST /x HTTP/1.1\r\nHost: r\r\nContent-Length: 5\r\nTransfer-Encoding: chunked\r\n\r\nhello",
824 state,
825 ),
826 );
827 }
828
829 {
830 // also rejects when both appear on a (would-be) WS upgrade
831 var state = try pool.acquire();
832 defer state.release();
833 try t.expectError(
834 error.AmbiguousBodyLength,
835 testHandshake(
836 "GET / HTTP/1.1\r\nConnection: upgrade\r\nUpgrade: websocket\r\nsec-websocket-version: 13\r\nsec-websocket-key: x\r\nContent-Length: 0\r\nTransfer-Encoding: chunked\r\n\r\n",
837 state,
838 ),
839 );
840 }
841}
842
843test "handshake: parse — rejects obsolete line folding (RFC 7230 §3.2.4)" {
844 // "A sender MUST NOT generate a message that includes line folding."
845 // We surface a folded continuation as InvalidHeader rather than try
846 // to reconstruct the obs-fold semantics.
847 var pool = try Pool.init(t.allocator, 1, 512, 10, 1);
848 defer pool.deinit();
849
850 var state = try pool.acquire();
851 defer state.release();
852 // continuation line starts with whitespace — would be folded value of "Host"
853 try t.expectError(
854 error.InvalidHeader,
855 testHandshake("GET / HTTP/1.1\r\nHost: r\r\n continued\r\n\r\n", state),
856 );
857}
858
859test "handshake: parse — idempotent under repeated calls without new bytes" {
860 // Calling parse twice in a row with the same state.len must produce
861 // the same answer. Otherwise the worker can race with itself when
862 // it polls before another byte arrives.
863 var pool = try Pool.init(t.allocator, 1, 512, 10, 1);
864 defer pool.deinit();
865
866 var state = try pool.acquire();
867 defer state.release();
868
869 const req = "GET /a HTTP/1.1\r\nConnection: upgrade\r\nUpgrade: websocket\r\nsec-websocket-version:13\r\nsec-websocket-key: x\r\n\r\n";
870 @memcpy(state.buf[0..req.len], req);
871 state.len = req.len;
872
873 const h1 = (try Handshake.parse(state)).?;
874 const h2 = (try Handshake.parse(state)).?;
875 try t.expectString(h1.key, h2.key);
876 try t.expectString(h1.method, h2.method);
877 try t.expectString(h1.url, h2.url);
878}
879
880test "handshake: parse POST with body signals MissingHeaders (for httpFallback)" {
881 var pool = try Pool.init(t.allocator, 1, 512, 10, 1);
882 defer pool.deinit();
883
884 {
885 // POST whose body has arrived in the same read as the headers.
886 // Must surface MissingHeaders (so the worker dispatches to httpFallback)
887 // — NOT return null (which would loop forever waiting for more data).
888 var state = try pool.acquire();
889 defer state.release();
890 try t.expectError(
891 error.MissingHeaders,
892 testHandshake(
893 "POST /xrpc/com.atproto.sync.requestCrawl HTTP/1.1\r\nHost: relay\r\nContent-Type: application/json\r\nContent-Length: 27\r\n\r\n{\"hostname\":\"pds.test.com\"}",
894 state,
895 ),
896 );
897 }
898
899 {
900 // POST whose body has NOT fully arrived yet — parse must return null
901 // so the worker reads more data before dispatching the fallback.
902 var state = try pool.acquire();
903 defer state.release();
904 try t.expectEqual(
905 null,
906 try testHandshake(
907 "POST /x HTTP/1.1\r\nHost: r\r\nContent-Length: 100\r\n\r\n{\"only\":\"part\"}",
908 state,
909 ),
910 );
911 }
912
913 {
914 // POST with no Content-Length and headers terminated — surface
915 // MissingHeaders immediately, no body expected.
916 var state = try pool.acquire();
917 defer state.release();
918 try t.expectError(
919 error.MissingHeaders,
920 testHandshake("POST /x HTTP/1.1\r\nHost: r\r\n\r\n", state),
921 );
922 }
923}
924
925test "handshake: Connection: keep-alive + partial body returns null (not InvalidConnection)" {
926 // Regression: an httpx-style POST sends `Connection: keep-alive` (no
927 // Upgrade token). Previously, parse returned `error.InvalidConnection`
928 // the moment it saw that header — even if the body hadn't fully arrived
929 // yet — and the worker dispatched httpFallback with a truncated body.
930 // Handlers observed body.len == 0 for a `Content-Length: 179` POST,
931 // returned 400, and the prefect client crashed the worker run.
932 //
933 // The fix defers `InvalidConnection`/`InvalidUpgrade` past the
934 // body-completeness check so we ask for more data first.
935
936 var pool = try Pool.init(t.allocator, 1, 1024, 10, 1);
937 defer pool.deinit();
938
939 {
940 // Headers only (no body yet) — parse must return null so the worker
941 // reads more. It must NOT return error.InvalidConnection.
942 var state = try pool.acquire();
943 defer state.release();
944 try t.expectEqual(
945 null,
946 try testHandshake(
947 "POST /api/flow_runs/x/labels HTTP/1.1\r\n" ++
948 "Host: localhost\r\n" ++
949 "Connection: keep-alive\r\n" ++
950 "Content-Type: application/json\r\n" ++
951 "Content-Length: 50\r\n" ++
952 "\r\n",
953 state,
954 ),
955 );
956 }
957
958 {
959 // Same request with body now present — parse must surface
960 // InvalidConnection (so the worker dispatches httpFallback with the
961 // full body now in state.buf).
962 var state = try pool.acquire();
963 defer state.release();
964 try t.expectError(
965 error.InvalidConnection,
966 testHandshake(
967 "POST /api/flow_runs/x/labels HTTP/1.1\r\n" ++
968 "Host: localhost\r\n" ++
969 "Connection: keep-alive\r\n" ++
970 "Content-Type: application/json\r\n" ++
971 "Content-Length: 50\r\n" ++
972 "\r\n" ++
973 "{\"prefect.flow-run.smoke\":\"\",\"prefect.worker\":\"w\"}",
974 state,
975 ),
976 );
977 }
978
979 {
980 // Non-websocket Upgrade token (e.g., Upgrade: h2c) on a POST — same
981 // pattern, this time triggering the InvalidUpgrade defer path.
982 var state = try pool.acquire();
983 defer state.release();
984 try t.expectEqual(
985 null,
986 try testHandshake(
987 "POST /x HTTP/1.1\r\n" ++
988 "Host: r\r\n" ++
989 "Upgrade: h2c\r\n" ++
990 "Connection: upgrade\r\n" ++
991 "Content-Length: 20\r\n" ++
992 "\r\n",
993 state,
994 ),
995 );
996 }
997}
998
999test "handshake: reply" {
1000 var buf: [512]u8 = undefined;
1001
1002 {
1003 // no compression
1004 const expected =
1005 "HTTP/1.1 101 Switching Protocols\r\n" ++
1006 "Upgrade: websocket\r\n" ++
1007 "Connection: upgrade\r\n" ++
1008 "Sec-Websocket-Accept: flzHu2DevQ2dSCSVqKSii5e9C2o=\r\n\r\n";
1009 try t.expectString(expected, try Handshake.createReply("this is my key", null, false, &buf));
1010 }
1011
1012 // Compatibility for http.zig and other integrations which own the HTTP
1013 // upgrade and pass the historical boolean negotiation flag.
1014 {
1015 const expected =
1016 "HTTP/1.1 101 Switching Protocols\r\n" ++
1017 "Upgrade: websocket\r\n" ++
1018 "Connection: upgrade\r\n" ++
1019 "Sec-Websocket-Accept: flzHu2DevQ2dSCSVqKSii5e9C2o=\r\n" ++
1020 "Sec-WebSocket-Extensions: permessage-deflate; server_no_context_takeover; client_no_context_takeover\r\n\r\n";
1021 try t.expectString(expected, try Handshake.createReply("this is my key", null, true, &buf));
1022 }
1023
1024 {
1025 // compression
1026 const expected =
1027 "HTTP/1.1 101 Switching Protocols\r\n" ++
1028 "Upgrade: websocket\r\n" ++
1029 "Connection: upgrade\r\n" ++
1030 "Sec-Websocket-Accept: flzHu2DevQ2dSCSVqKSii5e9C2o=\r\n" ++
1031 "Sec-WebSocket-Extensions: permessage-deflate; server_no_context_takeover; client_no_context_takeover\r\n\r\n";
1032 try t.expectString(expected, try Handshake.createReplyNegotiated("this is my key", null, .{
1033 .server_no_context_takeover = true,
1034 .client_no_context_takeover = true,
1035 }, &buf));
1036 }
1037
1038 // With custom headers
1039 var res_headers = try Handshake.KeyValue.init(t.allocator, 2);
1040 defer res_headers.deinit(t.allocator);
1041 res_headers.add("Set-Cookie", "Yummy!");
1042
1043 {
1044 // no compression
1045 const expected =
1046 "HTTP/1.1 101 Switching Protocols\r\n" ++
1047 "Upgrade: websocket\r\n" ++
1048 "Connection: upgrade\r\n" ++
1049 "Sec-Websocket-Accept: flzHu2DevQ2dSCSVqKSii5e9C2o=\r\n" ++
1050 "Set-Cookie: Yummy!\r\n\r\n";
1051 try t.expectString(expected, try Handshake.createReply("this is my key", &res_headers, false, &buf));
1052 }
1053
1054 {
1055 // compression
1056 const expected =
1057 "HTTP/1.1 101 Switching Protocols\r\n" ++
1058 "Upgrade: websocket\r\n" ++
1059 "Connection: upgrade\r\n" ++
1060 "Sec-Websocket-Accept: flzHu2DevQ2dSCSVqKSii5e9C2o=\r\n" ++
1061 "Sec-WebSocket-Extensions: permessage-deflate; server_no_context_takeover; client_no_context_takeover\r\n" ++
1062 "Set-Cookie: Yummy!\r\n\r\n";
1063 try t.expectString(expected, try Handshake.createReplyNegotiated("this is my key", &res_headers, .{
1064 .server_no_context_takeover = true,
1065 .client_no_context_takeover = true,
1066 }, &buf));
1067 }
1068}
1069
1070test "handshake: permessage-deflate offer selection" {
1071 const context = (try Handshake.parseExtension(
1072 "permessage-deflate; client_max_window_bits",
1073 )).?;
1074 try t.expectEqual(false, context.client_no_context_takeover);
1075 try t.expectEqual(false, context.server_no_context_takeover);
1076
1077 const reset = (try Handshake.parseExtension(
1078 "permessage-deflate; server_no_context_takeover; client_no_context_takeover; server_max_window_bits=15",
1079 )).?;
1080 try t.expectEqual(true, reset.client_no_context_takeover);
1081 try t.expectEqual(true, reset.server_no_context_takeover);
1082
1083 try t.expectEqual(null, try Handshake.parseExtension(
1084 "permessage-deflate; server_max_window_bits=12",
1085 ));
1086 try t.expectEqual(null, try Handshake.parseExtension(
1087 "permessage-deflate; made_up_parameter=yes",
1088 ));
1089 try std.testing.expect((try Handshake.parseExtension(
1090 "x-webkit-deflate-frame, permessage-deflate; client_max_window_bits=12",
1091 )) != null);
1092}
1093
1094test "KeyValue: get" {
1095 const allocator = t.allocator;
1096 var kv = try Handshake.KeyValue.init(allocator, 2);
1097 defer kv.deinit(t.allocator);
1098
1099 var key = "content-type".*;
1100 kv.add(&key, "application/json");
1101
1102 try t.expectString("application/json", kv.get("content-type").?);
1103
1104 kv.len = 0;
1105 try t.expectEqual(null, kv.get("content-type"));
1106 kv.add(&key, "application/json2");
1107 try t.expectString("application/json2", kv.get("content-type").?);
1108}
1109
1110test "KeyValue: ignores beyond max" {
1111 var kv = try Handshake.KeyValue.init(t.allocator, 2);
1112 defer kv.deinit(t.allocator);
1113
1114 var n1 = "content-length".*;
1115 kv.add(&n1, "cl");
1116
1117 var n2 = "host".*;
1118 kv.add(&n2, "www");
1119
1120 var n3 = "authorization".*;
1121 kv.add(&n3, "hack");
1122
1123 try t.expectString("cl", kv.get("content-length").?);
1124 try t.expectString("www", kv.get("host").?);
1125 try t.expectEqual(null, kv.get("authorization"));
1126}
1127
1128test "pool: acquire and release" {
1129 // not 100% sure this is testing exactly what I want, but it's ....something ?
1130 var p = try Pool.init(t.allocator, 2, 10, 3, 1);
1131 defer p.deinit();
1132
1133 var hs1a = p.acquire() catch unreachable;
1134 var hs2a = p.acquire() catch unreachable;
1135 var hs3a = p.acquire() catch unreachable; // this should be dynamically generated
1136
1137 try t.expectEqual(false, &hs1a.buf[0] == &hs2a.buf[0]);
1138 try t.expectEqual(false, &hs2a.buf[0] == &hs3a.buf[0]);
1139 try t.expectEqual(10, hs1a.buf.len);
1140 try t.expectEqual(10, hs2a.buf.len);
1141 try t.expectEqual(10, hs3a.buf.len);
1142 try t.expectEqual(0, hs1a.req_headers.len);
1143 try t.expectEqual(0, hs2a.req_headers.len);
1144 try t.expectEqual(0, hs3a.req_headers.len);
1145 try t.expectEqual(3, hs1a.req_headers.keys.len);
1146 try t.expectEqual(3, hs2a.req_headers.keys.len);
1147 try t.expectEqual(3, hs3a.req_headers.keys.len);
1148
1149 p.release(hs1a);
1150
1151 var hs1b = p.acquire() catch unreachable;
1152 try t.expectEqual(true, &hs1a.buf[0] == &hs1b.buf[0]);
1153
1154 p.release(hs3a);
1155 p.release(hs2a);
1156 p.release(hs1b);
1157}
1158
1159test "Handshake.Pool: threadsafety" {
1160 var p = try Pool.init(t.allocator, 4, 10, 2, 2);
1161 defer p.deinit();
1162
1163 for (p.states) |hs| {
1164 hs.buf[0] = 0;
1165 }
1166
1167 const t1 = try std.Thread.spawn(.{}, testPool, .{p});
1168 const t2 = try std.Thread.spawn(.{}, testPool, .{p});
1169 const t3 = try std.Thread.spawn(.{}, testPool, .{p});
1170 const t4 = try std.Thread.spawn(.{}, testPool, .{p});
1171
1172 t1.join();
1173 t2.join();
1174 t3.join();
1175 t4.join();
1176}
1177
1178fn testPool(p: *Pool) void {
1179 var r = t.getRandom();
1180 const random = r.random();
1181
1182 for (0..5000) |_| {
1183 var hs = p.acquire() catch unreachable;
1184 std.debug.assert(hs.buf[0] == 0);
1185 hs.buf[0] = 255;
1186 sleep(random.uintAtMost(u32, 100000));
1187 hs.buf[0] = 0;
1188 p.release(hs);
1189 }
1190}
1191
1192fn testHandshake(request: []const u8, state: *Handshake.State) !?Handshake {
1193 // simulate a fresh socket: clear streaming parse state before pushing
1194 // a new request. existing tests reuse one State across many distinct
1195 // requests for convenience, which would otherwise wedge head_parser
1196 // in whatever state the previous request left it in.
1197 state.head_parser = .{};
1198 state.head_fed = 0;
1199 state.req_headers.len = 0;
1200 @memcpy(state.buf[0..request.len], request);
1201 state.len = request.len;
1202 return Handshake.parse(state);
1203}
1204
1205/// Feed the request `chunk_size` bytes at a time, asserting that all
1206/// non-final calls return null. Returns the result of the final parse.
1207/// This is the core regression-test mechanism for TCP-split bugs:
1208/// any parser that's correct under all-at-once must also be correct
1209/// under every possible byte-split arrival pattern.
1210fn testHandshakeIncremental(request: []const u8, chunk_size: usize, state: *Handshake.State) !?Handshake {
1211 state.head_parser = .{};
1212 state.head_fed = 0;
1213 state.req_headers.len = 0;
1214 state.len = 0;
1215
1216 var fed: usize = 0;
1217 while (fed < request.len) {
1218 const take = @min(chunk_size, request.len - fed);
1219 @memcpy(state.buf[state.len .. state.len + take], request[fed .. fed + take]);
1220 state.len += take;
1221 fed += take;
1222 if (fed < request.len) {
1223 // not the last chunk yet — parser must not commit either way
1224 const r = Handshake.parse(state) catch |e| {
1225 // an early error is allowed only if the request is
1226 // detectably malformed regardless of remaining bytes
1227 // (e.g. wrong HTTP version in the request line).
1228 return e;
1229 };
1230 if (r != null) {
1231 // committing to success before the request is fully delivered
1232 // would be a bug — at minimum it would race with body bytes
1233 // for non-upgrade requests.
1234 return error.PrematureSuccess;
1235 }
1236 }
1237 }
1238 return Handshake.parse(state);
1239}