Skip to main content

sim_lib_agent_runner_http/probe/
transport.rs

1use super::{ProbeHttpRequest, ProbeHttpResponse, ProbeTransport};
2#[cfg(feature = "tls")]
3use rustls::{ClientConfig, ClientConnection, RootCertStore, StreamOwned, pki_types::ServerName};
4use sim_kernel::{Error, Result};
5use sim_lib_net_core::{
6    HeadOutcome, HttpBodyMode, UrlParts, body_mode, parse_http_head, read_head_until_double_crlf,
7};
8#[cfg(feature = "tls")]
9use std::sync::Arc;
10use std::{
11    io::{BufRead, BufReader, Read, Write},
12    net::TcpStream,
13};
14
15/// Plain HTTP/HTTPS probe transport used by the agent-facing `provider/probe` function.
16#[derive(Clone, Copy, Debug, Default)]
17pub struct HttpProbeTransport;
18
19impl ProbeTransport for HttpProbeTransport {
20    fn get(&self, request: ProbeHttpRequest<'_>) -> Result<ProbeHttpResponse> {
21        let stream = TcpStream::connect((
22            request.endpoint_parts.host.as_str(),
23            request.endpoint_parts.port,
24        ))
25        .map_err(|error| {
26            Error::HostError(format!("provider/probe io {:?}: {}", error.kind(), error))
27        })?;
28        stream
29            .set_read_timeout(Some(request.timeout))
30            .map_err(map_io_error)?;
31        stream
32            .set_write_timeout(Some(request.timeout))
33            .map_err(map_io_error)?;
34        let mut stream = connect_probe_stream(&request.endpoint_parts, stream)?;
35
36        let target = join_paths(&request.endpoint_parts.path, request.path);
37        write!(
38            stream,
39            "GET {target} HTTP/1.1\r\nHost: {}\r\nAccept: application/json\r\nConnection: close\r\n",
40            host_header(&request.endpoint_parts)
41        )
42        .map_err(map_io_error)?;
43        for (name, value) in &request.headers {
44            write!(stream, "{name}: {value}\r\n").map_err(map_io_error)?;
45        }
46        write!(stream, "\r\n").map_err(map_io_error)?;
47        stream.flush().map_err(map_io_error)?;
48        read_probe_response(&mut stream, request.max_response_bytes)
49    }
50}
51
52trait ProbeReadWrite: Read + Write {}
53
54impl<T: Read + Write> ProbeReadWrite for T {}
55
56#[cfg(feature = "tls")]
57fn connect_probe_stream(parts: &UrlParts, stream: TcpStream) -> Result<Box<dyn ProbeReadWrite>> {
58    match parts.scheme.as_str() {
59        "http" => Ok(Box::new(stream)),
60        "https" => {
61            let server_name = ServerName::try_from(parts.host.clone()).map_err(|_| {
62                Error::Eval(format!(
63                    "provider/probe invalid tls server name {}",
64                    parts.host
65                ))
66            })?;
67            let connection = ClientConnection::new(tls_client_config()?, server_name)
68                .map_err(|error| Error::HostError(error.to_string()))?;
69            Ok(Box::new(StreamOwned::new(connection, stream)))
70        }
71        other => Err(Error::Eval(format!(
72            "provider/probe unsupported url scheme {other}"
73        ))),
74    }
75}
76
77#[cfg(not(feature = "tls"))]
78fn connect_probe_stream(parts: &UrlParts, stream: TcpStream) -> Result<Box<dyn ProbeReadWrite>> {
79    match parts.scheme.as_str() {
80        "http" => Ok(Box::new(stream)),
81        "https" => Err(Error::Eval(
82            "provider/probe https endpoints require the tls feature".to_owned(),
83        )),
84        other => Err(Error::Eval(format!(
85            "provider/probe unsupported url scheme {other}"
86        ))),
87    }
88}
89
90#[cfg(feature = "tls")]
91fn tls_client_config() -> Result<Arc<ClientConfig>> {
92    let mut roots = RootCertStore::empty();
93    let cert_result = rustls_native_certs::load_native_certs();
94    for certificate in cert_result.certs {
95        roots
96            .add(certificate)
97            .map_err(|error| Error::HostError(error.to_string()))?;
98    }
99    if roots.is_empty() {
100        return Err(Error::HostError(
101            "provider/probe no tls root certificates available".to_owned(),
102        ));
103    }
104    Ok(Arc::new(
105        ClientConfig::builder()
106            .with_root_certificates(roots)
107            .with_no_client_auth(),
108    ))
109}
110
111fn read_probe_response(
112    stream: &mut dyn Read,
113    max_response_bytes: usize,
114) -> Result<ProbeHttpResponse> {
115    let mut reader = BufReader::new(stream);
116    let head = match read_head_until_double_crlf(&mut reader, 64 * 1024) {
117        Ok(HeadOutcome::Head(head)) => head,
118        Ok(HeadOutcome::TooLarge) => {
119            return Err(Error::HostError(
120                "provider/probe http headers exceed size limit".to_owned(),
121            ));
122        }
123        Ok(HeadOutcome::Eof | HeadOutcome::Truncated(_)) => {
124            return Err(Error::HostError(
125                "provider/probe http response ended before headers".to_owned(),
126            ));
127        }
128        Err(error) => return Err(Error::HostError(error.to_string())),
129    };
130    let head_text = std::str::from_utf8(&head)
131        .map_err(|_| Error::HostError("provider/probe http headers are not utf-8".to_owned()))?;
132    let parsed = parse_http_head(head_text).map_err(|error| Error::HostError(error.to_string()))?;
133    let mode = body_mode(&parsed).map_err(|error| Error::HostError(error.to_string()))?;
134    let body = match mode {
135        HttpBodyMode::ContentLength(length) => {
136            read_content_length(&mut reader, length, max_response_bytes)?
137        }
138        HttpBodyMode::Chunked => read_chunked(&mut reader, max_response_bytes)?,
139        HttpBodyMode::UntilEof | HttpBodyMode::Empty => {
140            read_to_end(&mut reader, max_response_bytes)?
141        }
142    };
143    Ok(ProbeHttpResponse { head: parsed, body })
144}
145
146fn read_content_length(
147    reader: &mut dyn Read,
148    length: usize,
149    max_response_bytes: usize,
150) -> Result<Vec<u8>> {
151    if length > max_response_bytes {
152        return Err(Error::Eval(format!(
153            "provider/probe response exceeded max output bytes {max_response_bytes}"
154        )));
155    }
156    let mut body = vec![0u8; length];
157    reader.read_exact(&mut body).map_err(map_io_error)?;
158    Ok(body)
159}
160
161fn read_to_end(reader: &mut dyn Read, max_response_bytes: usize) -> Result<Vec<u8>> {
162    let mut body = Vec::new();
163    let mut chunk = [0u8; 8192];
164    loop {
165        let read = reader.read(&mut chunk).map_err(map_io_error)?;
166        if read == 0 {
167            break;
168        }
169        body.extend_from_slice(&chunk[..read]);
170        if body.len() > max_response_bytes {
171            return Err(Error::Eval(format!(
172                "provider/probe response exceeded max output bytes {max_response_bytes}"
173            )));
174        }
175    }
176    Ok(body)
177}
178
179fn read_chunked(reader: &mut dyn BufRead, max_response_bytes: usize) -> Result<Vec<u8>> {
180    let mut body = Vec::new();
181    loop {
182        let mut size_line = String::new();
183        reader.read_line(&mut size_line).map_err(map_io_error)?;
184        let raw_size = size_line.trim_end_matches("\r\n");
185        let size_text = raw_size.split(';').next().unwrap_or(raw_size);
186        let size = usize::from_str_radix(size_text, 16)
187            .map_err(|_| Error::HostError("provider/probe invalid chunk size".to_owned()))?;
188        if size == 0 {
189            drain_chunk_trailers(reader)?;
190            return Ok(body);
191        }
192        if body.len().saturating_add(size) > max_response_bytes {
193            return Err(Error::Eval(format!(
194                "provider/probe response exceeded max output bytes {max_response_bytes}"
195            )));
196        }
197        let mut chunk = vec![0u8; size];
198        reader.read_exact(&mut chunk).map_err(map_io_error)?;
199        body.extend_from_slice(&chunk);
200        let mut crlf = [0u8; 2];
201        reader.read_exact(&mut crlf).map_err(map_io_error)?;
202        if crlf != *b"\r\n" {
203            return Err(Error::HostError(
204                "provider/probe invalid chunk terminator".to_owned(),
205            ));
206        }
207    }
208}
209
210fn drain_chunk_trailers(reader: &mut dyn BufRead) -> Result<()> {
211    loop {
212        let mut line = String::new();
213        reader.read_line(&mut line).map_err(map_io_error)?;
214        if line == "\r\n" || line.is_empty() {
215            return Ok(());
216        }
217    }
218}
219
220fn join_paths(base: &str, suffix: &str) -> String {
221    let base = base.trim_end_matches('/');
222    let suffix = suffix.trim_start_matches('/');
223    if base.is_empty() {
224        format!("/{suffix}")
225    } else {
226        format!("{base}/{suffix}")
227    }
228}
229
230fn host_header(parts: &UrlParts) -> String {
231    let default_port = (parts.scheme == "http" && parts.port == 80)
232        || (parts.scheme == "https" && parts.port == 443);
233    if default_port {
234        parts.host.clone()
235    } else {
236        format!("{}:{}", parts.host, parts.port)
237    }
238}
239
240fn map_io_error(error: std::io::Error) -> Error {
241    Error::HostError(format!("provider/probe io {:?}: {}", error.kind(), error))
242}