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