data_source/
lib.rs

1use std::{
2    collections::HashMap,
3    io,
4    sync::{Arc, Mutex},
5    time::SystemTime,
6};
7
8use log::{debug, warn};
9
10#[derive(thiserror::Error, Debug)]
11pub enum FetchError {
12    #[cfg(feature = "reqwest")]
13    #[error("reqwest err")]
14    R(#[from] reqwest::Error),
15    #[error("io err")]
16    I(#[from] io::Error),
17    #[error("time err")]
18    T(#[from] std::time::SystemTimeError),
19    #[error("size limit exceed")]
20    S,
21}
22
23#[cfg(feature = "reqwest")]
24#[derive(Clone, Debug, Default)]
25pub struct HttpSource {
26    pub url: String,
27    pub proxy: Option<String>,
28    pub custom_request_headers: Option<Vec<(String, String)>>,
29    pub should_use_proxy: bool,
30    pub size_limit_bytes: Option<usize>,
31    pub update_interval_seconds: Option<u64>,
32
33    pub cached_file: Arc<Mutex<Option<String>>>,
34
35    pub cache_file_path: Option<String>,
36}
37#[cfg(feature = "reqwest")]
38impl HttpSource {
39    pub fn get_url(
40        &self,
41        c: reqwest::blocking::Client,
42    ) -> reqwest::Result<reqwest::blocking::Response> {
43        let mut rb = c.get(&self.url);
44        if let Some(h) = &self.custom_request_headers {
45            for h in h.iter() {
46                rb = rb.header(&h.0, &h.1);
47            }
48        }
49        rb.send()
50    }
51    pub fn set_proxy(
52        &self,
53        mut cb: reqwest::blocking::ClientBuilder,
54    ) -> reqwest::Result<reqwest::blocking::ClientBuilder> {
55        let ps = self.proxy.as_ref().unwrap();
56        let proxy = reqwest::Proxy::https(ps)?;
57        cb = cb.proxy(proxy);
58        let proxy = reqwest::Proxy::http(ps)?;
59        Ok(cb.proxy(proxy))
60    }
61
62    pub fn fetch(&self) -> Result<Vec<u8>, FetchError> {
63        let ocf = { self.cached_file.lock().unwrap().take() };
64        if let Some(cf) = &ocf {
65            if std::fs::exists(cf)? {
66                let mut ok = true;
67                if let Some(uis) = self.update_interval_seconds {
68                    let m = std::fs::metadata(cf)?;
69                    let lu = m.modified()?;
70                    let d = SystemTime::now().duration_since(lu)?.as_secs();
71                    if d <= uis {
72                        let _ = self.cached_file.lock().unwrap().insert(cf.to_string());
73                    } else {
74                        ok = false;
75                    }
76                }
77                if ok {
78                    let s = std::fs::read_to_string(cf)?;
79                    return Ok(s.into_bytes());
80                }
81            }
82        }
83
84        let mut cb = reqwest::blocking::ClientBuilder::new();
85        if self.should_use_proxy {
86            cb = self.set_proxy(cb)?;
87        }
88        let c = cb.build()?;
89        let r = self.get_url(c);
90        let r = match r {
91            Ok(r) => r,
92            Err(e) => {
93                if !self.should_use_proxy && self.proxy.is_some() {
94                    let mut cb = reqwest::blocking::ClientBuilder::new();
95                    cb = self.set_proxy(cb)?;
96                    let c = cb.build()?;
97                    self.get_url(c)?
98                } else {
99                    return Err(FetchError::R(e));
100                }
101            }
102        };
103        if let Some(sl) = self.size_limit_bytes {
104            if let Some(s) = r.content_length() {
105                if s as usize > sl {
106                    return Err(FetchError::S);
107                }
108            }
109        }
110        let b = r.bytes()?;
111        let v = b.to_vec();
112        if let Some(cfp) = &self.cache_file_path {
113            let r = std::fs::write(cfp, &v);
114            match r {
115                Ok(_) => {
116                    let _ = self.cached_file.lock().unwrap().insert(cfp.to_string());
117                }
118                Err(e) => warn!("write to cache file failed, {e}"),
119            }
120        }
121        Ok(v)
122    }
123}
124
125#[derive(Clone, Debug)]
126pub enum SingleFileSource {
127    #[cfg(feature = "reqwest")]
128    Http(HttpSource),
129    FilePath(String),
130    Inline(Vec<u8>),
131}
132impl Default for SingleFileSource {
133    fn default() -> Self {
134        Self::Inline(Vec::new())
135    }
136}
137/// Defines where to get the content of the requested file name.
138///
139/// 很多配置中 都要再加载其他外部文件,
140/// FileSource 限定了 查找文件的 路径 和 来源, 读取文件时只会限制在这个范围内,
141/// 这样就增加了安全性
142#[derive(Clone, Debug, Default)]
143pub enum DataSource {
144    #[default]
145    StdReadFile,
146    Folders(Vec<String>), //从指定的一组路径来寻找文件
147
148    #[cfg(feature = "tar")]
149    Tar(Vec<u8>), // 从一个 已放到内存中的 tar 中 寻找文件
150
151    /// 与其它方式不同，FileMap 存储名称的映射表, 无需遍历目录
152    FileMap(HashMap<String, SingleFileSource>),
153}
154
155impl DataSource {
156    pub fn insert_current_working_dir(&mut self) -> io::Result<()> {
157        if let DataSource::Folders(ref mut v) = self {
158            v.push(std::env::current_dir()?.to_string_lossy().to_string())
159        }
160        Ok(())
161    }
162
163    pub fn read_to_string<P>(&self, file_name: P) -> io::Result<String>
164    where
165        P: AsRef<std::path::Path>,
166    {
167        let r = self.get_file_content(file_name)?;
168        Ok(String::from_utf8_lossy(r.0.as_slice()).to_string())
169    }
170
171    /// 返回读到的 数据。可能还会返回 成功找到的路径
172    pub fn get_file_content<P>(&self, file_name: P) -> io::Result<(Vec<u8>, Option<String>)>
173    where
174        P: AsRef<std::path::Path>,
175    {
176        match self {
177            #[cfg(feature = "tar")]
178            DataSource::Tar(tar_binary) => get_file_from_tar(file_name, tar_binary),
179
180            DataSource::Folders(possible_addrs) => {
181                for dir in possible_addrs {
182                    let real_file_name = std::path::Path::new(dir).join(file_name.as_ref());
183
184                    if real_file_name.exists() {
185                        return std::fs::read(&real_file_name).map(|v| (v, Some(dir.to_owned())));
186                    }
187                }
188                Err(std::io::Error::new(
189                    std::io::ErrorKind::NotFound,
190                    format!(
191                        "File not found in specified directories: {:?}",
192                        file_name.as_ref()
193                    ),
194                ))
195            }
196            DataSource::StdReadFile => {
197                let s = std::fs::read_to_string(file_name)?;
198                Ok((s.into_bytes(), None))
199            }
200
201            DataSource::FileMap(map) => {
202                let r = map.get(&file_name.as_ref().to_string_lossy().to_string());
203
204                match r {
205                    Some(sf) => match sf {
206                        #[cfg(feature = "reqwest")]
207                        SingleFileSource::Http(http_source) => {
208                            let r = http_source.fetch();
209                            match r {
210                                Ok(v) => Ok((v, None)),
211                                Err(e) => Err(io::Error::other(e)),
212                            }
213                        }
214                        SingleFileSource::FilePath(f) => {
215                            let s = std::fs::read_to_string(f)?;
216                            Ok((s.into_bytes(), None))
217                        }
218                        SingleFileSource::Inline(v) => Ok((v.clone(), None)),
219                    },
220                    None => Err(io::Error::new(io::ErrorKind::NotFound, "")),
221                }
222            }
223        }
224    }
225}
226
227#[cfg(feature = "tar")]
228pub fn get_file_from_tar<P>(
229    file_name_in_tar: P,
230    tar_binary: &Vec<u8>,
231) -> io::Result<(Vec<u8>, Option<String>)>
232where
233    P: AsRef<std::path::Path>,
234{
235    let mut a = tar::Archive::new(std::io::Cursor::new(tar_binary));
236
237    debug!(
238        "finding {} from tar, tar whole size is {}",
239        file_name_in_tar.as_ref().to_str().unwrap(),
240        tar_binary.len()
241    );
242
243    let mut e = a
244        .entries()
245        .unwrap()
246        .find(|a| {
247            a.as_ref()
248                .is_ok_and(|b| b.path().is_ok_and(|c| c == file_name_in_tar.as_ref()))
249        })
250        .ok_or_else(|| {
251            io::Error::new(
252                io::ErrorKind::NotFound,
253                format!(
254                    "get_file_from_tar: can't find the file, {}",
255                    file_name_in_tar.as_ref().to_str().unwrap()
256                ),
257            )
258        })??;
259
260    debug!("found {}", file_name_in_tar.as_ref().to_str().unwrap());
261
262    let mut result = vec![];
263    use std::io::Read;
264    e.read_to_end(&mut result)?;
265    Ok((
266        result,
267        Some(e.path().unwrap().to_str().unwrap().to_string()),
268    ))
269}
data_source/lib.rs

data_source/
lib.rs
