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 }