lib/http/src/chunk.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 
  3 pub const Error = error{
  4     IncompleteBody,
  5     MalformedBody,
  6     BodyTooLarge,
  7     OutOfMemory,
  8 };
  9 
 10 pub const Scan = union(enum) {
 11     incomplete,
 12     malformed,
 13     complete: usize,
 14 };
 15 
 16 pub const Decode = struct {
 17     body: []u8,
 18     consumed: usize,
 19 };
 20 
 21 pub const Capacity = struct {
 22     decoded_bytes: usize,
 23     consumed: usize,
 24 };
 25 
 26 pub const Decoder = struct {
 27     state: State = .size,
 28     size: usize = 0,
 29     size_started: bool = false,
 30     size_finished: bool = false,
 31     extension: bool = false,
 32     remaining: usize = 0,
 33     trailer_empty: bool = true,
 34     done: bool = false,
 35 
 36     const State = enum {
 37         size,
 38         size_lf,
 39         data,
 40         data_cr,
 41         data_lf,
 42         trailer,
 43         trailer_lf,
 44         done,
 45     };
 46 
 47     pub fn init() Decoder {
 48         return .{};
 49     }
 50 
 51     pub fn feed(self: *Decoder, data: []const u8, handler: anytype) anyerror!usize {
 52         if (self.done) return 0;
 53         var index: usize = 0;
 54         while (index < data.len and !self.done) {
 55             switch (self.state) {
 56                 .size => {
 57                     const byte = data[index];
 58                     index += 1;
 59                     if (byte == '\r') {
 60                         if (!self.size_started) return error.MalformedBody;
 61                         self.state = .size_lf;
 62                         continue;
 63                     }
 64                     if (byte == '\n') return error.MalformedBody;
 65                     if (self.extension) continue;
 66                     if (byte == ';') {
 67                         if (!self.size_started) return error.MalformedBody;
 68                         self.extension = true;
 69                         continue;
 70                     }
 71                     if (byte == ' ' or byte == '\t') {
 72                         if (self.size_started) self.size_finished = true;
 73                         continue;
 74                     }
 75                     const digit = hexDigit(byte) orelse return error.MalformedBody;
 76                     if (self.size_finished) return error.MalformedBody;
 77                     self.size = std.math.mul(usize, self.size, 16) catch {
 78                         return error.MalformedBody;
 79                     };
 80                     self.size = std.math.add(usize, self.size, digit) catch {
 81                         return error.MalformedBody;
 82                     };
 83                     self.size_started = true;
 84                 },
 85                 .size_lf => {
 86                     if (data[index] != '\n') return error.MalformedBody;
 87                     index += 1;
 88                     if (self.size == 0) {
 89                         self.trailer_empty = true;
 90                         self.state = .trailer;
 91                     } else {
 92                         self.remaining = self.size;
 93                         self.state = .data;
 94                     }
 95                 },
 96                 .data => {
 97                     const take = @min(self.remaining, data.len - index);
 98                     if (take != 0) try handler.emit(data[index..][0..take]);
 99                     index += take;
100                     self.remaining -= take;
101                     if (self.remaining == 0) self.state = .data_cr;
102                 },
103                 .data_cr => {
104                     if (data[index] != '\r') return error.MalformedBody;
105                     index += 1;
106                     self.state = .data_lf;
107                 },
108                 .data_lf => {
109                     if (data[index] != '\n') return error.MalformedBody;
110                     index += 1;
111                     self.resetSize();
112                 },
113                 .trailer => {
114                     const byte = data[index];
115                     index += 1;
116                     if (byte == '\r') {
117                         self.state = .trailer_lf;
118                     } else if (byte == '\n') {
119                         return error.MalformedBody;
120                     } else {
121                         self.trailer_empty = false;
122                     }
123                 },
124                 .trailer_lf => {
125                     if (data[index] != '\n') return error.MalformedBody;
126                     index += 1;
127                     if (self.trailer_empty) {
128                         self.done = true;
129                         self.state = .done;
130                     } else {
131                         self.trailer_empty = true;
132                         self.state = .trailer;
133                     }
134                 },
135                 .done => unreachable,
136             }
137         }
138         return index;
139     }
140 
141     fn resetSize(self: *Decoder) void {
142         self.state = .size;
143         self.size = 0;
144         self.size_started = false;
145         self.size_finished = false;
146         self.extension = false;
147     }
148 };
149 
150 fn hexDigit(byte: u8) ?usize {
151     return switch (byte) {
152         '0'...'9' => byte - '0',
153         'a'...'f' => byte - 'a' + 10,
154         'A'...'F' => byte - 'A' + 10,
155         else => null,
156     };
157 }
158 
159 pub fn parseSize(line: []const u8) ?usize {
160     if (line.len == 0) return null;
161     const clean = if (std.mem.indexOf(u8, line, ";")) |sc|
162         std.mem.trim(u8, line[0..sc], " \t")
163     else
164         line;
165     if (clean.len == 0) return null;
166     return std.fmt.parseInt(usize, clean, 16) catch null;
167 }
168 
169 pub fn scanBody(data: []const u8) Scan {
170     const capacity = survey(data, null) catch |err| switch (err) {
171         error.IncompleteBody => return .incomplete,
172         error.MalformedBody, error.BodyTooLarge => return .malformed,
173         error.OutOfMemory => unreachable,
174     };
175     return .{ .complete = capacity.consumed };
176 }
177 
178 pub fn survey(data: []const u8, max_decoded_size: ?usize) Error!Capacity {
179     var pos: usize = 0;
180     var decoded_bytes: usize = 0;
181     while (true) {
182         const size_end = std.mem.indexOf(u8, data[pos..], "\r\n") orelse {
183             return error.IncompleteBody;
184         };
185         const size_line = std.mem.trim(u8, data[pos .. pos + size_end], " \t");
186         const chunk_size = parseSize(size_line) orelse return error.MalformedBody;
187         pos += size_end + 2;
188 
189         if (chunk_size == 0) {
190             while (true) {
191                 const trailer_end = std.mem.indexOf(u8, data[pos..], "\r\n") orelse {
192                     return error.IncompleteBody;
193                 };
194                 pos += trailer_end + 2;
195                 if (trailer_end == 0) {
196                     return .{ .decoded_bytes = decoded_bytes, .consumed = pos };
197                 }
198             }
199         }
200 
201         decoded_bytes = std.math.add(usize, decoded_bytes, chunk_size) catch {
202             return error.BodyTooLarge;
203         };
204         if (max_decoded_size) |limit| {
205             if (decoded_bytes > limit) return error.BodyTooLarge;
206         }
207         if (chunk_size > data.len - pos) return error.IncompleteBody;
208         const chunk_end = pos + chunk_size;
209         if (data.len - chunk_end < 2) return error.IncompleteBody;
210         if (data[chunk_end] != '\r' or data[chunk_end + 1] != '\n') {
211             return error.MalformedBody;
212         }
213         pos = chunk_end + 2;
214     }
215 }
216 
217 pub fn decodeInto(
218     output: []u8,
219     data: []const u8,
220     max_decoded_size: ?usize,
221 ) Error!Decode {
222     const capacity = try survey(data, max_decoded_size);
223     return decodeSurveyedInto(output, data, capacity);
224 }
225 
226 pub fn decodeSurveyedInto(output: []u8, data: []const u8, capacity: Capacity) Error!Decode {
227     if (capacity.decoded_bytes > output.len) return error.BodyTooLarge;
228 
229     var pos: usize = 0;
230     var output_length: usize = 0;
231 
232     while (true) {
233         const size_end = std.mem.indexOf(u8, data[pos..], "\r\n") orelse {
234             return error.IncompleteBody;
235         };
236         const chunk_size = parseSize(std.mem.trim(
237             u8,
238             data[pos .. pos + size_end],
239             " \t",
240         )) orelse return error.MalformedBody;
241         pos += size_end + 2;
242         if (chunk_size == 0) {
243             if (output_length != capacity.decoded_bytes) return error.MalformedBody;
244             return .{
245                 .body = output[0..output_length],
246                 .consumed = capacity.consumed,
247             };
248         }
249         if (chunk_size > capacity.decoded_bytes - output_length) {
250             return error.MalformedBody;
251         }
252         if (chunk_size > data.len - pos) return error.IncompleteBody;
253         const chunk_end = pos + chunk_size;
254         std.mem.copyForwards(
255             u8,
256             output[output_length..][0..chunk_size],
257             data[pos..chunk_end],
258         );
259         output_length += chunk_size;
260         pos = chunk_end + 2;
261     }
262 }
263 
264 pub fn decodeAlloc(
265     allocator: std.mem.Allocator,
266     data: []const u8,
267     max_decoded_size: ?usize,
268 ) Error!Decode {
269     const capacity = try survey(data, max_decoded_size);
270     const body = allocator.alloc(u8, capacity.decoded_bytes) catch {
271         return error.OutOfMemory;
272     };
273     errdefer allocator.free(body);
274     return decodeSurveyedInto(body, data, capacity);
275 }
276 
277 const testing = std.testing;
278 
279 const TestCollector = struct {
280     allocator: std.mem.Allocator,
281     buf: std.ArrayListUnmanaged(u8) = .empty,
282 
283     fn deinit(self: *TestCollector) void {
284         self.buf.deinit(self.allocator);
285     }
286 
287     fn emit(self: *TestCollector, bytes: []const u8) !void {
288         try self.buf.appendSlice(self.allocator, bytes);
289     }
290 };
291 
292 const RejectingChunkHandler = struct {
293     fn emit(_: *@This(), _: []const u8) !void {
294         return error.ChunkRejected;
295     }
296 };
297 
298 const FixedChunkCollector = struct {
299     body: [32]u8 = undefined,
300     length: usize = 0,
301 
302     fn emit(self: *FixedChunkCollector, bytes: []const u8) !void {
303         if (bytes.len > self.body.len - self.length) return error.TestBodyCapacityExceeded;
304         @memcpy(self.body[self.length..][0..bytes.len], bytes);
305         self.length += bytes.len;
306     }
307 };
308 
309 test "chunk scan distinguishes embedded terminal bytes from structure" {
310     try testing.expectEqual(Scan.incomplete, scanBody("9\r\nabc0\r\n\r\n"));
311     try testing.expectEqual(@as(usize, 19), scanBody("9\r\nabc0\r\n\r\nx\r\n0\r\n\r\n").complete);
312 }
313 
314 test "chunk scan waits for trailer terminator" {
315     try testing.expectEqual(Scan.incomplete, scanBody("0\r\nTrailer: value\r\n"));
316     const data = "0\r\nTrailer: value\r\n\r\n";
317     try testing.expectEqual(data.len, scanBody(data).complete);
318 }
319 
320 test "chunk decode single chunk" {
321     const decoded = try decodeAlloc(testing.allocator, "5\r\nHello\r\n0\r\n\r\n", null);
322     defer testing.allocator.free(decoded.body);
323 
324     try testing.expectEqualStrings("Hello", decoded.body);
325 }
326 
327 test "chunk decode multiple chunks" {
328     const decoded = try decodeAlloc(testing.allocator, "5\r\nHello\r\n1\r\n \r\n5\r\nWorld\r\n0\r\n\r\n", null);
329     defer testing.allocator.free(decoded.body);
330 
331     try testing.expectEqualStrings("Hello World", decoded.body);
332 }
333 
334 test "chunk decode accepts trailers" {
335     const data = "5\r\nHello\r\n0\r\nDigest: sha-256=abc123\r\n\r\n";
336     const decoded = try decodeAlloc(testing.allocator, data, null);
337     defer testing.allocator.free(decoded.body);
338 
339     try testing.expectEqualStrings("Hello", decoded.body);
340     try testing.expectEqual(data.len, decoded.consumed);
341 }
342 
343 test "chunk decode empty returns empty owned slice" {
344     const decoded = try decodeAlloc(testing.allocator, "0\r\n\r\n", null);
345     defer testing.allocator.free(decoded.body);
346 
347     try testing.expectEqual(@as(usize, 0), decoded.body.len);
348 }
349 
350 test "chunk decode returns consumed length before extra bytes" {
351     const decoded = try decodeAlloc(testing.allocator, "5\r\nHello\r\n0\r\n\r\nGET /next HTTP/1.1\r\n\r\n", null);
352     defer testing.allocator.free(decoded.body);
353 
354     try testing.expectEqualStrings("Hello", decoded.body);
355     try testing.expectEqual(@as(usize, 15), decoded.consumed);
356 }
357 
358 test "chunk decode rejects incomplete final terminator" {
359     try testing.expectError(error.IncompleteBody, decodeAlloc(testing.allocator, "0\r\n", null));
360 }
361 
362 test "chunk decode preserves allocation failure" {
363     var failing = testing.FailingAllocator.init(testing.allocator, .{ .fail_index = 0 });
364     try testing.expectError(error.OutOfMemory, decodeAlloc(failing.allocator(), "5\r\nHello\r\n0\r\n\r\n", null));
365 }
366 
367 test "chunk decode enforces decoded size limit" {
368     try testing.expectError(error.BodyTooLarge, decodeAlloc(testing.allocator, "5\r\nHello\r\n0\r\n\r\n", 4));
369 }
370 
371 test "chunk survey reports decoded and consumed capacity" {
372     const data = "3\r\none\r\n2\r\n!!\r\n0\r\nTrailer: value\r\n\r\nnext";
373     try testing.expectEqual(Capacity{
374         .decoded_bytes = 5,
375         .consumed = data.len - "next".len,
376     }, try survey(data, 5));
377 }
378 
379 test "chunk decodeInto rejects short output before mutation" {
380     var output = @as([4]u8, @splat(0xA5));
381     try testing.expectError(
382         error.BodyTooLarge,
383         decodeInto(&output, "5\r\nHello\r\n0\r\n\r\n", null),
384     );
385     try testing.expectEqualSlices(u8, &.{ 0xA5, 0xA5, 0xA5, 0xA5 }, &output);
386 }
387 
388 test "chunk decodeInto accepts exact output" {
389     var output: [5]u8 = undefined;
390     const decoded = try decodeInto(&output, "2\r\nHe\r\n3\r\nllo\r\n0\r\n\r\n", 5);
391     try testing.expectEqualStrings("Hello", decoded.body);
392 }
393 
394 test "chunk decodeInto supports in-place compaction" {
395     var storage = [_]u8{ '2', '\r', '\n', 'H', 'e', '\r', '\n', '3', '\r', '\n', 'l', 'l', 'o', '\r', '\n', '0', '\r', '\n', '\r', '\n' };
396     const decoded = try decodeInto(&storage, &storage, storage.len);
397     try testing.expectEqualStrings("Hello", decoded.body);
398 }
399 
400 test "chunk streaming decoder emits data across boundaries" {
401     var collector = TestCollector{ .allocator = testing.allocator };
402     defer collector.deinit();
403 
404     var decoder = Decoder.init();
405 
406     _ = try decoder.feed("4\r\nWiki\r\n", &collector);
407     _ = try decoder.feed("5\r\npedia\r\n0\r\n\r\n", &collector);
408 
409     try testing.expectEqualStrings("Wikipedia", collector.buf.items);
410 }
411 
412 test "chunk streaming decoder reports every exact terminal boundary" {
413     const cases = [_]struct {
414         wire: []const u8,
415         encoded_length: usize,
416     }{
417         .{
418             .wire = "3\r\nabc\r\n0\r\n\r\nTAIL",
419             .encoded_length = "3\r\nabc\r\n0\r\n\r\n".len,
420         },
421         .{
422             .wire = "3\r\nabc\r\n0\r\nTrace: value\r\n\r\nTAIL",
423             .encoded_length = "3\r\nabc\r\n0\r\nTrace: value\r\n\r\n".len,
424         },
425     };
426 
427     for (cases) |case| {
428         for (0..case.encoded_length + 1) |split| {
429             var collector = FixedChunkCollector{};
430             var decoder = Decoder.init();
431             const before = try decoder.feed(case.wire[0..split], &collector);
432             const after = try decoder.feed(case.wire[split..], &collector);
433             const consumed = before + after;
434             try testing.expect(decoder.done);
435             try testing.expectEqual(case.encoded_length, consumed);
436             try testing.expectEqualStrings("abc", collector.body[0..collector.length]);
437             try testing.expectEqualStrings("TAIL", case.wire[consumed..]);
438         }
439     }
440 }
441 
442 test "chunk streaming decoder accepts every byte boundary without storage" {
443     const input = "4;kind=test\r\nWiki\r\n5\r\npedia\r\n0\r\nTrace: value\r\n\r\n";
444     var collector = FixedChunkCollector{};
445     var decoder = Decoder.init();
446     for (input, 0..) |_, index| {
447         _ = try decoder.feed(input[index..][0..1], &collector);
448     }
449     try testing.expect(decoder.done);
450     try testing.expectEqualStrings("Wikipedia", collector.body[0..collector.length]);
451 }
452 
453 test "chunk streaming decoder accepts split chunk terminator" {
454     var collector = TestCollector{ .allocator = testing.allocator };
455     defer collector.deinit();
456 
457     var decoder = Decoder.init();
458 
459     _ = try decoder.feed("4\r\nWiki", &collector);
460     try testing.expect(!decoder.done);
461     try testing.expectEqualStrings("Wiki", collector.buf.items);
462     _ = try decoder.feed("\r", &collector);
463     try testing.expect(!decoder.done);
464     _ = try decoder.feed("\n0\r\n\r\n", &collector);
465 
466     try testing.expect(decoder.done);
467     try testing.expectEqualStrings("Wiki", collector.buf.items);
468 }
469 
470 test "chunk streaming decoder waits for trailer terminator" {
471     var collector = TestCollector{ .allocator = testing.allocator };
472     defer collector.deinit();
473 
474     var decoder = Decoder.init();
475 
476     _ = try decoder.feed("0\r\n", &collector);
477     try testing.expect(!decoder.done);
478     _ = try decoder.feed("X-Trace: bench\r\n", &collector);
479     try testing.expect(!decoder.done);
480     _ = try decoder.feed("\r\n", &collector);
481     try testing.expect(decoder.done);
482     try testing.expectEqual(@as(usize, 0), collector.buf.items.len);
483 }
484 
485 test "chunk streaming decoder preserves callback error identity" {
486     var decoder = Decoder.init();
487 
488     var failing = RejectingChunkHandler{};
489     try testing.expectError(error.ChunkRejected, decoder.feed("5\r\nHello\r\n0\r\n\r\n", &failing));
490 }