dragonfly_client_rs/client/
methods.rs

1use super::{models, ScanResultSerializer};
2
3use reqwest::blocking::Client;
4
5pub fn fetch_bulk_job(
6    http_client: &Client,
7    base_url: &str,
8    n_jobs: usize,
9) -> reqwest::Result<Vec<models::Job>> {
10    http_client
11        .post(format!("{base_url}/jobs"))
12        .query(&[("batch", n_jobs)])
13        .send()?
14        .error_for_status()?
15        .json()
16}
17
18pub fn fetch_rules(http_client: &Client, base_url: &str) -> reqwest::Result<models::RulesResponse> {
19    http_client
20        .get(format!("{base_url}/rules"))
21        .send()?
22        .error_for_status()?
23        .json()
24}
25
26#[cfg(test)]
27pub fn send_result(
28    http_client: &Client,
29    base_url: &str,
30    body: models::ScanResult,
31) -> reqwest::Result<()> {
32    send_result_with_metrics(http_client, base_url, body, None)
33}
34
35pub fn send_result_with_metrics(
36    http_client: &Client,
37    base_url: &str,
38    body: models::ScanResult,
39    stats: Option<&crate::reuse_cache::CacheStats>,
40) -> reqwest::Result<()> {
41    #[derive(serde::Serialize)]
42    struct Report<'a> {
43        #[serde(flatten)]
44        result: ScanResultSerializer,
45        #[serde(skip_serializing_if = "Option::is_none")]
46        scan_reuse: Option<&'a crate::reuse_cache::CacheStats>,
47    }
48    let body = Report {
49        result: body.into(),
50        scan_reuse: stats,
51    };
52    http_client
53        .put(format!("{base_url}/package"))
54        .json(&body)
55        .send()?
56        .error_for_status()?;
57    Ok(())
58}
59
60#[cfg(test)]
61mod tests {
62    use super::{fetch_bulk_job, fetch_rules, send_result};
63    use crate::client::{build_api_http_client, SubmitJobResultsError};
64    use std::{
65        io::{Read, Write},
66        net::TcpListener,
67        sync::mpsc,
68        thread,
69    };
70
71    const CLIENT_ID: &str = "test-client.access";
72    const CLIENT_SECRET: &str = "test-secret";
73
74    fn serve_once(response_body: &str) -> (String, mpsc::Receiver<String>) {
75        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
76        let address = listener.local_addr().unwrap();
77        let response_body = response_body.to_owned();
78        let (sender, receiver) = mpsc::channel();
79
80        thread::spawn(move || {
81            let (mut stream, _) = listener.accept().unwrap();
82            let mut request = Vec::new();
83            let mut buffer = [0_u8; 4096];
84
85            loop {
86                let bytes_read = stream.read(&mut buffer).unwrap();
87                if bytes_read == 0 {
88                    break;
89                }
90                request.extend_from_slice(&buffer[..bytes_read]);
91
92                let request_text = String::from_utf8_lossy(&request);
93                let Some(headers_end) = request_text.find("\r\n\r\n") else {
94                    continue;
95                };
96                let content_length = request_text[..headers_end]
97                    .lines()
98                    .find_map(|line| {
99                        line.to_ascii_lowercase()
100                            .strip_prefix("content-length: ")
101                            .and_then(|value| value.parse::<usize>().ok())
102                    })
103                    .unwrap_or_default();
104                if request.len() >= headers_end + 4 + content_length {
105                    break;
106                }
107            }
108
109            sender.send(String::from_utf8(request).unwrap()).unwrap();
110            write!(
111                stream,
112                "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
113                response_body.len(),
114                response_body,
115            )
116            .unwrap();
117        });
118
119        (format!("http://{address}"), receiver)
120    }
121
122    fn assert_cloudflare_access_headers(request: &str) {
123        let lowercase_request = request.to_ascii_lowercase();
124        assert!(lowercase_request.contains(&format!(
125            "\r\ncf-access-client-id: {}\r\n",
126            CLIENT_ID.to_ascii_lowercase()
127        )));
128        assert!(lowercase_request
129            .contains(&format!("\r\ncf-access-client-secret: {CLIENT_SECRET}\r\n")));
130        assert!(!lowercase_request.contains("\r\nauthorization:"));
131    }
132
133    #[test]
134    fn telemetry_is_flattened_into_success_and_failure_wire_payloads() {
135        use crate::{
136            client::SubmitJobResultsSuccess,
137            reuse_cache::{CacheMode, CacheStats},
138        };
139        for success in [true, false] {
140            let (base_url, request) = serve_once("");
141            let client = build_api_http_client(CLIENT_ID, CLIENT_SECRET).unwrap();
142            let result = if success {
143                Ok(SubmitJobResultsSuccess {
144                    name: "example".into(),
145                    version: "1".into(),
146                    attempt: 1,
147                    assignment_id: "lease".into(),
148                    score: 0,
149                    inspector_url: None,
150                    rules_matched: Vec::new(),
151                    commit: "rules".into(),
152                })
153            } else {
154                Err(SubmitJobResultsError {
155                    name: "example".into(),
156                    version: "1".into(),
157                    attempt: 1,
158                    assignment_id: "lease".into(),
159                    reason: "failure".into(),
160                })
161            };
162            let mut stats = CacheStats::new("yara", CacheMode::Reuse);
163            stats.reused_files = 7;
164            stats.engine_us = 123;
165            super::send_result_with_metrics(&client, &base_url, result, Some(&stats)).unwrap();
166            let request = request.recv().unwrap();
167            assert!(request.starts_with("PUT /package HTTP/1.1\r\n"));
168            assert_cloudflare_access_headers(&request);
169            let body: serde_json::Value =
170                serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap();
171            assert_eq!(body["name"], "example");
172            assert_eq!(body["assignment_id"], "lease");
173            assert_eq!(body["scan_reuse"]["mode"], "reuse");
174            assert_eq!(body["scan_reuse"]["reused_files"], 7);
175            assert_eq!(body["scan_reuse"]["engine_us"], 123);
176            assert_eq!(body.get("commit").is_some(), success);
177            assert_eq!(body.get("reason").is_some(), !success);
178            assert!(body.get("result").is_none());
179        }
180    }
181
182    #[test]
183    fn jobs_route_uses_cloudflare_access_service_token() {
184        let (base_url, request) = serve_once("[]");
185        let client = build_api_http_client(CLIENT_ID, CLIENT_SECRET).unwrap();
186
187        let jobs = fetch_bulk_job(&client, &base_url, 3).unwrap();
188
189        assert!(jobs.is_empty());
190        let request = request.recv().unwrap();
191        assert!(request.starts_with("POST /jobs?batch=3 HTTP/1.1\r\n"));
192        assert_cloudflare_access_headers(&request);
193    }
194
195    #[test]
196    fn rules_route_uses_cloudflare_access_service_token() {
197        let (base_url, request) = serve_once(r#"{"hash":"abc123","rules":{}}"#);
198        let client = build_api_http_client(CLIENT_ID, CLIENT_SECRET).unwrap();
199
200        let rules = fetch_rules(&client, &base_url).unwrap();
201
202        assert_eq!(rules.hash, "abc123");
203        let request = request.recv().unwrap();
204        assert!(request.starts_with("GET /rules HTTP/1.1\r\n"));
205        assert_cloudflare_access_headers(&request);
206    }
207
208    #[test]
209    fn package_route_uses_cloudflare_access_service_token() {
210        let (base_url, request) = serve_once("");
211        let client = build_api_http_client(CLIENT_ID, CLIENT_SECRET).unwrap();
212        let result = Err(SubmitJobResultsError {
213            name: "example".to_owned(),
214            version: "1.0.0".to_owned(),
215            attempt: 2,
216            assignment_id: "0962b72e-f197-41c3-a059-e10fdc149cce".to_owned(),
217            reason: "test failure".to_owned(),
218        });
219
220        send_result(&client, &base_url, result).unwrap();
221
222        let request = request.recv().unwrap();
223        assert!(request.starts_with("PUT /package HTTP/1.1\r\n"));
224        assert_cloudflare_access_headers(&request);
225        assert!(request.contains(
226            r#"{"name":"example","version":"1.0.0","attempt":2,"assignment_id":"0962b72e-f197-41c3-a059-e10fdc149cce","reason":"test failure"}"#
227        ));
228    }
229}
dragonfly_client_rs/client/methods.rs

dragonfly_client_rs/client/
methods.rs
