Skip to content

Commit 27aba2d

Browse files
achan1989andrewrk
authored andcommitted
Fix infinite loop in Reader.Limited
1 parent 163ebe0 commit 27aba2d

File tree

1 file changed

+63
-3
lines changed

1 file changed

+63
-3
lines changed

lib/std/Io/Reader/Limited.zig

Lines changed: 63 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,11 @@ pub fn init(reader: *Reader, limit: Limit, buffer: []u8) Limited {
2828
fn stream(r: *Reader, w: *Writer, limit: Limit) Reader.StreamError!usize {
2929
const l: *Limited = @fieldParentPtr("interface", r);
3030
const combined_limit = limit.min(l.remaining);
31-
const n = try l.unlimited.stream(w, combined_limit);
32-
l.remaining = l.remaining.subtract(n).?;
33-
return n;
31+
if (combined_limit.nonzero()) {
32+
const n = try l.unlimited.stream(w, combined_limit);
33+
l.remaining = l.remaining.subtract(n).?;
34+
return n;
35+
} else return error.EndOfStream;
3436
}
3537

3638
test stream {
@@ -49,6 +51,64 @@ test stream {
4951
try std.testing.expectEqualStrings("test", result_buf[0..streamed]);
5052
}
5153

54+
test "readSliceAll from infinite source" {
55+
const InfSource = struct {
56+
reader: std.Io.Reader,
57+
58+
pub fn init(buffer: []u8) @This() {
59+
return @This(){
60+
.reader = .{
61+
.vtable = &.{
62+
.stream = streamA,
63+
},
64+
.buffer = buffer,
65+
.seek = 0,
66+
.end = 0,
67+
},
68+
};
69+
}
70+
71+
fn streamA(io_reader: *std.Io.Reader, w: *std.Io.Writer, limit: std.Io.Limit) std.Io.Reader.StreamError!usize {
72+
_ = io_reader;
73+
74+
std.debug.assert(limit.nonzero());
75+
76+
const n_bytes_remaining = limit.minInt(2);
77+
for (0..n_bytes_remaining) |_| {
78+
try w.writeByte('A');
79+
}
80+
return n_bytes_remaining;
81+
}
82+
};
83+
84+
// Exact size
85+
{
86+
var inf_buf: [10]u8 = undefined;
87+
var inf_stream = InfSource.init(&inf_buf);
88+
89+
var limit_buf: [2]u8 = undefined;
90+
var limited: std.Io.Reader.Limited = .init(&inf_stream.reader, .limited(2), &limit_buf);
91+
const limited_reader = &limited.interface;
92+
93+
var out_buffer: [2]u8 = undefined;
94+
try std.testing.expectEqual({}, limited_reader.readSliceAll(&out_buffer));
95+
try std.testing.expectEqualStrings("AA", &out_buffer);
96+
}
97+
98+
// Too large
99+
{
100+
var inf_buf: [10]u8 = undefined;
101+
var inf_stream = InfSource.init(&inf_buf);
102+
103+
var limit_buf: [2]u8 = undefined;
104+
var limited: std.Io.Reader.Limited = .init(&inf_stream.reader, .limited(2), &limit_buf);
105+
const limited_reader = &limited.interface;
106+
107+
var out_buffer: [8]u8 = undefined;
108+
try std.testing.expectError(error.EndOfStream, limited_reader.readSliceAll(&out_buffer));
109+
}
110+
}
111+
52112
fn discard(r: *Reader, limit: Limit) Reader.Error!usize {
53113
const l: *Limited = @fieldParentPtr("interface", r);
54114
const combined_limit = limit.min(l.remaining);

0 commit comments

Comments
 (0)