diff --git a/src/filter/fetcher.zig b/src/filter/fetcher.zig index 5a97192..56a03c0 100644 --- a/src/filter/fetcher.zig +++ b/src/filter/fetcher.zig @@ -101,22 +101,45 @@ pub const Fetcher = struct { } const body = resp.reader(self.transfer_buf); - var total: u64 = 0; - while (true) { - const n = body.readSliceShort(self.transfer_buf) catch |err| - return mapError(err, .receive); - if (n == 0) break; - total += n; - if (total > max_body_bytes) return error.BodyTooLarge; - // The caller owns `w` and can read the concrete failure from its - // own writer; this taxonomy has no member for a failing sink. - w.writeAll(self.transfer_buf[0..n]) catch return error.Unexpected; - } - - return .{ .bytes_read = total, .status = resp.head.status }; + return .{ .bytes_read = try pumpBody(body, w, max_body_bytes), .status = resp.head.status }; } }; +/// Streams the body into `w`, refusing more than `cap` bytes. `fetch` always +/// passes `max_body_bytes`; the parameter exists so the refusal is testable +/// without a 64 MiB body. +/// +/// `w` is the only destination: `Reader.stream` hands the reader's buffered +/// bytes straight to the writer, so the body is never copied twice and the +/// reader's own buffer is never a destination slice. Reading into that buffer +/// instead — `readSliceShort(self.transfer_buf)` on a reader constructed over +/// `self.transfer_buf` — makes `@memcpy` copy the buffer onto itself +/// (Reader.zig:677) and panics the process on the first read that finds bytes +/// already buffered. +fn pumpBody(body: *std.Io.Reader, w: *std.Io.Writer, cap: usize) Error!u64 { + var total: usize = 0; + while (total < cap) { + total += body.stream(w, .limited(cap - total)) catch |err| switch (err) { + error.EndOfStream => return total, + // The caller owns `w` and can read the concrete failure from its + // own writer; this taxonomy has no member for a failing sink. + error.WriteFailed => return error.Unexpected, + else => |e| return mapError(e, .receive), + }; + } + + // `cap` written exactly. One more byte separates a body that fits from one + // that was cut off, and it must not reach `w`. + var probe_buf: [1]u8 = undefined; + var probe: std.Io.Writer.Discarding = .init(&probe_buf); + const extra = body.stream(&probe.writer, .limited(1)) catch |err| switch (err) { + error.EndOfStream => return total, + error.WriteFailed => return error.Unexpected, + else => |e| return mapError(e, .receive), + }; + return if (extra == 0) total else error.BodyTooLarge; +} + /// A scheme other than `http`/`https`, an unparseable URL and a URL with no /// host are one fault to the operator: the source row is unusable. fn parseUrl(url: []const u8) Error!std.Uri { @@ -170,6 +193,52 @@ fn narrowLocal(local: transport.ExchangeError) Error { const testing = std.testing; +// `Reader.fixed` hands back a reader whose whole payload is already buffered +// (seek 0, end len) — the state that made the previous implementation copy the +// reader's buffer onto itself and abort the process mid-download. + +test "pumpBody streams a fully buffered body without aliasing its source" { + const payload = "0.0.0.0 ads.example.com\n0.0.0.0 tracker.example\n"; + var body: std.Io.Reader = .fixed(payload); + var out: [payload.len]u8 = undefined; + var sink: std.Io.Writer = .fixed(&out); + + const n = try pumpBody(&body, &sink, max_body_bytes); + try testing.expectEqual(@as(u64, payload.len), n); + try testing.expectEqualStrings(payload, out[0..payload.len]); +} + +test "pumpBody accepts a body of exactly the cap" { + const payload = "abcdefgh"; + var body: std.Io.Reader = .fixed(payload); + var out: [payload.len]u8 = undefined; + var sink: std.Io.Writer = .fixed(&out); + + const n = try pumpBody(&body, &sink, payload.len); + try testing.expectEqual(@as(u64, payload.len), n); + try testing.expectEqualStrings(payload, out[0..payload.len]); +} + +test "pumpBody refuses a body one byte over the cap" { + const payload = "abcdefgh"; + var body: std.Io.Reader = .fixed(payload); + var out: [payload.len]u8 = undefined; + var sink: std.Io.Writer = .fixed(&out); + + try testing.expectError( + error.BodyTooLarge, + pumpBody(&body, &sink, payload.len - 1), + ); +} + +test "pumpBody reports an empty body as zero bytes" { + var body: std.Io.Reader = .fixed(""); + var sink_buf: [0]u8 = .{}; + var discarding: std.Io.Writer.Discarding = .init(&sink_buf); + + try testing.expectEqual(@as(u64, 0), try pumpBody(&body, &discarding.writer, max_body_bytes)); +} + fn undefinedFetcher(transfer_buf: []u8, redirect_buf: []u8) Fetcher { // `http` is never driven: every test below asserts a rejection that // happens before the first client call.