use super::super::processor::{DataSource, DiagnosticManifest, SourceContext, StreamingDataSource};
use super::{RawResponse, Receive, ReceiveRaw};
use crate::{
client::KibanaClient,
data::{KnownHost, Product},
};
use eyre::{Result, eyre};
use futures::stream::BoxStream;
use reqwest::Method;
use serde::Deserialize;
use serde::de::DeserializeOwned;
use serde_json::Value;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Instant;
use tokio::sync::OnceCell;
use url::Url;
#[derive(Clone)]
pub struct KibanaReceiver {
client: KibanaClient,
url: Url,
version: Arc<OnceCell<semver::Version>>,
spaces: Arc<OnceCell<Vec<String>>>,
}
#[derive(Deserialize)]
struct KibanaStatusVersion {
number: String,
}
#[derive(Deserialize)]
struct KibanaStatusResponse {
version: KibanaStatusVersion,
}
#[derive(Deserialize)]
struct KibanaSpace {
id: String,
}
#[derive(Debug)]
#[non_exhaustive]
pub struct KibanaRequestError {
pub status: reqwest::StatusCode,
pub body: String,
pub response_time_ms: u64,
pub response_size_bytes: u64,
}
impl std::fmt::Display for KibanaRequestError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "http {} - {}", self.status, self.body)
}
}
impl std::error::Error for KibanaRequestError {}
impl KibanaReceiver {
pub fn new(url: Url, client: KibanaClient) -> Self {
Self {
client,
url,
version: Arc::new(OnceCell::new()),
spaces: Arc::new(OnceCell::new()),
}
}
async fn get_status(&self) -> Result<KibanaStatusResponse> {
let started = Instant::now();
let response = self
.client
.request(Method::GET, &HashMap::new(), "/api/status", None)
.await?;
let status = response.status();
let body = response.text().await?;
let response_time_ms = started.elapsed().as_millis() as u64;
let response_size_bytes = body.len() as u64;
if !status.is_success() {
return Err(KibanaRequestError {
status,
body,
response_time_ms,
response_size_bytes,
}
.into());
}
serde_json::from_str(&body).map_err(Into::into)
}
pub async fn get_version(&self) -> Result<&semver::Version> {
self.version
.get_or_try_init(|| async {
let status = self.get_status().await?;
kibana_sync::parse_kibana_version(&status.version.number)
.map_err(|e| eyre!("Failed to parse Kibana version: {}", e))
})
.await
}
pub async fn get_spaces(&self) -> Result<&Vec<String>> {
self.spaces
.get_or_try_init(|| async {
let started = Instant::now();
let response = self
.client
.request(Method::GET, &HashMap::new(), "/api/spaces/space", None)
.await?;
let status = response.status();
let body = response.text().await?;
let response_time_ms = started.elapsed().as_millis() as u64;
let response_size_bytes = body.len() as u64;
if !status.is_success() {
return Err(KibanaRequestError {
status,
body,
response_time_ms,
response_size_bytes,
}
.into());
}
let spaces: Vec<KibanaSpace> = serde_json::from_str(&body)?;
Ok(spaces.into_iter().map(|space| space.id).collect())
})
.await
}
pub async fn get_raw_response_by_path(&self, path: &str, extension: &str) -> Result<RawResponse> {
tracing::debug!("Getting raw Kibana API path: {}", path);
let started = Instant::now();
let mut headers = HashMap::new();
if extension == ".txt" {
headers.insert("Accept".to_string(), "text/plain".to_string());
} else {
headers.insert("Accept".to_string(), "application/json".to_string());
}
let response = self.client.request(Method::GET, &headers, path, None).await?;
let status = response.status();
let body = response.text().await?;
let response_time_ms = started.elapsed().as_millis() as u64;
let response_size_bytes = body.len() as u64;
if !status.is_success() {
return Err(KibanaRequestError {
status,
body,
response_time_ms,
response_size_bytes,
}
.into());
}
Ok(RawResponse {
body,
status: Some(status.as_u16()),
response_time_ms,
response_size_bytes,
})
}
pub async fn get_raw_by_path(&self, path: &str, extension: &str) -> Result<String> {
self.get_raw_response_by_path(path, extension)
.await
.map(|response| response.body)
}
}
impl TryFrom<KnownHost> for KibanaReceiver {
type Error = eyre::Report;
fn try_from(host: KnownHost) -> Result<Self> {
let url = host.get_url()?;
let client = KibanaClient::try_from(host)?;
Ok(Self::new(url, client))
}
}
impl Receive for KibanaReceiver {
async fn collection_date(&self) -> String {
chrono::Utc::now().to_rfc3339()
}
async fn is_connected(&self) -> bool {
self.client.test_connection().await.is_ok()
}
fn filename(&self) -> Option<String> {
None
}
async fn get<T>(&self) -> Result<T>
where
T: DataSource + DeserializeOwned,
{
let ctx = SourceContext::new("kibana", Some(self.get_version().await?.clone()));
let path = T::resolve_source_request_path(&ctx)?;
let started = Instant::now();
let response = self.client.request(Method::GET, &HashMap::new(), &path, None).await?;
let status = response.status();
let body = response.text().await?;
let response_time_ms = started.elapsed().as_millis() as u64;
let response_size_bytes = body.len() as u64;
if status.is_success() {
serde_json::from_str(&body).map_err(Into::into)
} else {
let body_json = serde_json::from_str::<Value>(&body).unwrap_or(Value::String(body));
Err(KibanaRequestError {
status,
body: body_json.to_string(),
response_time_ms,
response_size_bytes,
}
.into())
}
}
async fn get_stream<T>(&self) -> Result<BoxStream<'static, Result<T::Item>>>
where
T: StreamingDataSource + DeserializeOwned,
T::Item: DeserializeOwned + Send + 'static,
{
Err(eyre!("Streaming is not yet implemented for Kibana receiver"))
}
async fn try_get_manifest(&self) -> Result<DiagnosticManifest> {
let status = self.get_status().await?;
Ok(DiagnosticManifest::new(
chrono::Utc::now().to_rfc3339(),
Some(format!("esdiag-{}", env!("CARGO_PKG_VERSION"))),
None,
None,
Some("compatible".to_string()),
Product::Kibana,
Some("kibana_diagnostic".to_string()),
Some("esdiag".to_string()),
Some(status.version.number),
))
}
}
impl ReceiveRaw for KibanaReceiver {
async fn get_raw<T>(&self) -> Result<String>
where
T: DataSource,
{
self.get_raw_response::<T>().await.map(|response| response.body)
}
async fn get_raw_response<T>(&self) -> Result<RawResponse>
where
T: DataSource,
{
let ctx = SourceContext::new("kibana", Some(self.get_version().await?.clone()));
let path = T::resolve_source_request_path(&ctx)?;
let extension = T::resolve_source_extension(&ctx)?;
self.get_raw_response_by_path(&path, &extension).await
}
}
impl std::fmt::Display for KibanaReceiver {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
write!(f, "Kibana {}", self.url)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::data::Auth;
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::{TcpListener, TcpStream},
};
#[tokio::test]
async fn version_parsing_uses_kibana_sync_normalization() {
let (receiver, server) = receiver_with_single_response(
b"HTTP/1.1 200 OK\r\nconnection: close\r\ncontent-type: application/json\r\n\r\n{\"version\":{\"number\":\"v9.6\"}}\n",
)
.await;
let version = receiver.get_version().await.expect("version");
server.await.expect("server");
assert_eq!(version, &semver::Version::new(9, 6, 0));
}
#[tokio::test]
async fn raw_response_errors_preserve_status_body_and_metrics() {
let (receiver, server) = receiver_with_single_response(
b"HTTP/1.1 429 Too Many Requests\r\nconnection: close\r\ncontent-type: text/plain\r\n\r\nslow down",
)
.await;
let err = receiver
.get_raw_response_by_path("/api/saved_objects/_find", ".json")
.await
.expect_err("429 should be an error");
server.await.expect("server");
let request_error = err.downcast_ref::<KibanaRequestError>().expect("request error");
assert_eq!(request_error.status, reqwest::StatusCode::TOO_MANY_REQUESTS);
assert_eq!(request_error.body, "slow down");
assert_eq!(request_error.response_size_bytes, 9);
assert!(request_error.response_time_ms < 60_000);
}
async fn receiver_with_single_response(response: &'static [u8]) -> (KibanaReceiver, tokio::task::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("listener");
let url = Url::parse(&format!("http://{}", listener.local_addr().expect("addr"))).expect("url");
let client = KibanaClient::try_new(url.clone(), Auth::None).expect("client");
let receiver = KibanaReceiver::new(url, client);
let server = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.expect("accept");
read_request(&mut stream).await;
stream.write_all(response).await.expect("write response");
});
(receiver, server)
}
async fn read_request(stream: &mut TcpStream) {
const MAX_TEST_REQUEST_BYTES: usize = 64 * 1024;
let mut request = Vec::new();
let mut buf = [0_u8; 1024];
loop {
let read = stream.read(&mut buf).await.expect("read");
assert_ne!(read, 0, "connection closed before request headers completed");
request.extend_from_slice(&buf[..read]);
assert!(
request.len() <= MAX_TEST_REQUEST_BYTES,
"request headers exceeded test helper limit"
);
if request.windows(4).any(|window| window == b"\r\n\r\n") {
break;
}
}
}
}