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; use std::num::NonZeroUsize; /// DefaultHost is the default URL to upload logs to when Builder.base_url isn't provided. const DEFAULT_HOST: &'static str = "https://log.tailscale.io"; #[derive(Default)] pub struct Builder { collection: String, private_id: Option, user_agent: Option, base_url: Option, client: Option, buffer_size: usize, } impl Builder { pub fn make() -> Self { Builder::default() } pub fn collection(mut self, collection: String) -> Self { self.collection = collection; self } pub fn private_id(mut self, id: logtail::PrivateID) -> Self { self.private_id = Some(id); self } pub fn user_agent(mut self, ua: String) -> Self { self.user_agent = Some(ua); self } pub fn base_url(mut self, base_url: String) -> Self { self.base_url = Some(base_url); self } pub fn client(mut self, client: Client) -> Self { self.client = Some(client); self } pub fn build(self) -> Result<(Ingress, Egress), Error> { let buf_size: usize = if self.buffer_size != 0 { self.buffer_size } else { 256 }; let (tx, rx) = ring_channel::ring_channel(NonZeroUsize::new(buf_size).unwrap()); 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) .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("can't put to in-memory buffer")] TXFail, #[error("can't get from in-memory buffer: {0}")] RXFail(#[from] ring_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, } pub struct Ingress { tx: ring_channel::RingSender, } impl Ingress { pub fn send(&mut 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(_) => Err(Error::TXFail), } } } #[derive(Clone, serde::Serialize)] pub struct LogtailHeader { pub client_time: DateTime, } pub struct Egress { url: String, client: reqwest::Client, rx: ring_channel::RingReceiver, } impl Egress { fn pull(&mut self) -> Result, Error> { let mut values: Vec = vec![]; // self.rx.collect::>().await; loop { match self.rx.try_recv() { Ok(val) => values.push(val), Err(why) => { use ring_channel::TryRecvError::*; match why { Empty => break, Disconnected => return Err(Error::RXFail(why)), }; } }; } Ok(values) } pub async fn post(&mut self) -> Result<(), Error> { let values = self.pull()?; let bytes = serde_json::to_vec(&values)?; let orig_len = bytes.len(); let compressed = zstd::block::compress(&bytes, 5).map_err(|_| Error::ZstdError)?; 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 logpoke() { let (mut ing, mut eg) = Builder::make() .collection("rebterlai.logtail-poster.test".to_string()) .base_url("http://127.0.0.1:48283".to_string()) .build() .unwrap(); ing.send( serde_json::to_value(Data { foo: "bar".to_string(), }) .unwrap(), ) .unwrap(); eg.post().await.unwrap(); } }