lobe_core/engine/
proxy.rs

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

lobe_core/engine/
proxy.rs
