lib/quic/src/connection/assemble.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const quic = @import("../root.zig");
3
4 const KeyStorage = [quic.crypto.Keys.storage_bytes_max]u8;
5
6 pub const PacketType = enum { initial, handshake, one_rtt };
7
8 pub const Spec = struct {
9 packet_type: PacketType,
10 destination: quic.packet.ConnectionId,
11 source: quic.packet.ConnectionId,
12 number: quic.connection.NumberEncoding,
13 key_phase: bool = false,
14 minimum_bytes: u16 = 0,
15 };
16
17 pub const Built = struct {
18 length: u16,
19 packet_number_offset: usize,
20 packet_number_length: u3,
21 payload_length: u16,
22 };
23
24 pub const Error = quic.packet.EncodeError || quic.crypto.packet.SealError ||
25 error{ NoSpace, PayloadTooLarge };
26
27 /// Four bytes is the floor that a packet number and the payload behind it have to cover between
28 /// them, so that header protection finds a whole sample to read. `packet` grows a short payload
29 /// until the two together reach that floor, so a caller never has to pad by hand. One PING byte
30 /// carried under a one-byte packet number therefore reaches the wire with three payload bytes.
31 pub const sampled_bytes_min: usize = 4;
32
33 pub fn packet(
34 keys: *quic.crypto.Keys,
35 spec: Spec,
36 payload: []const u8,
37 out: []u8,
38 ) Error!Built {
39 const pn_len: u3 = @intCast(spec.number.bits / 8);
40 std.debug.assert(pn_len >= 1);
41 std.debug.assert(pn_len <= sampled_bytes_min);
42 const payload_floor = @max(payload.len, sampled_bytes_min - pn_len);
43 var payload_len = payload_floor;
44 for (0..4) |_| {
45 const total = try packetLength(spec, pn_len, payload_len);
46 if (total == spec.minimum_bytes or spec.minimum_bytes == 0) break;
47 if (total < spec.minimum_bytes) {
48 payload_len = std.math.add(usize, payload_len, spec.minimum_bytes - total) catch
49 return error.PayloadTooLarge;
50 } else {
51 const excess = total - spec.minimum_bytes;
52 if (excess > payload_len - payload_floor) break;
53 payload_len -= excess;
54 }
55 }
56 const total = try packetLength(spec, pn_len, payload_len);
57 if (total > out.len) return error.NoSpace;
58 if (payload_len > std.math.maxInt(u16)) return error.PayloadTooLarge;
59 var output = quic.cursor.Write.init(out[0..total]);
60 const pn_offset = try writeHeader(spec, pn_len, payload_len, &output);
61 try writePacketNumber(spec.number.packet_number, pn_len, &output);
62 try output.put(payload);
63 for (payload.len..payload_len) |_| try output.byte(0);
64 const tag_end = std.math.add(usize, output.index, quic.crypto.packet.tag_bytes) catch
65 return error.PayloadTooLarge;
66 if (tag_end != total) return error.PayloadTooLarge;
67 @memset(out[output.index..tag_end], 0);
68 try quic.crypto.packet.seal(
69 keys,
70 spec.number.packet_number,
71 out[0..total],
72 pn_offset,
73 pn_len,
74 @intCast(payload_len),
75 );
76 return .{
77 .length = @intCast(total),
78 .packet_number_offset = pn_offset,
79 .packet_number_length = pn_len,
80 .payload_length = @intCast(payload_len),
81 };
82 }
83
84 fn packetLength(spec: Spec, pn_len: u3, payload_len: usize) Error!usize {
85 const protected_len = std.math.add(
86 usize,
87 payload_len,
88 quic.crypto.packet.tag_bytes,
89 ) catch return error.PayloadTooLarge;
90 const body_len = std.math.add(usize, pn_len, protected_len) catch
91 return error.PayloadTooLarge;
92 const header_len = switch (spec.packet_type) {
93 .initial, .handshake => try longHeaderLength(spec, body_len),
94 .one_rtt => 1 + spec.destination.length,
95 };
96 return std.math.add(usize, header_len, body_len) catch error.PayloadTooLarge;
97 }
98
99 fn longHeaderLength(spec: Spec, body_len: usize) Error!usize {
100 if (body_len > std.math.maxInt(u62)) return error.PayloadTooLarge;
101 const destination: usize = spec.destination.length;
102 const source: usize = spec.source.length;
103 const common = 1 + 4 + 1 + destination + 1 + source;
104 const token: usize = if (spec.packet_type == .initial) 1 else 0;
105 return common + token + quic.varint.encodedLength(@intCast(body_len));
106 }
107
108 fn writeHeader(
109 spec: Spec,
110 pn_len: u3,
111 payload_len: usize,
112 output: *quic.cursor.Write,
113 ) Error!usize {
114 if (spec.packet_type == .one_rtt) {
115 const first: u8 = 0x40 | (@as(u8, @intFromBool(spec.key_phase)) << 2) | (pn_len - 1);
116 try quic.packet.encodeShort(.{
117 .first = first,
118 .destination = spec.destination,
119 .packet_number_offset = 0,
120 }, output);
121 return output.index;
122 }
123 const protected_len = payload_len + quic.crypto.packet.tag_bytes;
124 const length: u62 = @intCast(@as(usize, pn_len) + protected_len);
125 const first: u8 = switch (spec.packet_type) {
126 .initial => 0xc0 | @as(u8, pn_len - 1),
127 .handshake => 0xe0 | @as(u8, pn_len - 1),
128 .one_rtt => unreachable,
129 };
130 const common = quic.packet.Common{
131 .first = first,
132 .version = 1,
133 .destination = spec.destination,
134 .source = spec.source,
135 };
136 const value: quic.packet.Long = switch (spec.packet_type) {
137 .initial => .{ .initial = .{
138 .common = common,
139 .token = &.{},
140 .length = length,
141 .packet_number_offset = 0,
142 } },
143 .handshake => .{ .handshake = .{
144 .common = common,
145 .length = length,
146 .packet_number_offset = 0,
147 } },
148 .one_rtt => unreachable,
149 };
150 try quic.packet.encodeLong(value, output);
151 return output.index;
152 }
153
154 fn writePacketNumber(value: u62, length: u3, output: *quic.cursor.Write) !void {
155 var bytes: [8]u8 = undefined;
156 std.mem.writeInt(u64, &bytes, value, .big);
157 try output.put(bytes[bytes.len - length ..]);
158 }
159
160 test "RFC 9000 section 12.2 assembled Initial layout decodes through packet codec" {
161 const secret: quic.crypto.Secret = @splat(0x42);
162 var key_bytes: KeyStorage align(quic.crypto.Keys.storage_alignment) = undefined;
163 var keys = try quic.crypto.Keys.derive(&key_bytes, .aes_128_gcm_sha256, secret);
164 defer _ = keys.deinit();
165 const cid = try quic.packet.ConnectionId.init("12345678");
166 var out: [1200]u8 = undefined;
167 const built = try packet(&keys, .{
168 .packet_type = .initial,
169 .destination = cid,
170 .source = cid,
171 .number = .{ .packet_number = 0, .bits = 32, .value = 0 },
172 .minimum_bytes = 1200,
173 }, &.{0x01}, &out);
174 try std.testing.expectEqual(@as(u16, 1200), built.length);
175 const decoded = try quic.packet.decodeLong(out[0..built.length]);
176 try std.testing.expectEqual(built.packet_number_offset, decoded.initial.packet_number_offset);
177 try std.testing.expectEqual(
178 @as(u62, @intCast(1200 - built.packet_number_offset)),
179 decoded.initial.length,
180 );
181 }
182
183 test "RFC 9001 section 5.4.2 one-byte PING pads to a full header protection sample" {
184 const secret: quic.crypto.Secret = @splat(0x42);
185 var key_bytes: KeyStorage align(quic.crypto.Keys.storage_alignment) = undefined;
186 var keys = try quic.crypto.Keys.derive(&key_bytes, .aes_128_gcm_sha256, secret);
187 defer _ = keys.deinit();
188 const cid = try quic.packet.ConnectionId.init("12345678");
189 var out: [64]u8 = undefined;
190 const built = try packet(&keys, .{
191 .packet_type = .one_rtt,
192 .destination = cid,
193 .source = cid,
194 .number = .{ .packet_number = 7, .bits = 8, .value = 7 },
195 }, &.{0x01}, &out);
196 try std.testing.expectEqual(@as(u16, sampled_bytes_min - 1), built.payload_length);
197 try std.testing.expectEqual(@as(u16, 1 + 8 + sampled_bytes_min + 16), built.length);
198 }