netget/client/influxdb/
transport.rs

1use crate::server::influxdb::codec::{self, WriteBatch};
2use anyhow::{ensure, Context, Result};
3use bytes::Bytes;
4use http_body_util::{BodyExt, Full, Limited};
5use hyper::{
6    client::conn::http1,
7    header::{
8        AUTHORIZATION, CONNECTION, CONTENT_ENCODING, CONTENT_LENGTH, CONTENT_TYPE, HOST,
9        RETRY_AFTER, TRANSFER_ENCODING,
10    },
11    Request, Uri,
12};
13use hyper_util::rt::TokioIo;
14use serde::{Deserialize, Serialize};
15use std::time::Duration;
16pub const IO_TIMEOUT: Duration = Duration::from_secs(10);
17pub const MAX_RESPONSE_BYTES: usize = 64 * 1024;
18#[derive(Clone)]
19pub struct Origin {
20    pub authority: String,
21    pub connect_addr: String,
22}
23impl Origin {
24    pub fn parse(value: &str) -> Result<Self> {
25        ensure!(
26            !value.contains(['#', '@']),
27            "HTTP origin cannot contain credentials/fragment"
28        );
29        let value = if value.starts_with("http://") {
30            value.to_string()
31        } else {
32            format!("http://{value}")
33        };
34        let uri: Uri = value.parse().context("invalid HTTP origin")?;
35        ensure!(
36            uri.scheme_str() == Some("http")
37                && uri.query().is_none()
38                && matches!(uri.path(), "" | "/"),
39            "cleartext HTTP origin only, without path/query"
40        );
41        let authority = uri.authority().context("HTTP host required")?;
42        ensure!(authority.as_str().len() <= 1024, "HTTP origin byte limit");
43        let host = authority.host();
44        let connect_addr = format!("{}:{}", host, authority.port_u16().unwrap_or(8086));
45        Ok(Self {
46            authority: authority.as_str().to_string(),
47            connect_addr,
48        })
49    }
50}
51pub(crate) fn query(value: &str) -> String {
52    let mut out = String::new();
53    for b in value.bytes() {
54        if b.is_ascii_alphanumeric() || b"-._~".contains(&b) {
55            out.push(b as char);
56        } else {
57            use std::fmt::Write;
58            let _ = write!(out, "%{b:02X}");
59        }
60    }
61    out
62}
63#[derive(Clone, Debug, Serialize, Deserialize)]
64pub struct WriteError {
65    pub code: String,
66    pub message: String,
67    #[serde(skip_serializing_if = "Option::is_none")]
68    pub line: Option<usize>,
69    #[serde(skip_serializing_if = "Option::is_none")]
70    pub accepted_points: Option<usize>,
71    #[serde(skip_serializing_if = "Option::is_none")]
72    pub rejected_points: Option<usize>,
73}
74#[derive(Clone, Debug, Serialize)]
75pub struct WriteResponse {
76    pub org: String,
77    pub bucket: String,
78    pub point_count: usize,
79    pub status: u16,
80    pub error: Option<WriteError>,
81    pub retry_after_seconds: Option<u16>,
82}
83pub async fn write(
84    origin: Origin,
85    token: Option<String>,
86    batch: WriteBatch,
87    body: Vec<u8>,
88) -> Result<WriteResponse> {
89    tokio::time::timeout(IO_TIMEOUT, async move {
90        let socket = tokio::net::TcpStream::connect(&origin.connect_addr).await?;
91        let (mut sender, connection) = http1::Builder::new()
92            .max_headers(64)
93            .max_buf_size(32 * 1024)
94            .handshake(TokioIo::new(socket))
95            .await?;
96        let mut request = Request::builder()
97            .method("POST")
98            .uri(format!(
99                "/api/v2/write?org={}&bucket={}&precision={}",
100                query(&batch.org),
101                query(&batch.bucket),
102                batch.precision.as_str()
103            ))
104            .header(HOST, &origin.authority)
105            .header(CONNECTION, "close")
106            .header(CONTENT_TYPE, "text/plain; charset=utf-8")
107            .header(
108                CONTENT_ENCODING,
109                if batch.gzip { "gzip" } else { "identity" },
110            );
111        if let Some(token) = token {
112            request = request.header(AUTHORIZATION, format!("Token {token}"));
113        }
114        let request = request.body(Full::new(Bytes::from(body)))?;
115        let exchange = async {
116            let response = sender.send_request(request).await?;
117            let (parts, body) = response.into_parts();
118            ensure!(
119                parts.headers.get_all(RETRY_AFTER).iter().count() <= 1,
120                "duplicate Retry-After"
121            );
122            let retry = parts
123                .headers
124                .get(RETRY_AFTER)
125                .map(|h| {
126                    h.to_str()
127                        .context("invalid Retry-After")?
128                        .parse::<u16>()
129                        .context("only numeric Retry-After supported")
130                })
131                .transpose()?;
132            if let Some(n) = retry {
133                ensure!(
134                    (1..=3600).contains(&n) && matches!(parts.status.as_u16(), 429 | 503),
135                    "invalid Retry-After range/status"
136                );
137            }
138            ensure!(
139                parts.headers.get_all(CONTENT_ENCODING).iter().count() <= 1,
140                "duplicate response encoding"
141            );
142            ensure!(
143                parts
144                    .headers
145                    .get(CONTENT_ENCODING)
146                    .map(|v| v.as_bytes() == b"identity")
147                    .unwrap_or(true),
148                "compressed HTTP responses unsupported"
149            );
150            let bytes = Limited::new(body, MAX_RESPONSE_BYTES)
151                .collect()
152                .await
153                .map_err(|_| anyhow::anyhow!("invalid/oversized response body"))?
154                .to_bytes();
155            let status = parts.status.as_u16();
156            let error = if status == 204 {
157                ensure!(
158                    bytes.is_empty()
159                        && retry.is_none()
160                        && parts
161                            .headers
162                            .get_all(CONTENT_LENGTH)
163                            .iter()
164                            .all(|h| h.as_bytes() == b"0")
165                        && !parts.headers.contains_key(TRANSFER_ENCODING),
166                    "204 response must be empty"
167                );
168                None
169            } else {
170                ensure!(
171                    [400, 401, 403, 404, 405, 408, 413, 415, 422, 429, 500, 503].contains(&status),
172                    "unsupported InfluxDB write status"
173                );
174                let ct = parts
175                    .headers
176                    .get(CONTENT_TYPE)
177                    .and_then(|v| v.to_str().ok())
178                    .context("JSON error content type required")?;
179                ensure!(
180                    ct.split(';')
181                        .next()
182                        .is_some_and(|v| v.trim().eq_ignore_ascii_case("application/json")),
183                    "JSON error content type required"
184                );
185                let error: WriteError =
186                    serde_json::from_slice(&bytes).context("invalid InfluxDB JSON error")?;
187                ensure!(
188                    !error.code.is_empty()
189                        && error.code.len() <= 128
190                        && !crate::utils::sanitize::has_controls(&error.code)
191                        && error.message.len() <= 1024
192                        && !crate::utils::sanitize::has_controls(&error.message),
193                    "error text limit"
194                );
195                if let Some(line) = error.line {
196                    ensure!((1..=codec::MAX_LINES).contains(&line), "invalid error line");
197                }
198                for n in [error.accepted_points, error.rejected_points]
199                    .into_iter()
200                    .flatten()
201                {
202                    ensure!(n <= batch.points.len(), "invalid partial-write count");
203                }
204                if let (Some(a), Some(r)) = (error.accepted_points, error.rejected_points) {
205                    ensure!(
206                        a + r == batch.points.len(),
207                        "inconsistent partial-write counts"
208                    );
209                }
210                Some(error)
211            };
212            Ok::<_, anyhow::Error>(WriteResponse {
213                org: batch.org,
214                bucket: batch.bucket,
215                point_count: batch.points.len(),
216                status,
217                error,
218                retry_after_seconds: retry,
219            })
220        };
221        tokio::pin!(connection);
222        tokio::pin!(exchange);
223        // Both futures own the single TCP exchange; dropping either path cancels IO.
224        tokio::select! {
225            result = &mut exchange => result,
226            result = &mut connection => {
227                result.context("HTTP driver failed")?;
228                exchange.await
229            }
230        }
231    })
232    .await
233    .context("InfluxDB exchange deadline exceeded")?
234}
netget/client/influxdb/transport.rs

netget/client/influxdb/
transport.rs
