1use std::sync::Arc;
2use std::time::Instant;
3
4use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
5use tokio::net::{lookup_host, TcpStream};
6use tokio_rustls::rustls::pki_types::ServerName;
7use tokio_rustls::rustls::{ClientConfig, RootCertStore};
8use tokio_rustls::TlsConnector;
9use url::Url;
10use webpki_roots::TLS_SERVER_ROOTS;
11
12use crate::engine::timing::{build_report_from_phases, TimingPhases};
13use crate::engine::tls::{report_for_http, report_for_https_failure, report_for_https_success};
14use crate::error::{Result, TloxError};
15use crate::models::{CertificateInfo, NewProbeResult, ProbeResult, ProbeStatus, TimingReport, TlsReport};
16use crate::storage::SqliteStore;
17
18#[derive(Debug, Clone, Default)]
19pub struct ProxyMeasurement {
20 phases: TimingPhases,
21}
22
23impl ProxyMeasurement {
24 pub fn new() -> Self {
25 Self::default()
26 }
27
28 pub fn record_dns_ms(&mut self, dns_ms: u64) {
29 self.phases.dns_ms = dns_ms;
30 }
31
32 pub fn record_tcp_ms(&mut self, tcp_ms: u64) {
33 self.phases.tcp_ms = tcp_ms;
34 }
35
36 pub fn record_tls_ms(&mut self, tls_ms: u64) {
37 self.phases.tls_ms = tls_ms;
38 }
39
40 pub fn record_ttfb_ms(&mut self, ttfb_ms: u64) {
41 self.phases.ttfb_ms = ttfb_ms;
42 }
43
44 pub fn record_download_ms(&mut self, download_ms: u64) {
45 self.phases.download_ms = download_ms;
46 }
47
48 pub fn finish_report(&self) -> TimingReport {
49 build_report_from_phases(self.phases)
50 }
51}
52
53#[derive(Debug, Clone, Default)]
54pub struct SimpleProbe;
55
56#[derive(Debug, Clone, PartialEq, Eq)]
57struct HttpResponseMetadata {
58 total_ms: u64,
59 status_code: Option<u16>,
60 status_text: Option<String>,
61 bytes: u64,
62}
63
64impl SimpleProbe {
65 pub fn new() -> Self {
66 Self
67 }
68
69 pub async fn probe(&self, target: &str) -> Result<NewProbeResult> {
70 let url = Url::parse(target).map_err(|_| TloxError::InvalidTarget(target.to_string()))?;
71 let is_https = url.scheme() == "https";
72 let host = url
73 .host_str()
74 .ok_or_else(|| TloxError::InvalidTarget(target.to_string()))?
75 .to_string();
76 let port = url
77 .port_or_known_default()
78 .ok_or_else(|| TloxError::InvalidTarget(target.to_string()))?;
79 let request_target = build_request_target(&url);
80
81 let mut measurement = ProxyMeasurement::new();
82 let total_start = Instant::now();
83
84 let dns_start = Instant::now();
85 let mut addresses = match lookup_host((host.as_str(), port)).await {
86 Ok(addresses) => addresses,
87 Err(error) => {
88 return Ok(failed_probe_result(
89 target,
90 &request_target,
91 finish_report(&measurement, total_start.elapsed().as_millis() as u64),
92 tls_report_for_failure(is_https, error.to_string()),
93 error.to_string(),
94 ));
95 }
96 };
97 measurement.record_dns_ms(dns_start.elapsed().as_millis() as u64);
98
99 let Some(address) = addresses.next() else {
100 let error_message = "no addresses resolved".to_string();
101
102 return Ok(failed_probe_result(
103 target,
104 &request_target,
105 finish_report(&measurement, total_start.elapsed().as_millis() as u64),
106 tls_report_for_failure(is_https, error_message.clone()),
107 error_message,
108 ));
109 };
110
111 let tcp_start = Instant::now();
112 let stream = match TcpStream::connect(address).await {
113 Ok(stream) => stream,
114 Err(error) => {
115 return Ok(failed_probe_result(
116 target,
117 &request_target,
118 finish_report(&measurement, total_start.elapsed().as_millis() as u64),
119 tls_report_for_failure(is_https, error.to_string()),
120 error.to_string(),
121 ));
122 }
123 };
124 measurement.record_tcp_ms(tcp_start.elapsed().as_millis() as u64);
125
126 if is_https {
127 self.probe_https(target, &host, &request_target, stream, measurement, total_start)
128 .await
129 } else {
130 self.probe_http(target, &host, &request_target, stream, measurement, total_start)
131 .await
132 }
133 }
134
135 pub async fn probe_and_store(&self, store: &SqliteStore, target: &str) -> Result<ProbeResult> {
136 let result = self.probe(target).await?;
137
138 store.save_result(&result)
139 }
140
141 async fn probe_http(
142 &self,
143 target: &str,
144 host: &str,
145 request_target: &str,
146 mut stream: TcpStream,
147 mut measurement: ProxyMeasurement,
148 total_start: Instant,
149 ) -> Result<NewProbeResult> {
150 match execute_request(
151 &mut stream,
152 host,
153 request_target,
154 &mut measurement,
155 total_start,
156 )
157 .await
158 {
159 Ok(response) => Ok(NewProbeResult {
160 target: target.to_string(),
161 request_path: request_target.to_string(),
162 report: finish_report(&measurement, response.total_ms),
163 status: ProbeStatus::Succeeded,
164 response_status_code: response.status_code,
165 response_status_text: response.status_text,
166 response_bytes: Some(response.bytes),
167 error_message: None,
168 tls: Some(report_for_http()),
169 }),
170 Err(error) => Ok(failed_probe_result(
171 target,
172 request_target,
173 finish_report(&measurement, total_start.elapsed().as_millis() as u64),
174 Some(report_for_http()),
175 error.to_string(),
176 )),
177 }
178 }
179
180 async fn probe_https(
181 &self,
182 target: &str,
183 host: &str,
184 request_target: &str,
185 stream: TcpStream,
186 mut measurement: ProxyMeasurement,
187 total_start: Instant,
188 ) -> Result<NewProbeResult> {
189 let tls_start = Instant::now();
190 let connector = build_tls_connector();
191 let server_name = ServerName::try_from(host.to_string())
192 .map_err(|_| TloxError::InvalidTarget(target.to_string()))?;
193
194 let mut tls_stream = match connector.connect(server_name, stream).await {
195 Ok(stream) => stream,
196 Err(error) => {
197 return Ok(failed_probe_result(
198 target,
199 request_target,
200 finish_report(&measurement, total_start.elapsed().as_millis() as u64),
201 Some(report_for_https_failure(error.to_string())),
202 error.to_string(),
203 ));
204 }
205 };
206 measurement.record_tls_ms(tls_start.elapsed().as_millis() as u64);
207
208 match execute_request(
209 &mut tls_stream,
210 host,
211 request_target,
212 &mut measurement,
213 total_start,
214 )
215 .await
216 {
217 Ok(response) => {
218 let (_, connection) = tls_stream.get_ref();
219 let tls = Some(report_for_https_success(
220 connection.protocol_version().map(|value| format!("{value:?}")),
221 connection
222 .negotiated_cipher_suite()
223 .map(|value| format!("{:?}", value.suite())),
224 extract_certificate_info(connection.peer_certificates()),
225 ));
226
227 Ok(NewProbeResult {
228 target: target.to_string(),
229 request_path: request_target.to_string(),
230 report: finish_report(&measurement, response.total_ms),
231 status: ProbeStatus::Succeeded,
232 response_status_code: response.status_code,
233 response_status_text: response.status_text,
234 response_bytes: Some(response.bytes),
235 error_message: None,
236 tls,
237 })
238 }
239 Err(error) => Ok(failed_probe_result(
240 target,
241 request_target,
242 finish_report(&measurement, total_start.elapsed().as_millis() as u64),
243 Some(report_for_https_failure(error.to_string())),
244 error.to_string(),
245 )),
246 }
247 }
248}
249
250fn failed_probe_result(
251 target: &str,
252 request_path: &str,
253 report: TimingReport,
254 tls: Option<TlsReport>,
255 error_message: String,
256) -> NewProbeResult {
257 NewProbeResult {
258 target: target.to_string(),
259 request_path: request_path.to_string(),
260 report,
261 status: ProbeStatus::Failed,
262 response_status_code: None,
263 response_status_text: None,
264 response_bytes: None,
265 error_message: Some(error_message),
266 tls,
267 }
268}
269
270fn finish_report(measurement: &ProxyMeasurement, total_ms: u64) -> TimingReport {
271 let mut report = measurement.finish_report();
272 report.total_ms = total_ms;
273 report
274}
275
276fn tls_report_for_failure(is_https: bool, error_message: String) -> Option<TlsReport> {
277 if is_https {
278 Some(report_for_https_failure(error_message))
279 } else {
280 Some(report_for_http())
281 }
282}
283
284fn build_request_target(url: &Url) -> String {
285 let mut target = url.path().to_string();
286 if target.is_empty() {
287 target.push('/');
288 }
289
290 if let Some(query) = url.query() {
291 target.push('?');
292 target.push_str(query);
293 }
294
295 target
296}
297
298fn build_tls_connector() -> TlsConnector {
299 let roots = RootCertStore::from_iter(TLS_SERVER_ROOTS.iter().cloned());
300 let config = ClientConfig::builder()
301 .with_root_certificates(roots)
302 .with_no_client_auth();
303
304 TlsConnector::from(Arc::new(config))
305}
306
307async fn execute_request<S>(
308 stream: &mut S,
309 host: &str,
310 request_target: &str,
311 measurement: &mut ProxyMeasurement,
312 total_start: Instant,
313) -> std::io::Result<HttpResponseMetadata>
314where
315 S: AsyncRead + AsyncWrite + Unpin,
316{
317 let request = format!(
318 "GET {request_target} HTTP/1.1\r\nHost: {host}\r\nUser-Agent: lobe/0.1\r\nAccept: */*\r\nConnection: close\r\n\r\n"
319 );
320
321 let ttfb_start = Instant::now();
322 stream.write_all(request.as_bytes()).await?;
323 stream.flush().await?;
324
325 let mut buffer = Vec::new();
326 let mut chunk = [0_u8; 1024];
327 let mut download_start: Option<Instant> = None;
328
329 loop {
330 let bytes_read = stream.read(&mut chunk).await?;
331 if bytes_read == 0 {
332 break;
333 }
334
335 if download_start.is_none() {
336 measurement.record_ttfb_ms(ttfb_start.elapsed().as_millis() as u64);
337 download_start = Some(Instant::now());
338 }
339
340 buffer.extend_from_slice(&chunk[..bytes_read]);
341 }
342
343 if let Some(start) = download_start {
344 measurement.record_download_ms(start.elapsed().as_millis() as u64);
345 }
346
347 let metadata = parse_http_response(&buffer);
348
349 Ok(HttpResponseMetadata {
350 total_ms: total_start.elapsed().as_millis() as u64,
351 status_code: metadata.status_code,
352 status_text: metadata.status_text,
353 bytes: metadata.bytes,
354 })
355}
356
357fn parse_http_response(buffer: &[u8]) -> HttpResponseMetadata {
358 let Some(header_end) = find_header_end(buffer) else {
359 return HttpResponseMetadata {
360 total_ms: 0,
361 status_code: None,
362 status_text: None,
363 bytes: buffer.len() as u64,
364 };
365 };
366
367 let header_bytes = &buffer[..header_end];
368 let body_bytes = &buffer[header_end + 4..];
369 let status_line = header_bytes
370 .split(|byte| *byte == b'\n')
371 .next()
372 .and_then(|line| std::str::from_utf8(line).ok())
373 .map(str::trim)
374 .unwrap_or_default();
375
376 let mut parts = status_line.splitn(3, ' ');
377 let _http_version = parts.next();
378 let status_code = parts.next().and_then(|part| part.parse::<u16>().ok());
379 let status_text = parts.next().map(|part| part.trim().to_string());
380
381 HttpResponseMetadata {
382 total_ms: 0,
383 status_code,
384 status_text,
385 bytes: body_bytes.len() as u64,
386 }
387}
388
389fn find_header_end(buffer: &[u8]) -> Option<usize> {
390 buffer.windows(4).position(|window| window == b"\r\n\r\n")
391}
392
393fn extract_certificate_info(
394 certificates: Option<&[tokio_rustls::rustls::pki_types::CertificateDer<'_>]>,
395) -> Option<CertificateInfo> {
396 let certificate = certificates?.first()?;
397 let (_, parsed) = x509_parser::parse_x509_certificate(certificate.as_ref()).ok()?;
398
399 Some(CertificateInfo {
400 subject: parsed.subject().to_string(),
401 issuer: parsed.issuer().to_string(),
402 not_before: Some(parsed.validity().not_before.to_string()),
403 not_after: Some(parsed.validity().not_after.to_string()),
404 })
405}
406
407#[cfg(test)]
408mod tests {
409 use url::Url;
410
411 use crate::models::ProbeStatus;
412 use crate::models::TlsStatus;
413 use crate::storage::SqliteStore;
414
415 use super::{build_request_target, ProxyMeasurement, SimpleProbe};
416
417 #[test]
418 fn finish_report_builds_report_from_recorded_phases() {
419 let mut measurement = ProxyMeasurement::new();
420 measurement.record_dns_ms(12);
421 measurement.record_tcp_ms(18);
422 measurement.record_tls_ms(24);
423 measurement.record_ttfb_ms(30);
424 measurement.record_download_ms(6);
425
426 let report = measurement.finish_report();
427
428 assert_eq!(report.dns_ms, 12);
429 assert_eq!(report.tcp_ms, 18);
430 assert_eq!(report.tls_ms, 24);
431 assert_eq!(report.ttfb_ms, 30);
432 assert_eq!(report.download_ms, 6);
433 assert_eq!(report.total_ms, 90);
434 }
435
436 #[test]
437 fn build_request_target_includes_query_string() {
438 let url = Url::parse("https://example.com/path?hello=world").expect("url should parse");
439
440 let target = build_request_target(&url);
441
442 assert_eq!(target, "/path?hello=world");
443 }
444
445 #[tokio::test]
446 async fn simple_probe_returns_failed_result_for_unreachable_target() {
447 let probe = SimpleProbe::new();
448 let result = probe
449 .probe("http://127.0.0.1:1")
450 .await
451 .expect("probe should return a result");
452
453 assert_eq!(result.status, ProbeStatus::Failed);
454 assert_eq!(result.report.tls_ms, 0);
455 assert_eq!(result.report.ttfb_ms, 0);
456 assert!(result.report.dns_ms <= result.report.total_ms);
457 assert!(result.error_message.is_some());
458 assert_eq!(
459 result.tls.expect("tls report should exist").status,
460 TlsStatus::NotUsed
461 );
462 }
463
464 #[tokio::test]
465 async fn simple_probe_rejects_invalid_targets() {
466 let probe = SimpleProbe::new();
467 let result = probe.probe("not-a-url").await;
468
469 assert!(matches!(
470 result,
471 Err(crate::error::TloxError::InvalidTarget(_))
472 ));
473 }
474
475 #[tokio::test]
476 async fn probe_and_store_persists_probe_result() {
477 let probe = SimpleProbe::new();
478 let store = SqliteStore::in_memory().expect("store should initialize");
479
480 let saved = probe
481 .probe_and_store(&store, "http://127.0.0.1:1")
482 .await
483 .expect("probe result should save");
484
485 let loaded = store
486 .latest_result_for_target("http://127.0.0.1:1")
487 .expect("latest result should load")
488 .expect("target should have a saved result");
489
490 assert_eq!(loaded.id, saved.id);
491 assert_eq!(loaded.target, "http://127.0.0.1:1");
492 assert_eq!(loaded.status, ProbeStatus::Failed);
493 assert!(loaded.error_message.is_some());
494 }
495}