limit.ha (1895B)
1 // SPDX-License-Identifier: MPL-2.0 2 // (c) Hare authors <https://harelang.org> 3 4 export type limitstream = struct { 5 vtable: stream, 6 source: handle, 7 limit: size, 8 }; 9 10 const limit_vtable_reader: vtable = vtable { 11 reader = &limit_read, 12 ... 13 }; 14 15 const limit_vtable_writer: vtable = vtable { 16 writer = &limit_write, 17 ... 18 }; 19 20 // Create an overlay stream that only allows a limited amount of bytes to be 21 // read from the underlying stream. This stream does not need to be closed, and 22 // closing it does not close the underlying stream. Reading any data beyond the 23 // given limit causes the reader to return [[EOF]]. 24 export fn limitreader(source: handle, limit: size) limitstream = { 25 return limitstream { 26 vtable = &limit_vtable_reader, 27 source = source, 28 limit = limit, 29 }; 30 }; 31 32 // Create an overlay stream that only allows a limited amount of bytes to be 33 // written to the underlying stream. This stream does not need to be closed, and 34 // closing it does not close the underlying stream. Writing beyond the given 35 // limit causes the writer to return short writes (as few as zero bytes). 36 export fn limitwriter(source: handle, limit: size) limitstream = { 37 return limitstream { 38 vtable = &limit_vtable_writer, 39 source = source, 40 limit = limit, 41 }; 42 }; 43 44 fn limit_read(s: *stream, buf: []u8) (size | EOF | error) = { 45 let stream = s: *limitstream; 46 if (stream.limit == 0) { 47 return EOF; 48 }; 49 if (len(buf) > stream.limit) { 50 buf = buf[..stream.limit]; 51 }; 52 match (read(stream.source, buf)) { 53 case EOF => 54 return EOF; 55 case let z: size => 56 stream.limit -= z; 57 return z; 58 }; 59 }; 60 61 fn limit_write(s: *stream, buf: const []u8) (size | error) = { 62 let stream = s: *limitstream; 63 if (stream.limit == 0) { 64 return 0z; 65 }; 66 let slice = if (len(buf) > stream.limit) { 67 yield buf[..stream.limit]; 68 } else { 69 yield buf[..]; 70 }; 71 const z = write(stream.source, slice)?; 72 stream.limit -= z; 73 return z; 74 };