1use std::{cell::RefCell, time::Duration};
8
9use curl::easy::{Easy2, Handler, WriteError};
10use mime::Mime;
11use tracing::{event, instrument, Level};
12use url::Url;
13
14use super::{filter_schemes, MimeData, ResourceUrlHandler};
15
16#[derive(Debug, Clone, Default)]
18pub struct CollectBuffer {
19 read_limit: u64,
20 buffer: Vec<u8>,
21}
22
23impl Handler for CollectBuffer {
24 fn write(&mut self, data: &[u8]) -> Result<usize, WriteError> {
25 if self.read_limit < (self.buffer.len() + data.len()).try_into().unwrap() {
26 Ok(0)
29 } else {
30 self.buffer.extend_from_slice(data);
31 Ok(data.len())
32 }
33 }
34}
35
36pub struct CurlResourceHandler {
38 easy: RefCell<Easy2<CollectBuffer>>,
39}
40
41impl CurlResourceHandler {
42 pub fn create(read_limit: u64, useragent: &str) -> std::io::Result<Self> {
47 let mut easy = Easy2::new(CollectBuffer {
48 buffer: Vec::new(),
49 read_limit,
50 });
51 easy.timeout(Duration::from_secs(1))?;
55 easy.connect_timeout(Duration::from_secs(1))?;
56 easy.follow_location(true)?;
57 easy.fail_on_error(true)?;
58 easy.tcp_nodelay(true)?;
59 easy.useragent(useragent)?;
60 Ok(Self::new(easy))
61 }
62
63 pub fn new(easy: Easy2<CollectBuffer>) -> Self {
65 Self {
66 easy: RefCell::new(easy),
67 }
68 }
69}
70
71impl ResourceUrlHandler for CurlResourceHandler {
72 #[instrument(level = "debug", skip(self), fields(url = %url))]
73 fn read_resource(&self, url: &Url) -> std::io::Result<MimeData> {
74 filter_schemes(&["http", "https", "ftp", "ftps", "smb"], url).and_then(|url| {
77 let mut easy = self.easy.borrow_mut();
78 easy.url(url.as_str())?;
79 easy.perform()?;
80
81 let mime_type = easy.content_type()?.and_then(|content_type| {
82 event!(
83 Level::DEBUG,
84 "Raw Content-Type of remote resource {}: {:?}",
85 &url,
86 content_type
87 );
88 content_type.parse::<Mime>().ok()
89 });
90 let data = easy.get_ref().buffer.clone();
91 easy.get_mut().buffer.clear();
92 Ok(MimeData { mime_type, data })
93 })
94 }
95}