use chrono::{DateTime, Utc}; /** This facilitates writing logs to a logtail server. This is a port of [github.com/tailscale/tailscale/logtail](https://github.com/tailscale/tailscale/blob/main/logtail/logtail.go)'s `logtail.go`. */ use reqwest::Client; mod config; pub use self::config::*; /// DefaultHost is the default URL to upload logs to when Builder.base_url isn't provided. pub const DEFAULT_HOST: &'static str = "https://log.tailscale.io"; /** Builds a send/recv pair for the logtail service. Create a new Builder with the [Builder::default] method. The only mandatory field is the `collection`. */ #[derive(Default)] pub struct Builder { collection: Option, private_id: Option, user_agent: Option, base_url: Option, client: Option, buffer_size: usize, } impl Builder { /// The logtail collection to register logs to. This **MUST** be a hostname, even though /// it is not used as a hostname. This is used to disambiguate multiple different programs /// from eachother. pub fn collection(mut self, collection: String) -> Self { self.collection = Some(collection); self } /// The private ID for this logtail identity. If one is not set then one will be auto-generated /// but not saved. Users should store their local [logtail::PrivateID] on disk for later use. pub fn private_id(mut self, id: logtail::PrivateID) -> Self { self.private_id = Some(id); self } /// The user agent to attribute logs to. If not set, one will not be sent. pub fn user_agent(mut self, ua: String) -> Self { self.user_agent = Some(ua); self } /// The base logcatcher URL. If not set, this will default to [DEFAULT_HOST]. pub fn base_url(mut self, base_url: String) -> Self { self.base_url = Some(base_url); self } /// A custom [reqwest::Client] to use for all interactions. If set this makes /// [Builder::user_agent] calls ineffectual. pub fn client(mut self, client: Client) -> Self { self.client = Some(client); self } /// The number of log messages to buffer in memory. By default this is set to /// 256 messages buffered until new ones are dropped. pub fn buffer_size(mut self, buffer_size: usize) -> Self { self.buffer_size = buffer_size; self } /// A "low-memory" friendly value for the buffer size. This will only queue up to /// 64 messages until new ones are dropped. pub fn low_mem(self) -> Self { self.buffer_size(64) } /// Trades the Builder in for an Ingress/Egress pair. Ingress will be safe to `clone` /// as many times as you need to. Egress must have [Egress::post] called periodically /// in order for log messages to get sent to the server. pub fn build(self) -> Result<(Ingress, Egress), Error> { let buf_size: usize = if self.buffer_size != 0 { self.buffer_size } else { 256 }; if let None = self.collection { return Err(Error::NoCollection); } let (tx, rx) = crossbeam::channel::bounded(buf_size); let private_id = self.private_id.unwrap_or(logtail::PrivateID::new()); let base_url = self.base_url.unwrap_or(DEFAULT_HOST.to_string()); let mut u = url::Url::parse(&base_url)?; u.path_segments_mut() .unwrap() .push("c") .push(&self.collection.unwrap()) .push(&private_id.as_hex()); let ing = Ingress { tx }; let eg = Egress { url: u.as_str().to_string(), client: self.client.unwrap_or({ let mut builder = Client::builder(); if let Some(ua) = self.user_agent { builder = builder.user_agent(ua); } builder.build().unwrap() }), rx, }; Ok((ing, eg)) } } #[derive(thiserror::Error, Debug)] pub enum Error { #[error("no collection defined")] NoCollection, #[error("can't put to in-memory buffer: {0}")] TXFail(String), #[error("can't get from in-memory buffer: {0}")] RXFail(#[from] crossbeam::channel::TryRecvError), #[error("can't parse a URL: {0}")] URLParseError(#[from] url::ParseError), #[error("can't post logs: {0}")] ReqwestError(#[from] reqwest::Error), #[error("can't encode to json: {0}")] JsonError(#[from] serde_json::Error), #[error("can't compress")] ZstdError, #[error("must be json object")] MustBeJsonObject, #[error("can't do compression task: {0}")] JoinError(#[from] tokio::task::JoinError), } /// The sink you dump log messages to. You can clone this as many times as you /// need to. #[derive(Clone)] pub struct Ingress { tx: crossbeam::channel::Sender, } impl Ingress { /// Sends a JSON object to the log server. This MUST be a JSON object. pub fn send(&self, val: serde_json::Value) -> Result<(), Error> { if !val.is_object() { return Err(Error::MustBeJsonObject); } let mut val = val.clone(); let header = LogtailHeader { client_time: Utc::now(), }; let obj = val.as_object_mut().unwrap(); obj.insert("logtail".to_string(), serde_json::to_value(header)?); match self.tx.send(val) { Ok(_) => Ok(()), Err(why) => Err(Error::TXFail(format!("{}", why))), } } } #[derive(Clone, serde::Serialize)] struct LogtailHeader { pub client_time: DateTime, } /// The egressor of log messages buffered by its matching [Ingress]. pub struct Egress { url: String, client: reqwest::Client, rx: crossbeam::channel::Receiver, } impl Egress { fn pull(&self) -> Vec { let mut values: Vec = vec![]; loop { match self.rx.try_recv() { Ok(val) => values.push(val), Err(_) => { break; } }; } values } /// Pushes log messages to logtail. This will push everything buffered into the /// log server. This should be called periodically. pub async fn post(&self) -> Result<(), Error> { let values = self.pull(); self.push(values).await?; Ok(()) } async fn push(&self, values: Vec) -> Result<(), Error> { let bytes = serde_json::to_vec(&values)?; let orig_len = bytes.len(); let compressed = tokio::task::spawn_blocking(move || { zstd::block::compress(&bytes, 5).map_err(|_| Error::ZstdError) }) .await??; let resp = self .client .post(&self.url) .header("Content-Encoding", "zstd") .header("Orig-Content-Length", orig_len) .body(compressed) .timeout(std::time::Duration::from_secs(1)) .send() .await?; resp.error_for_status()?; Ok(()) } } #[cfg(test)] mod tests { use super::Builder; #[derive(Clone, serde::Serialize)] struct Data { pub foo: String, } #[tokio::test] async fn end_to_end() { let (ing, mut eg) = Builder::default() .collection("rebterlai.logtail-poster.test".to_string()) .user_agent("rebterlai/test".to_string()) .base_url("http://127.0.0.1:3848".to_string()) .build() .unwrap(); ing.send( serde_json::to_value(Data { foo: "bar".to_string(), }) .unwrap(), ) .unwrap(); eg.post().await.unwrap(); } }