dragonfly_client_rs/
client.rs

1mod methods;
2mod models;
3
4use flate2::read::GzDecoder;
5pub use methods::*;
6pub use models::*;
7use tempfile::{tempdir, tempfile, TempDir};
8
9use crate::APP_CONFIG;
10use color_eyre::{
11    eyre::{bail, ensure},
12    Result,
13};
14use reqwest::{
15    blocking::Client,
16    header::{HeaderMap, HeaderValue},
17    redirect::Policy,
18    Url,
19};
20use std::{
21    fs::File,
22    io::{self, Read, Seek},
23};
24
25#[derive(Clone, Copy)]
26struct ArchiveLimits {
27    entries: usize,
28    download_size: u64,
29    expanded_size: u64,
30    scan_size: u64,
31}
32
33impl ArchiveLimits {
34    fn configured() -> Self {
35        Self {
36            entries: APP_CONFIG.max_archive_entries,
37            download_size: APP_CONFIG.max_download_size,
38            expanded_size: APP_CONFIG.max_expanded_size,
39            scan_size: APP_CONFIG.max_scan_size,
40        }
41    }
42}
43
44pub struct RulesState {
45    pub rules: yara::Rules,
46    pub hash: String,
47}
48
49#[warn(clippy::module_name_repetitions)]
50pub struct DragonflyClient {
51    api_client: Client,
52    download_client: Client,
53    pub rules_state: RulesState,
54    pub reuse_cache: crate::reuse_cache::ReuseCache,
55    base_url: String,
56}
57
58impl DragonflyClient {
59    pub fn new() -> Result<Self> {
60        let api_client = build_api_http_client(
61            &APP_CONFIG.cf_access_client_id,
62            &APP_CONFIG.cf_access_client_secret,
63        )?;
64        let download_client = build_download_http_client()?;
65
66        let rules_response = fetch_rules(&api_client, &APP_CONFIG.base_url)?;
67
68        let mut reuse_cache = crate::reuse_cache::ReuseCache::new(
69            APP_CONFIG.reuse_cache_mode,
70            APP_CONFIG.reuse_cache_entries,
71            APP_CONFIG.reuse_cache_bytes,
72        );
73        if APP_CONFIG.reuse_cache_database {
74            reuse_cache.set_database(crate::durable_cache::DurableCache::new(
75                api_client.clone(),
76                &APP_CONFIG.base_url,
77                "yara",
78                &rules_response.hash,
79                &rules_response.rules,
80                None,
81            )?);
82        }
83        let rules_state = RulesState {
84            rules: rules_response.compile()?,
85            hash: rules_response.hash,
86        };
87
88        Ok(Self {
89            api_client,
90            download_client,
91            rules_state,
92            reuse_cache,
93            base_url: APP_CONFIG.base_url.clone(),
94        })
95    }
96
97    /// Update the global ruleset. Waits for a write lock.
98    pub fn update_rules(&mut self) -> Result<()> {
99        let response = fetch_rules(&self.api_client, &self.base_url)?;
100        let compiled_rules = response.compile()?;
101        if APP_CONFIG.reuse_cache_database {
102            self.reuse_cache
103                .set_database(crate::durable_cache::DurableCache::new(
104                    self.api_client.clone(),
105                    &self.base_url,
106                    "yara",
107                    &response.hash,
108                    &response.rules,
109                    None,
110                )?);
111        }
112        self.reuse_cache.clear();
113        self.rules_state.rules = compiled_rules;
114        self.rules_state.hash = response.hash;
115
116        Ok(())
117    }
118
119    pub fn bulk_get_job(&self, n_jobs: usize) -> reqwest::Result<Vec<Job>> {
120        fetch_bulk_job(&self.api_client, &self.base_url, n_jobs)
121    }
122
123    /// Send a [`crate::client::models::ScanResult`] to mainframe
124    pub fn send_result(
125        &self,
126        body: models::ScanResult,
127        stats: &crate::reuse_cache::CacheStats,
128    ) -> reqwest::Result<()> {
129        send_result_with_metrics(&self.api_client, &self.base_url, body, Some(stats))
130    }
131
132    /// Return the client used for uncredentialed distribution downloads.
133    pub(crate) fn download_client(&self) -> &Client {
134        &self.download_client
135    }
136}
137
138fn build_api_http_client(client_id: &str, client_secret: &str) -> Result<Client> {
139    let mut headers = HeaderMap::new();
140    headers.insert("CF-Access-Client-Id", HeaderValue::from_str(client_id)?);
141
142    let mut secret = HeaderValue::from_str(client_secret)?;
143    secret.set_sensitive(true);
144    headers.insert("CF-Access-Client-Secret", secret);
145
146    Ok(Client::builder()
147        .gzip(true)
148        .redirect(Policy::custom(|attempt| {
149            attempt.error("Dragonfly API redirects are not allowed")
150        }))
151        .default_headers(headers)
152        .build()?)
153}
154
155fn build_download_http_client() -> reqwest::Result<Client> {
156    Client::builder().gzip(true).build()
157}
158
159fn stage_download<R: Read>(response: R, limit: u64) -> Result<File> {
160    let read_limit = limit
161        .checked_add(1)
162        .ok_or_else(|| color_eyre::eyre::eyre!("download size limit is too large"))?;
163    let mut file = tempfile()?;
164    let downloaded = io::copy(&mut response.take(read_limit), &mut file)?;
165    ensure!(
166        downloaded <= limit,
167        "compressed distribution exceeds the {limit}-byte download limit"
168    );
169    file.rewind()?;
170
171    Ok(file)
172}
173
174/// Unpack a tarball within the configured resource limits.
175fn extract_tarball(file: File, limits: ArchiveLimits) -> Result<TempDir> {
176    let mut tarball = tar::Archive::new(GzDecoder::new(file));
177    let tmpdir = tempdir()?;
178    let mut entries = 0_usize;
179    let mut expanded_size = 0_u64;
180
181    for entry in tarball.entries()? {
182        let mut entry = entry?;
183        entries = entries
184            .checked_add(1)
185            .ok_or_else(|| color_eyre::eyre::eyre!("tar entry count overflowed"))?;
186        ensure!(
187            entries <= limits.entries,
188            "tar archive exceeds the {}-entry limit",
189            limits.entries
190        );
191
192        let entry_size = entry.size();
193        ensure!(
194            entry_size <= limits.scan_size,
195            "tar entry {} is {entry_size} bytes, exceeding the {}-byte scan limit",
196            entry.path()?.display(),
197            limits.scan_size
198        );
199        expanded_size = expanded_size
200            .checked_add(entry_size)
201            .ok_or_else(|| color_eyre::eyre::eyre!("expanded tar size overflowed"))?;
202        ensure!(
203            expanded_size <= limits.expanded_size,
204            "tar archive exceeds the {}-byte expanded-size limit",
205            limits.expanded_size
206        );
207        ensure!(
208            entry.unpack_in(tmpdir.path())?,
209            "tar entry would unpack outside the temporary directory"
210        );
211    }
212
213    Ok(tmpdir)
214}
215
216fn zip_entry_count(file: &mut File) -> Result<usize> {
217    const END_HEADER_SIZE: usize = 22;
218    const MAX_COMMENT_SIZE: usize = u16::MAX as usize;
219    const SIGNATURE: [u8; 4] = [0x50, 0x4b, 0x05, 0x06];
220
221    let file_size = file.seek(io::SeekFrom::End(0))?;
222    let tail_size = file_size.min((END_HEADER_SIZE + MAX_COMMENT_SIZE) as u64);
223    file.seek(io::SeekFrom::End(
224        -i64::try_from(tail_size).expect("ZIP footer window fits in i64"),
225    ))?;
226
227    let mut tail = vec![0_u8; usize::try_from(tail_size)?];
228    file.read_exact(&mut tail)?;
229    let Some(header_start) =
230        (0..=tail.len().saturating_sub(SIGNATURE.len()))
231            .rev()
232            .find(|&index| {
233                if tail[index..].get(..SIGNATURE.len()) != Some(&SIGNATURE) {
234                    return false;
235                }
236                let Some(comment_size) = tail
237                    .get(index + 20..index + 22)
238                    .map(|bytes| u16::from_le_bytes([bytes[0], bytes[1]]) as usize)
239                else {
240                    return false;
241                };
242                tail.len() - index == END_HEADER_SIZE + comment_size
243            })
244    else {
245        bail!("ZIP end-of-central-directory record is missing");
246    };
247    let header = &tail[header_start..];
248    let disk_number = u16::from_le_bytes([header[4], header[5]]);
249    let directory_disk = u16::from_le_bytes([header[6], header[7]]);
250    let entries_on_disk = u16::from_le_bytes([header[8], header[9]]);
251    let entries = u16::from_le_bytes([header[10], header[11]]);
252    ensure!(
253        disk_number == 0 && directory_disk == 0 && entries_on_disk == entries,
254        "multi-disk ZIP archives are not supported"
255    );
256    ensure!(
257        entries != u16::MAX,
258        "ZIP64 archives are not accepted in the constrained scanner"
259    );
260    file.rewind()?;
261
262    Ok(entries as usize)
263}
264
265/// Extract a ZIP archive within the configured resource limits.
266fn extract_zipfile(mut file: File, limits: ArchiveLimits) -> Result<TempDir> {
267    let entry_count = zip_entry_count(&mut file)?;
268    ensure!(
269        entry_count <= limits.entries,
270        "ZIP archive contains {entry_count} entries, exceeding the {}-entry limit",
271        limits.entries
272    );
273    let mut zip = zip::ZipArchive::new(file)?;
274    ensure!(
275        zip.len() == entry_count,
276        "ZIP central-directory entry count changed while parsing"
277    );
278
279    let mut expanded_size = 0_u64;
280    for index in 0..zip.len() {
281        let entry = zip.by_index(index)?;
282        ensure!(
283            matches!(
284                entry.compression(),
285                zip::CompressionMethod::Stored | zip::CompressionMethod::Deflated
286            ),
287            "ZIP entry {} uses unsupported compression method {:?}",
288            entry.name(),
289            entry.compression()
290        );
291        let entry_size = entry.size();
292        ensure!(
293            entry_size <= limits.scan_size,
294            "ZIP entry {} is {entry_size} bytes, exceeding the {}-byte scan limit",
295            entry.name(),
296            limits.scan_size
297        );
298        expanded_size = expanded_size
299            .checked_add(entry_size)
300            .ok_or_else(|| color_eyre::eyre::eyre!("expanded ZIP size overflowed"))?;
301        ensure!(
302            expanded_size <= limits.expanded_size,
303            "ZIP archive exceeds the {}-byte expanded-size limit",
304            limits.expanded_size
305        );
306    }
307
308    let tmpdir = tempdir()?;
309    zip.extract(tmpdir.path())?;
310
311    Ok(tmpdir)
312}
313
314pub fn download_distribution(http_client: &Client, download_url: Url) -> Result<TempDir> {
315    let is_tarball = download_url.as_str().ends_with(".tar.gz");
316    let response = http_client.get(download_url).send()?.error_for_status()?;
317    let limits = ArchiveLimits::configured();
318    let file = stage_download(response, limits.download_size)?;
319
320    if is_tarball {
321        extract_tarball(file, limits)
322    } else {
323        extract_zipfile(file, limits)
324    }
325}
326
327#[cfg(test)]
328mod tests {
329    use super::{
330        build_api_http_client, build_download_http_client, extract_tarball, extract_zipfile,
331        stage_download, ArchiveLimits, DragonflyClient, RulesState,
332    };
333    use flate2::{write::GzEncoder, Compression};
334    use std::{
335        io::{Cursor, Read, Write},
336        net::TcpListener,
337        sync::mpsc,
338        thread,
339    };
340    use yara::Compiler;
341    use zip::{write::SimpleFileOptions, ZipWriter};
342
343    const CLIENT_ID: &str = "test-client.access";
344    const CLIENT_SECRET: &str = "test-secret";
345
346    fn archive_limits() -> ArchiveLimits {
347        ArchiveLimits {
348            entries: 4,
349            download_size: 1024,
350            expanded_size: 1024,
351            scan_size: 1024,
352        }
353    }
354
355    fn build_zip(files: &[(&str, &[u8])]) -> Vec<u8> {
356        let mut zip = ZipWriter::new(Cursor::new(Vec::new()));
357        for (name, contents) in files {
358            zip.start_file(*name, SimpleFileOptions::default()).unwrap();
359            zip.write_all(contents).unwrap();
360        }
361        zip.finish().unwrap().into_inner()
362    }
363
364    fn build_tarball(files: &[(&str, &[u8])]) -> Vec<u8> {
365        let encoder = GzEncoder::new(Vec::new(), Compression::fast());
366        let mut tarball = tar::Builder::new(encoder);
367        for (name, contents) in files {
368            let mut header = tar::Header::new_gnu();
369            header.set_size(contents.len().try_into().unwrap());
370            header.set_cksum();
371            tarball.append_data(&mut header, name, *contents).unwrap();
372        }
373        tarball.into_inner().unwrap().finish().unwrap()
374    }
375
376    fn serve_once(response: String) -> (String, mpsc::Receiver<String>) {
377        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
378        let address = listener.local_addr().unwrap();
379        let (sender, receiver) = mpsc::channel();
380
381        thread::spawn(move || {
382            let (mut stream, _) = listener.accept().unwrap();
383            let mut request = Vec::new();
384            let mut buffer = [0_u8; 4096];
385
386            loop {
387                let bytes_read = stream.read(&mut buffer).unwrap();
388                if bytes_read == 0 {
389                    break;
390                }
391                request.extend_from_slice(&buffer[..bytes_read]);
392                if request.windows(4).any(|window| window == b"\r\n\r\n") {
393                    break;
394                }
395            }
396
397            sender.send(String::from_utf8(request).unwrap()).unwrap();
398            stream.write_all(response.as_bytes()).unwrap();
399        });
400
401        (format!("http://{address}/example.tar.gz"), receiver)
402    }
403
404    #[test]
405    fn distribution_downloads_do_not_send_cloudflare_access_credentials() {
406        let response =
407            String::from("HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n");
408        let (download_url, request) = serve_once(response);
409        let rules = Compiler::new()
410            .unwrap()
411            .add_rules_str("rule test_rule { condition: false }")
412            .unwrap()
413            .compile_rules()
414            .unwrap();
415        let client = DragonflyClient {
416            api_client: build_api_http_client(CLIENT_ID, CLIENT_SECRET).unwrap(),
417            download_client: build_download_http_client().unwrap(),
418            rules_state: RulesState {
419                rules,
420                hash: String::new(),
421            },
422            reuse_cache: crate::reuse_cache::ReuseCache::new(
423                crate::reuse_cache::CacheMode::Off,
424                0,
425                0,
426            ),
427            base_url: String::from("https://dragonfly.example"),
428        };
429
430        client
431            .download_client()
432            .get(download_url)
433            .send()
434            .unwrap()
435            .error_for_status()
436            .unwrap();
437
438        let request = request.recv().unwrap().to_ascii_lowercase();
439        assert!(!request.contains("\r\ncf-access-client-id:"));
440        assert!(!request.contains("\r\ncf-access-client-secret:"));
441        assert!(!request.contains("\r\nauthorization:"));
442    }
443
444    #[test]
445    fn api_client_rejects_redirects_before_forwarding_credentials() {
446        let redirect_target = TcpListener::bind("127.0.0.1:0").unwrap();
447        redirect_target.set_nonblocking(true).unwrap();
448        let target_url = format!("http://{}/capture", redirect_target.local_addr().unwrap());
449        let response = format!(
450            "HTTP/1.1 302 Found\r\nLocation: {target_url}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"
451        );
452        let (api_url, source_request) = serve_once(response);
453        let api_client = build_api_http_client(CLIENT_ID, CLIENT_SECRET).unwrap();
454
455        let error = api_client.get(api_url).send().unwrap_err();
456
457        assert!(error.is_redirect());
458        assert_eq!(
459            redirect_target.accept().unwrap_err().kind(),
460            std::io::ErrorKind::WouldBlock
461        );
462        let source_request = source_request.recv().unwrap().to_ascii_lowercase();
463        assert!(source_request.contains("\r\ncf-access-client-id:"));
464        assert!(source_request.contains("\r\ncf-access-client-secret:"));
465    }
466
467    #[test]
468    fn stage_download_rejects_oversized_input() {
469        let error = stage_download(Cursor::new(vec![0_u8; 5]), 4).unwrap_err();
470
471        assert!(error.to_string().contains("4-byte download limit"));
472    }
473
474    #[test]
475    fn zip_entry_limit_is_checked_before_extraction() {
476        let bytes = build_zip(&[("one", b"1"), ("two", b"2")]);
477        let file = stage_download(Cursor::new(bytes), 1024).unwrap();
478        let limits = ArchiveLimits {
479            entries: 1,
480            ..archive_limits()
481        };
482
483        let error = extract_zipfile(file, limits).unwrap_err();
484
485        assert!(error.to_string().contains("2 entries"));
486    }
487
488    #[test]
489    fn zip_expanded_size_is_bounded() {
490        let bytes = build_zip(&[("one", b"1234"), ("two", b"5678")]);
491        let file = stage_download(Cursor::new(bytes), 1024).unwrap();
492        let limits = ArchiveLimits {
493            expanded_size: 7,
494            ..archive_limits()
495        };
496
497        let error = extract_zipfile(file, limits).unwrap_err();
498
499        assert!(error.to_string().contains("expanded-size limit"));
500    }
501
502    #[test]
503    fn zip_file_scan_size_is_bounded() {
504        let bytes = build_zip(&[("large", b"12345")]);
505        let file = stage_download(Cursor::new(bytes), 1024).unwrap();
506        let limits = ArchiveLimits {
507            scan_size: 4,
508            ..archive_limits()
509        };
510
511        let error = extract_zipfile(file, limits).unwrap_err();
512
513        assert!(error.to_string().contains("4-byte scan limit"));
514    }
515
516    #[test]
517    fn tar_expanded_size_is_bounded() {
518        let bytes = build_tarball(&[("one", b"1234"), ("two", b"5678")]);
519        let file = stage_download(Cursor::new(bytes), 1024).unwrap();
520        let limits = ArchiveLimits {
521            expanded_size: 7,
522            ..archive_limits()
523        };
524
525        let error = extract_tarball(file, limits).unwrap_err();
526
527        assert!(error.to_string().contains("expanded-size limit"));
528    }
529}
dragonfly_client_rs/client.rs

dragonfly_client_rs/
client.rs
