calabro.io
forked from
zzstoatzz.io/websocket.zig
websocket
7.6 kB
248 lines
1const std = @import("std");
2const proto = @import("proto.zig");
3
4const Io = std.Io;
5const net = Io.net;
6const ArrayList = std.ArrayList;
7
8const Message = proto.Message;
9
10pub const allocator = std.testing.allocator;
11
12pub fn expectEqual(expected: anytype, actual: anytype) !void {
13 try std.testing.expectEqual(expected, actual);
14}
15
16pub const expectError = std.testing.expectError;
17pub const expectString = std.testing.expectEqualStrings;
18pub const expectSlice = std.testing.expectEqualSlices;
19
20pub fn getRandom() std.Random.DefaultPrng {
21 const io = std.Options.debug_io;
22 var seed_bytes: [8]u8 = undefined;
23 io.random(&seed_bytes);
24 const seed: u64 = @bitCast(seed_bytes);
25 return std.Random.DefaultPrng.init(seed);
26}
27
28pub var arena = std.heap.ArenaAllocator.init(allocator);
29pub fn reset() void {
30 _ = arena.reset(.free_all);
31}
32
33pub const Writer = struct {
34 pos: usize,
35 buf: std.ArrayList(u8),
36 random: std.Random.DefaultPrng,
37
38 pub fn init() Writer {
39 return .{
40 .pos = 0,
41 .buf = .empty,
42 .random = getRandom(),
43 };
44 }
45
46 pub fn deinit(self: *Writer) void {
47 self.buf.deinit(allocator);
48 }
49
50 pub fn ping(self: *Writer) void {
51 return self.pingPayload("");
52 }
53
54 pub fn pong(self: *Writer) void {
55 return self.frame(true, 10, "", 0);
56 }
57
58 pub fn pingPayload(self: *Writer, payload: []const u8) void {
59 return self.frame(true, 9, payload, 0);
60 }
61
62 pub fn textFrame(self: *Writer, fin: bool, payload: []const u8) void {
63 return self.frame(fin, 1, payload, 0);
64 }
65
66 pub fn cont(self: *Writer, fin: bool, payload: []const u8) void {
67 return self.frame(fin, 0, payload, 0);
68 }
69
70 pub fn frame(self: *Writer, fin: bool, op_code: u8, payload: []const u8, reserved: u8) void {
71 var buf = &self.buf;
72
73 const l = payload.len;
74 var length_of_length: usize = 0;
75
76 if (l > 125) {
77 if (l < 65536) {
78 length_of_length = 2;
79 } else {
80 length_of_length = 8;
81 }
82 }
83
84 // 2 byte header + length_of_length + mask + payload_length
85 const needed = 2 + length_of_length + 4 + l;
86 buf.ensureUnusedCapacity(allocator, needed) catch unreachable;
87
88 if (fin) {
89 buf.appendAssumeCapacity(128 | op_code | reserved);
90 } else {
91 buf.appendAssumeCapacity(op_code | reserved);
92 }
93
94 if (length_of_length == 0) {
95 buf.appendAssumeCapacity(128 | @as(u8, @intCast(l)));
96 } else if (length_of_length == 2) {
97 buf.appendAssumeCapacity(128 | 126);
98 buf.appendAssumeCapacity(@intCast((l >> 8) & 0xFF));
99 buf.appendAssumeCapacity(@intCast(l & 0xFF));
100 } else {
101 buf.appendAssumeCapacity(128 | 127);
102 buf.appendAssumeCapacity(@intCast((l >> 56) & 0xFF));
103 buf.appendAssumeCapacity(@intCast((l >> 48) & 0xFF));
104 buf.appendAssumeCapacity(@intCast((l >> 40) & 0xFF));
105 buf.appendAssumeCapacity(@intCast((l >> 32) & 0xFF));
106 buf.appendAssumeCapacity(@intCast((l >> 24) & 0xFF));
107 buf.appendAssumeCapacity(@intCast((l >> 16) & 0xFF));
108 buf.appendAssumeCapacity(@intCast((l >> 8) & 0xFF));
109 buf.appendAssumeCapacity(@intCast(l & 0xFF));
110 }
111
112 var mask: [4]u8 = undefined;
113 self.random.random().bytes(&mask);
114
115 buf.appendSliceAssumeCapacity(&mask);
116 for (payload, 0..) |b, i| {
117 buf.appendAssumeCapacity(b ^ mask[i & 3]);
118 }
119 }
120
121 pub fn bytes(self: *const Writer) []const u8 {
122 return self.buf.items;
123 }
124
125 pub fn clear(self: *Writer) void {
126 self.pos = 0;
127 self.buf.clearRetainingCapacity();
128 }
129
130 pub fn read(
131 self: *Writer,
132 buf: []u8,
133 ) !usize {
134 const data = self.buf.items[self.pos..];
135
136 if (data.len == 0 or buf.len == 0) {
137 return 0;
138 }
139
140 // randomly fragment the data
141 const to_read = self.random.random().intRangeAtMost(usize, 1, @min(data.len, buf.len));
142 @memcpy(buf[0..to_read], data[0..to_read]);
143 self.pos += to_read;
144 return to_read;
145 }
146};
147
148/// Reads directly from a socket, preserving SO_RCVTIMEO as WouldBlock.
149pub const StreamReader = struct {
150 handle: net.Socket.Handle,
151
152 pub fn read(self: *StreamReader, buf: []u8) !usize {
153 return std.posix.read(self.handle, buf) catch |err| {
154 return switch (err) {
155 error.ConnectionResetByPeer => error.ConnectionResetByPeer,
156 error.WouldBlock => error.WouldBlock,
157 else => error.Unexpected,
158 };
159 };
160 }
161};
162
163pub const SocketPair = struct {
164 writer: Writer,
165 io: Io,
166 client: net.Stream,
167 server: net.Stream,
168 server_taken: bool = false,
169
170 const Opts = struct {
171 port: ?u16 = null,
172 };
173
174 pub fn init(opts: Opts) SocketPair {
175 const io = std.Options.debug_io;
176 const port: u16 = opts.port orelse 0;
177
178 // use Io.net to create a listener, connect a client, and accept
179 var listen_addr: net.IpAddress = .{ .ip4 = .loopback(port) };
180 var server = net.IpAddress.listen(&listen_addr, io, .{ .reuse_address = true }) catch unreachable;
181
182 // get the actual bound address (for ephemeral port)
183 const bound_addr = server.socket.address;
184
185 // connect client
186 const client = net.IpAddress.connect(&bound_addr, io, .{ .mode = .stream }) catch unreachable;
187
188 // accept server-side connection
189 const accepted = server.accept(io) catch unreachable;
190 server.deinit(io);
191
192 return .{
193 .io = io,
194 .client = client,
195 .server = accepted,
196 .writer = Writer.init(),
197 };
198 }
199
200 pub fn deinit(self: *SocketPair) void {
201 self.writer.deinit();
202 self.client.close(self.io);
203 if (!self.server_taken) self.server.close(self.io);
204 }
205
206 pub fn pingPayload(self: *SocketPair, payload: []const u8) void {
207 self.writer.pingPayload(payload);
208 }
209
210 pub fn textFrame(self: *SocketPair, fin: bool, payload: []const u8) void {
211 self.writer.textFrame(fin, payload);
212 }
213
214 pub fn cont(self: *SocketPair, fin: bool, payload: []const u8) void {
215 self.writer.cont(fin, payload);
216 }
217
218 pub fn sendBuf(self: *SocketPair) void {
219 self.ioWriteAll(self.writer.bytes()) catch unreachable;
220 self.writer.clear();
221 }
222
223 pub fn clientWriteAll(self: *SocketPair, data: []const u8) !void {
224 try self.ioWriteAll(data);
225 }
226
227 fn ioWriteAll(self: *SocketPair, data: []const u8) !void {
228 var remaining = data;
229 while (remaining.len > 0) {
230 // netWrite: header is sent first, data array's last element is the splat pattern.
231 // Pass remaining as header, empty pattern with splat=0.
232 const empty = [_][]const u8{""};
233 const n = self.io.vtable.netWrite(self.io.userdata, self.client.socket.handle, remaining, &empty, 0) catch {
234 return error.SendError;
235 };
236 if (n == 0) return error.SendError;
237 remaining = remaining[n..];
238 }
239 }
240
241 pub fn serverReader(self: *SocketPair) StreamReader {
242 return .{ .handle = self.server.socket.handle };
243 }
244
245 pub fn clientReader(self: *SocketPair) StreamReader {
246 return .{ .handle = self.client.socket.handle };
247 }
248};