dragonfly_client_rs/
main.rs

1mod app_config;
2mod client;
3mod durable_cache;
4mod exts;
5mod reuse_cache;
6mod scan_cache;
7mod scanner;
8mod utils;
9
10use std::time::{Duration, Instant};
11
12use client::DragonflyClient;
13use color_eyre::eyre::{ensure, Result};
14use tracing::{error, info, span, warn, Level};
15use tracing_subscriber::EnvFilter;
16
17use crate::{
18    app_config::APP_CONFIG,
19    client::{Job, ScanResult, SubmitJobResultsError},
20    scanner::{scan_all_distributions, PackageScanResults},
21};
22
23fn scan_package(
24    client: &DragonflyClient,
25    job: &Job,
26    stats: &mut crate::reuse_cache::CacheStats,
27) -> Option<ScanResult> {
28    if job.hash != client.rules_state.hash {
29        warn!(
30            event = "scan_deferred",
31            required_rules_commit = %job.hash,
32            loaded_rules_commit = %client.rules_state.hash,
33            "Deferring job requiring rules commit {}; scanner loaded {}",
34            job.hash, client.rules_state.hash
35        );
36        // Do not submit a failure: mainframe will requeue the pending job after its timeout.
37        return None;
38    }
39
40    info!(
41        event = "scan_started",
42        distribution_count = job.distributions.len(),
43        rules_commit = %job.hash,
44        "Started package scan"
45    );
46
47    let result = match scan_all_distributions(
48        client.download_client(),
49        &client.rules_state.rules,
50        job,
51        Some(&client.reuse_cache),
52        stats,
53    ) {
54        Ok(results) => {
55            let package_scan_results = PackageScanResults::new(
56                job.name.clone(),
57                job.version.clone(),
58                job.attempt,
59                job.assignment_id.clone(),
60                results,
61                job.hash.clone(),
62            );
63            let body = package_scan_results.build_body();
64
65            Ok(body)
66        }
67        Err(err) => Err(SubmitJobResultsError {
68            name: job.name.clone(),
69            version: job.version.clone(),
70            attempt: job.attempt,
71            assignment_id: job.assignment_id.clone(),
72            reason: format!("{err}"),
73        }),
74    };
75    Some(result)
76}
77
78fn run_job(client: &DragonflyClient, job: &Job) {
79    let span = span!(
80        Level::INFO,
81        "scan_job",
82        package = %job.name,
83        version = %job.version,
84        attempt = job.attempt,
85        assignment_id = %job.assignment_id,
86        rules_commit = %job.hash,
87    );
88    let _enter = span.enter();
89    let started_at = Instant::now();
90
91    client.reuse_cache.begin_job(job);
92    let mut stats = crate::reuse_cache::CacheStats::new("yara", client.reuse_cache.mode);
93    let Some(scan_result) = scan_package(client, job, &mut stats) else {
94        return;
95    };
96
97    stats.emit();
98    let outcome = match &scan_result {
99        Ok(result) => {
100            info!(
101                event = "scan_completed",
102                elapsed_ms = started_at.elapsed().as_millis(),
103                distribution_count = job.distributions.len(),
104                score = result.score,
105                matched_rule_count = result.rules_matched.len(),
106                "Completed package scan"
107            );
108            "finished"
109        }
110        Err(result) => {
111            error!(
112                event = "scan_failed",
113                elapsed_ms = started_at.elapsed().as_millis(),
114                distribution_count = job.distributions.len(),
115                reason = %result.reason,
116                "Package scan failed"
117            );
118            "failed"
119        }
120    };
121
122    match client.send_result(scan_result, &stats) {
123        Ok(()) => info!(
124            event = "result_submitted",
125            elapsed_ms = started_at.elapsed().as_millis(),
126            outcome,
127            "Submitted package scan result"
128        ),
129        Err(err) => {
130            error!(
131                event = "result_submission_failed",
132                elapsed_ms = started_at.elapsed().as_millis(),
133                outcome,
134                error = %err,
135                "Error while sending package scan result to API"
136            );
137        }
138    }
139}
140
141fn run_jobs(client: &DragonflyClient, jobs: Vec<Job>, worker_count: usize) {
142    if worker_count == 1 {
143        for job in jobs {
144            run_job(client, &job);
145        }
146        return;
147    }
148
149    std::thread::scope(|scope| {
150        for job in jobs {
151            scope.spawn(move || run_job(client, &job));
152        }
153    });
154}
155
156fn main() -> Result<()> {
157    color_eyre::install()?;
158
159    let default_env_filter = EnvFilter::builder()
160        .parse("warn,dragonfly_client_rs=info")
161        .unwrap();
162    let env_filter = EnvFilter::try_from_default_env().unwrap_or(default_env_filter);
163
164    tracing_subscriber::fmt().with_env_filter(env_filter).init();
165    let mut client = DragonflyClient::new()?;
166    ensure!(
167        !APP_CONFIG.reuse_cache_database || APP_CONFIG.threads == 1,
168        "Database cache requires one scan thread"
169    );
170    ensure!(APP_CONFIG.reuse_cache_mode != crate::reuse_cache::CacheMode::Reuse || APP_CONFIG.threads == 1,
171        "Cross-package reuse requires DRAGONFLY_THREADS=1 so validation can invalidate the entire active job");
172    ensure!(
173        APP_CONFIG.threads > 0,
174        "DRAGONFLY_THREADS must be greater than zero"
175    );
176    ensure!(
177        APP_CONFIG.bulk_size > 0,
178        "DRAGONFLY_BULK_SIZE must be greater than zero"
179    );
180    let batch_size = APP_CONFIG.bulk_size.min(APP_CONFIG.threads);
181
182    loop {
183        info!("Fetching up to {batch_size} jobs");
184        match client.bulk_get_job(batch_size) {
185            Ok(jobs) if jobs.is_empty() => {
186                info!("No jobs found");
187                std::thread::sleep(Duration::from_secs(APP_CONFIG.load_duration));
188            }
189            Ok(jobs) => {
190                info!(
191                    event = "jobs_fetched",
192                    job_count = jobs.len(),
193                    "Fetched package scan jobs"
194                );
195
196                if jobs.iter().any(|job| job.hash != client.rules_state.hash) {
197                    info!(
198                        "At least one job requires rules other than {}, updating rules",
199                        client.rules_state.hash
200                    );
201
202                    if let Err(err) = client.update_rules() {
203                        error!("Error while updating rules: {err}");
204                    }
205                }
206
207                run_jobs(&client, jobs, APP_CONFIG.threads);
208            }
209
210            Err(err) => {
211                error!("Unexpected HTTP error: {err}");
212                std::thread::sleep(Duration::from_secs(APP_CONFIG.load_duration));
213            }
214        }
215    }
216}
dragonfly_client_rs/main.rs

dragonfly_client_rs/
main.rs
