dia_files/
limit.rs

1/*
2==--==--==--==--==--==--==--==--==--==--==--==--==--==--==--==--
3
4Dia-Files
5
6Copyright (C) 2019-2024  Anonymous
7
8There are several releases over multiple years,
9they are listed as ranges, such as: "2019-2024".
10
11This program is free software: you can redistribute it and/or modify
12it under the terms of the GNU Lesser General Public License as published by
13the Free Software Foundation, either version 3 of the License, or
14(at your option) any later version.
15
16This program is distributed in the hope that it will be useful,
17but WITHOUT ANY WARRANTY; without even the implied warranty of
18MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
19GNU Lesser General Public License for more details.
20
21You should have received a copy of the GNU Lesser General Public License
22along with this program.  If not, see <https://www.gnu.org/licenses/>.
23
24::--::--::--::--::--::--::--::--::--::--::--::--::--::--::--::--
25*/
26
27use {
28    std::io::{Error, ErrorKind, StdinLock},
29    crate::Result,
30};
31
32#[cfg(not(feature="async-std"))]
33use std::{
34    fs::File,
35    io::{self, BufRead, BufReader, Read, Stdin},
36    path::Path,
37};
38
39#[cfg(feature="async-std")]
40use {
41    core::{
42        pin::Pin,
43        task::{Context, Poll},
44    },
45    crate::AsyncBufReader as BufReader,
46    async_std::{
47        fs::File,
48        io::{self, Read, ReadExt, Stdin},
49        path::Path,
50    },
51};
52
53#[cfg(test)]
54mod tests;
55
56macro_rules! read_to_end { ($self: ident, $capacity: ident) => {{
57    let mut result = Vec::with_capacity(match $capacity {
58        None => usize::MIN,
59        Some(capacity) => match $self.limit.checked_sub($self.read) {
60            None => return Err(err!()),
61            Some(max) => match usize::try_from(max) {
62                Ok(max) => max.min(capacity),
63                Err(_) => capacity,
64            },
65        },
66    });
67    async_call!($self.inner.read_to_end(&mut result))?;
68    Ok(result)
69}}}
70
71macro_rules! read { ($stream: ident, $limit: ident, $capacity: ident) => {{
72    async_call!(Self::new($stream, $limit).read_to_end($capacity))
73}}}
74
75macro_rules! read_to_string { ($stream: ident, $limit: ident, $capacity: ident) => {{
76    let data = read!($stream, $limit, $capacity)?;
77    match String::from_utf8(data) {
78        Ok(s) => Ok(s),
79        Err(err) => Err(Error::new(ErrorKind::InvalidInput, __!("{}", err.utf8_error()))),
80    }
81}}}
82
83macro_rules! open_file { ($path: ident, $limit: ident) => {{
84    let path = $path.as_ref();
85
86    match async_call!(path.metadata())?.len() {
87        size => if size > $limit {
88            return Err(err!("File too large: {size} (limit: {limit})", limit=$limit));
89        },
90    };
91
92    Ok(Self::new(async_call!(File::open(path))?, $limit))
93}}}
94
95macro_rules! read_file { ($path: ident, $limit: ident) => {{
96    let path = $path.as_ref();
97    let mut reader = BufReader::new(async_call!(Self::open_file(path, $limit))?);
98    let mut result = Vec::with_capacity(async_call!(path.metadata())?.len().try_into().map_err(|_| err!())?);
99    loop {
100        let count = {
101            let tmp = async_call!(reader.fill_buf())?;
102            if tmp.is_empty() {
103                break;
104            }
105            result.extend(tmp);
106            tmp.len()
107        };
108        reader.consume(count);
109    }
110    Result::Ok(result)
111}}}
112
113macro_rules! read_file_to_string { ($path: ident, $limit: ident) => {{
114    let data = read_file!($path, $limit)?;
115    match String::from_utf8(data) {
116        Ok(s) => Ok(s),
117        Err(err) => Err(Error::new(ErrorKind::InvalidInput, __!("{}", err.utf8_error()))),
118    }
119}}}
120
121/// # Limit
122///
123/// This struct can be used for limiting data read from some source.
124#[derive(Debug)]
125pub struct Limit<R> {
126    inner: R,
127    read: u64,
128    limit: u64,
129}
130
131impl<R> Limit<R> {
132
133    /// # Makes new object
134    pub fn new(inner: R, limit: u64) -> Self {
135        Self {
136            inner,
137            read: u64::MIN,
138            limit,
139        }
140    }
141
142}
143
144#[cfg(not(feature="async-std"))]
145#[doc(cfg(not(feature="async-std")))]
146impl<R> Limit<R> where R: Read {
147
148    /// # Reads all bytes
149    ///
150    /// Be careful using this function (in case your source produces a lot of data).
151    pub fn read_to_end(&mut self, capacity: Option<usize>) -> Result<Vec<u8>> {
152        read_to_end!(self, capacity)
153    }
154
155    /// # Reads a whole stream
156    pub fn read(stream: R, limit: u64, capacity: Option<usize>) -> Result<Vec<u8>> {
157        read!(stream, limit, capacity)
158    }
159
160    /// # Reads a whole stream, then converts data into string
161    pub fn read_to_string(stream: R, limit: u64, capacity: Option<usize>) -> Result<String> {
162        read_to_string!(stream, limit, capacity)
163    }
164
165}
166
167#[cfg(feature="async-std")]
168#[doc(cfg(feature="async-std"))]
169impl<R> Limit<R> where R: Unpin + Read {
170
171    /// # Reads all bytes
172    ///
173    /// Be careful using this function (in case your source produces a lot of data).
174    pub async fn read_to_end(&mut self, capacity: Option<usize>) -> Result<Vec<u8>> {
175        read_to_end!(self, capacity)
176    }
177
178    /// # Reads a whole stream
179    pub async fn read(stream: R, limit: u64, capacity: Option<usize>) -> Result<Vec<u8>> {
180        read!(stream, limit, capacity)
181    }
182
183    /// # Reads a whole stream, then converts data into string
184    pub async fn read_to_string(stream: R, limit: u64, capacity: Option<usize>) -> Result<String> {
185        read_to_string!(stream, limit, capacity)
186    }
187
188}
189
190impl Limit<File> {
191
192    /// # Opens a file for reading with a limit
193    #[cfg(not(feature="async-std"))]
194    #[doc(cfg(not(feature="async-std")))]
195    pub fn open_file<P>(path: P, limit: u64) -> Result<Self> where P: AsRef<Path> {
196        open_file!(path, limit)
197    }
198
199    /// # Opens a file for reading with a limit
200    #[cfg(feature="async-std")]
201    #[doc(cfg(feature="async-std"))]
202    pub async fn open_file<P>(path: P, limit: u64) -> Result<Self> where P: AsRef<Path> {
203        open_file!(path, limit)
204    }
205
206    /// # Reads a whole file
207    #[cfg(not(feature="async-std"))]
208    #[doc(cfg(not(feature="async-std")))]
209    pub fn read_file<P>(path: P, limit: u64) -> Result<Vec<u8>> where P: AsRef<Path> {
210        read_file!(path, limit)
211    }
212
213    /// # Reads a whole file
214    #[cfg(feature="async-std")]
215    #[doc(cfg(feature="async-std"))]
216    pub async fn read_file<P>(path: P, limit: u64) -> Result<Vec<u8>> where P: AsRef<Path> {
217        read_file!(path, limit)
218    }
219
220    /// # Reads a whole file to a string
221    #[cfg(not(feature="async-std"))]
222    #[doc(cfg(not(feature="async-std")))]
223    pub fn read_file_to_string<P>(path: P, limit: u64) -> Result<String> where P: AsRef<Path> {
224        read_file_to_string!(path, limit)
225    }
226
227    /// # Reads a whole file to a string
228    #[cfg(feature="async-std")]
229    #[doc(cfg(feature="async-std"))]
230    pub async fn read_file_to_string<P>(path: P, limit: u64) -> Result<String> where P: AsRef<Path> {
231        read_file_to_string!(path, limit)
232    }
233
234}
235
236impl Limit<Stdin> {
237
238    /// # Makes new object from standard input
239    pub fn stdin(limit: u64) -> Self {
240        Self::new(io::stdin(), limit)
241    }
242
243}
244
245impl Limit<StdinLock<'_>> {
246
247    /// # Locks standard input and makes new object from it
248    pub fn lock_stdin(limit: u64) -> Self {
249        Self::new(std::io::stdin().lock(), limit)
250    }
251}
252
253macro_rules! check { ($self: ident) => {
254    if $self.read > $self.limit {
255        let err = Err(err!("Limit: {limit}, already read: {read}", limit=$self.limit, read=$self.read));
256        #[cfg(feature="async-std")]
257        let err = Poll::Ready(err);
258        return err;
259    }
260}}
261
262macro_rules! update_read { ($self: ident, $more: ident) => {{
263    match u64::try_from($more).map_err(|_| err!("Cannot convert {more} into u64", more=$more))?.checked_add($self.read) {
264        None => {
265            let err = Err(err!("Failed: {read} + {more}", read=$self.read, more=$more));
266            #[cfg(feature="async-std")]
267            let err = Poll::Ready(err);
268            return err;
269        },
270        Some(read) => $self.read = read,
271    };
272}}}
273
274#[cfg(not(feature="async-std"))]
275#[doc(cfg(not(feature="async-std")))]
276impl<R> Read for Limit<R> where R: Read {
277
278    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
279        check!(self);
280        let result = self.inner.read(buf)?;
281
282        update_read!(self, result);
283        check!(self);
284
285        Ok(result)
286    }
287
288}
289
290#[cfg(feature="async-std")]
291#[doc(cfg(feature="async-std"))]
292impl<R> Read for Limit<R> where R: Unpin + Read {
293
294    fn poll_read(self: Pin<&mut Self>, context: &mut Context<'_>, buf: &mut [u8]) -> Poll<Result<usize>> {
295        let limit = self.get_mut();
296
297        check!(limit);
298        let result = match Read::poll_read(Pin::new(&mut limit.inner), context, buf) {
299            Poll::Ready(Ok(count)) => count,
300            Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
301            Poll::Pending => return Poll::Pending,
302        };
303
304        update_read!(limit, result);
305        check!(limit);
306
307        Poll::Ready(Ok(result))
308    }
309
310}
dia_files/limit.rs

dia_files/
limit.rs
