hare

commit 5a92722a862ff8afa85a5be830910f1e02bd5fbe
parent 4786fb916a404d672f1617988c8b57d16e8b28c8
Author: Bor Grošelj Simić <bor.groseljsimic@telemach.net>
Date:   Sun,  7 Feb 2021 16:14:40 +0100

io: implement limit stream

Diffstat:
A
io/+test/limit.ha
 | 
36
++++++++++++++++++++++++++++++++++++
A
io/limit.ha
 | 
61
+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

2 files changed, 97 insertions(+), 0 deletions(-)
diff --git a/io/+test/limit.ha b/io/+test/limit.ha
@@ -0,0 +1,36 @@
+@test fn limit() void = {
+	let buf: [15z]u8 = [0u8...];
+	let source_stream = test_stream_open();
+	defer close(&source_stream.stream);
+
+	let r_stream = limit_reader(&source_stream.stream, 20z);
+	match (write(r_stream, buf)) {
+		unsupported => void,
+		* => abort(),
+	};
+	match (read(r_stream, buf)) {
+		n: size => assert(n == 15z),
+		error => abort(),
+	};
+	match (read(r_stream, buf)) {
+		n: size => assert(n == 5z),
+		error => abort(),
+	};
+	close(r_stream);
+
+	let w_stream = limit_writer(&source_stream.stream, 20z);
+	match (read(w_stream, buf)) {
+		unsupported => void,
+		* => abort(),
+	};
+	match (write(w_stream, buf)) {
+		n: size => assert(n == 15z),
+		error => abort(),
+	};
+	match (write(w_stream, buf)) {
+		n: size => assert(n == 5z),
+		error => abort(),
+	};
+	close(w_stream);
+
+};
diff --git a/io/limit.ha b/io/limit.ha
@@ -0,0 +1,61 @@
+use strings;
+
+type limited_stream = struct {
+	stream: stream,
+	source: *stream,
+	limit: size,
+};
+
+fn limited_stream_new(source: *stream, limit: size) *limited_stream = {
+	return alloc(*limited_stream,
+		limited_stream {
+			stream = stream {
+				name = strings::dup(source.name),
+				closer = &limited_close,
+				...
+			},
+			source = source,
+			limit = limit,
+	});
+};
+
+// Create an overlay stream that only allows a limited amount of bytes to be
+// read from the underlying stream.
+export fn limit_reader(source: *stream, limit: size) *stream = {
+	let stream = limited_stream_new(source, limit);
+	stream.stream.reader = &limited_read;
+	return &stream.stream;
+};
+
+// Create an overlay stream that only allows a limited amount of bytes to be
+// written to the underlying stream.
+export fn limit_writer(source: *stream, limit: size) *stream = {
+	let stream = limited_stream_new(source, limit);
+	stream.stream.writer = &limited_write;
+	return &stream.stream;
+};
+
+fn limited_read(s: *stream, buf: []u8) (size | EOF | error) = {
+	let stream = s: *limited_stream;
+	if (len(buf) > stream.limit) {
+		buf = buf[..stream.limit];
+	};
+	stream.limit -= len(buf);
+	return read(stream.source, buf);
+};
+
+fn limited_write(s: *stream, buf: const []u8) (size | error) = {
+	let stream = s: *limited_stream;
+	let slice = if (len(buf) > stream.limit) {
+		buf[..stream.limit];
+	} else {
+		buf[..];
+	};
+	stream.limit -= len(slice);
+	return write(stream.source, slice);
+};
+
+fn limited_close(s: *stream) void = {
+	free(s.name);
+	free(s);
+};