1use crate::server::loki::codec::{self, PushBatch};
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;
18pub const MAX_RESPONSE_HEADERS: usize = 64;
19pub const MAX_RESPONSE_HEADER_BYTES: usize = 32 * 1024;
20#[derive(Clone)]
21pub struct Origin {
22 pub authority: String,
23 pub connect_addr: String,
24}
25impl Origin {
26 pub fn parse(value: &str) -> Result<Self> {
27 ensure!(
28 !value.contains(['#', '@']),
29 "HTTP origin cannot contain credentials/fragment"
30 );
31 let value = if value.starts_with("http://") {
32 value.to_string()
33 } else {
34 format!("http://{value}")
35 };
36 let uri: Uri = value.parse().context("invalid HTTP origin")?;
37 ensure!(
38 uri.scheme_str() == Some("http")
39 && uri.query().is_none()
40 && matches!(uri.path(), "" | "/"),
41 "cleartext HTTP origin only, without path/query"
42 );
43 let authority = uri.authority().context("HTTP host required")?;
44 ensure!(authority.as_str().len() <= 1024, "HTTP origin byte limit");
45 let host = authority.host();
46 let connect_addr = format!("{}:{}", host, authority.port_u16().unwrap_or(3100));
47 Ok(Self {
48 authority: authority.as_str().to_string(),
49 connect_addr,
50 })
51 }
52}
53#[derive(Clone, Debug, Serialize, Deserialize)]
54pub struct PushError {
55 pub message: String,
56}
57#[derive(Clone, Debug, Serialize)]
58pub struct PushResponse {
59 pub tenant_id: String,
60 pub encoding: codec::Encoding,
61 pub stream_count: usize,
62 pub entry_count: usize,
63 pub status: u16,
64 pub error: Option<PushError>,
65 pub retry_after_seconds: Option<u16>,
66}
67pub async fn write(
68 origin: Origin,
69 token: Option<String>,
70 batch: PushBatch,
71 body: Vec<u8>,
72) -> Result<PushResponse> {
73 tokio::time::timeout(IO_TIMEOUT, async move {
74 let socket = tokio::net::TcpStream::connect(&origin.connect_addr).await?;
75 let (mut sender, connection) = http1::Builder::new()
76 .max_headers(MAX_RESPONSE_HEADERS)
77 .max_buf_size(MAX_RESPONSE_HEADER_BYTES)
78 .handshake(TokioIo::new(socket))
79 .await?;
80 let mut request = Request::builder()
81 .method("POST")
82 .uri("/loki/api/v1/push")
83 .header(HOST, &origin.authority)
84 .header(CONNECTION, "close")
85 .header(
86 CONTENT_TYPE,
87 if batch.encoding == codec::Encoding::SnappyProtobuf {
88 "application/x-protobuf"
89 } else {
90 "application/json"
91 },
92 );
93 if batch.encoding == codec::Encoding::GzipJson {
94 request = request.header(CONTENT_ENCODING, "gzip");
95 }
96 if let Some(tenant) = &batch.tenant_id {
97 request = request.header("X-Scope-OrgID", tenant);
98 }
99 if let Some(token) = token {
100 request = request.header(AUTHORIZATION, format!("Bearer {token}"));
101 }
102 let request = request.body(Full::new(Bytes::from(body)))?;
103 let exchange = async {
104 let response = sender.send_request(request).await?;
105 let (parts, body) = response.into_parts();
106 ensure!(
107 parts.headers.len() <= MAX_RESPONSE_HEADERS
108 && parts
109 .headers
110 .iter()
111 .map(|(k, v)| k.as_str().len() + v.as_bytes().len() + 4)
112 .sum::<usize>()
113 <= MAX_RESPONSE_HEADER_BYTES,
114 "response header count/byte limit"
115 );
116 ensure!(
117 parts.headers.get_all(RETRY_AFTER).iter().count() <= 1,
118 "duplicate Retry-After"
119 );
120 let retry = parts
121 .headers
122 .get(RETRY_AFTER)
123 .map(|h| {
124 h.to_str()
125 .context("invalid Retry-After")?
126 .parse::<u16>()
127 .context("only numeric Retry-After supported")
128 })
129 .transpose()?;
130 if let Some(n) = retry {
131 ensure!(
132 (1..=3600).contains(&n) && matches!(parts.status.as_u16(), 429 | 503),
133 "invalid Retry-After range/status"
134 );
135 }
136 ensure!(
137 parts.headers.get_all(CONTENT_ENCODING).iter().count() <= 1,
138 "duplicate response encoding"
139 );
140 ensure!(
141 parts
142 .headers
143 .get(CONTENT_ENCODING)
144 .map(|v| v.as_bytes() == b"identity")
145 .unwrap_or(true),
146 "compressed HTTP responses unsupported"
147 );
148 let bytes = Limited::new(body, MAX_RESPONSE_BYTES)
149 .collect()
150 .await
151 .map_err(|_| anyhow::anyhow!("invalid/oversized response body"))?
152 .to_bytes();
153 let status = parts.status.as_u16();
154 let error = if status == 204 {
155 ensure!(
156 bytes.is_empty()
157 && retry.is_none()
158 && parts
159 .headers
160 .get_all(CONTENT_LENGTH)
161 .iter()
162 .all(|h| h.as_bytes() == b"0")
163 && !parts.headers.contains_key(TRANSFER_ENCODING),
164 "204 response must be empty"
165 );
166 None
167 } else {
168 ensure!(
169 [260, 400, 401, 403, 404, 405, 408, 413, 415, 422, 429, 500, 503]
170 .contains(&status),
171 "unsupported Loki push status"
172 );
173 let message = std::str::from_utf8(&bytes).context("Loki error must be UTF-8")?;
174 ensure!(
175 !message.is_empty()
176 && message.len() <= 4096
177 && !message
178 .chars()
179 .any(|c| c.is_control() && !matches!(c, '\t' | '\n' | '\r')),
180 "error text byte/control limit"
181 );
182 let error = PushError {
183 message: message.to_owned(),
184 };
185 Some(error)
186 };
187 Ok::<_, anyhow::Error>(PushResponse {
188 tenant_id: batch.tenant_id.unwrap_or_else(|| "fake".into()),
189 encoding: batch.encoding,
190 stream_count: batch.streams.len(),
191 entry_count: batch.streams.iter().map(|s| s.entries.len()).sum(),
192 status,
193 error,
194 retry_after_seconds: retry,
195 })
196 };
197 tokio::pin!(connection);
198 tokio::pin!(exchange);
199 tokio::select! {
201 result = &mut exchange => result,
202 result = &mut connection => {
203 result.context("HTTP driver failed")?;
204 exchange.await
205 }
206 }
207 })
208 .await
209 .context("Loki exchange deadline exceeded")?
210}