active observability templates

src/active_observability_template/io.rs

@@ -0,0 +1,270 @@
+use std::io::{self, Read as _};
+use std::path::Path;
+
+use anyhow::{bail, Context, Result};
+use serde::{de::DeserializeOwned, Serialize};
+
+use crate::http::{build_http_client, DEFAULT_HTTP_TIMEOUT};
+use crate::utils::write_json_atomic;
+
+pub(super) const SCHEMA_VERSION: u32 = 1;
+const MAX_TEMPLATE_BYTES: usize = 10 * 1024 * 1024;
+
+pub(super) async fn read<T: DeserializeOwned>(source: &str) -> Result<T> {
+    let contents = if source == "-" {
+        let mut contents = String::new();
+        io::stdin()
+            .read_to_string(&mut contents)
+            .context("failed to read template from stdin")?;
+        contents
+    } else if is_http_url(source) {
+        fetch_url(source).await?
+    } else {
+        std::fs::read_to_string(source)
+            .with_context(|| format!("failed to read template {}", Path::new(source).display()))?
+    };
+
+    let value: serde_json::Value = serde_json::from_str(&contents).with_context(|| {
+        if source == "-" {
+            "failed to parse template from stdin as JSON".to_string()
+        } else if is_http_url(source) {
+            "failed to parse template URL as JSON; for GitHub Gists, use the Raw URL".to_string()
+        } else {
+            format!(
+                "failed to parse template {} as JSON",
+                Path::new(source).display()
+            )
+        }
+    })?;
+    serde_json::from_value(value).context("template does not match the expected schema")
+}
+
+fn is_http_url(source: &str) -> bool {
+    source.starts_with("http://") || source.starts_with("https://")
+}
+
+async fn fetch_url(source: &str) -> Result<String> {
+    let url = reqwest::Url::parse(source).context("invalid template URL")?;
+    let mut response = build_http_client(DEFAULT_HTTP_TIMEOUT)?
+        .get(url)
+        .send()
+        .await
+        .context("failed to download template URL")?;
+    let status = response.status();
+    if !status.is_success() {
+        bail!("failed to download template URL: HTTP {status}");
+    }
+    if response
+        .content_length()
+        .is_some_and(|length| length > MAX_TEMPLATE_BYTES as u64)
+    {
+        bail!("template URL is larger than the 10 MiB limit");
+    }
+
+    let capacity = response
+        .content_length()
+        .unwrap_or_default()
+        .min(MAX_TEMPLATE_BYTES as u64) as usize;
+    let mut bytes = Vec::with_capacity(capacity);
+    while let Some(chunk) = response
+        .chunk()
+        .await
+        .context("failed to read template URL response")?
+    {
+        if chunk.len() > MAX_TEMPLATE_BYTES - bytes.len() {
+            bail!("template URL is larger than the 10 MiB limit");
+        }
+        bytes.extend_from_slice(&chunk);
+    }
+    String::from_utf8(bytes).context("template URL response is not valid UTF-8")
+}
+
+pub(super) fn validate_version(version: u32) -> Result<()> {
+    if version != SCHEMA_VERSION {
+        bail!(
+            "unsupported template schema version {version}; this version of bt supports schema version {SCHEMA_VERSION}"
+        );
+    }
+    Ok(())
+}
+
+pub(super) fn write<T: Serialize>(value: &T, output: Option<&Path>, force: bool) -> Result<()> {
+    match output {
+        Some(path) if path != Path::new("-") => {
+            if !force
+                && path
+                    .try_exists()
+                    .with_context(|| format!("failed to check whether {} exists", path.display()))?
+            {
+                bail!(
+                    "output file {} already exists; use --force to overwrite it",
+                    path.display()
+                );
+            }
+            write_json_atomic(path, value)
+        }
+        _ => {
+            println!("{}", serde_json::to_string_pretty(value)?);
+            Ok(())
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::io::{Read as _, Write as _};
+    use std::net::TcpListener;
+    use std::path::PathBuf;
+    use std::thread;
+
+    use serde::Deserialize;
+    use tempfile::tempdir;
+
+    use super::*;
+
+    #[derive(Debug, Deserialize, PartialEq)]
+    struct TestTemplate {
+        kind: String,
+        schema_version: u32,
+    }
+
+    #[tokio::test]
+    async fn reads_template_from_file() {
+        let dir = tempdir().expect("tempdir");
+        let path = dir.path().join("template.json");
+        std::fs::write(
+            &path,
+            r#"{"kind":"active_observability_template","schema_version":1}"#,
+        )
+        .expect("write");
+
+        let template: TestTemplate = read(path.to_str().expect("path")).await.expect("template");
+
+        assert_eq!(
+            template,
+            TestTemplate {
+                kind: "active_observability_template".to_string(),
+                schema_version: 1,
+            }
+        );
+    }
+
+    #[tokio::test]
+    async fn reads_template_from_http_url() {
+        let listener = TcpListener::bind(("127.0.0.1", 0)).expect("bind test server");
+        let address = listener.local_addr().expect("test server address");
+        let body = r#"{"kind":"active_observability_template","schema_version":1}"#;
+        let server = thread::spawn(move || {
+            let (mut stream, _) = listener.accept().expect("accept request");
+            let mut request = [0_u8; 1024];
+            let _ = stream.read(&mut request).expect("read request");
+            write!(
+                stream,
+                "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
+                body.len()
+            )
+            .expect("write response");
+        });
+
+        let template: TestTemplate = read(&format!("http://{address}/template.json"))
+            .await
+            .expect("template");
+        server.join().expect("test server");
+
+        assert_eq!(template.kind, "active_observability_template");
+        assert_eq!(template.schema_version, SCHEMA_VERSION);
+    }
+
+    #[tokio::test]
+    async fn stops_streaming_chunked_urls_at_the_size_limit() {
+        let listener = TcpListener::bind(("127.0.0.1", 0)).expect("bind test server");
+        let address = listener.local_addr().expect("test server address");
+        let server = thread::spawn(move || {
+            let (mut stream, _) = listener.accept().expect("accept request");
+            let mut request = [0_u8; 1024];
+            let _ = stream.read(&mut request).expect("read request");
+            stream
+                .write_all(
+                    b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n",
+                )
+                .expect("write response headers");
+
+            let chunk = vec![b'x'; 1024 * 1024];
+            for _ in 0..10 {
+                if write_chunk(&mut stream, &chunk).is_err() {
+                    return;
+                }
+            }
+            let _ = write_chunk(&mut stream, b"x");
+            let _ = stream.write_all(b"0\r\n\r\n");
+        });
+
+        let err = fetch_url(&format!("http://{address}/template.json"))
+            .await
+            .expect_err("oversized chunked response");
+        server.join().expect("test server");
+
+        assert!(err.to_string().contains("larger than the 10 MiB limit"));
+    }
+
+    fn write_chunk(stream: &mut std::net::TcpStream, chunk: &[u8]) -> std::io::Result<()> {
+        write!(stream, "{:X}\r\n", chunk.len())?;
+        stream.write_all(chunk)?;
+        stream.write_all(b"\r\n")
+    }
+
+    #[test]
+    fn validates_schema_version() {
+        validate_version(SCHEMA_VERSION).expect("current version");
+        let err = validate_version(SCHEMA_VERSION + 1).expect_err("future version");
+        assert!(err
+            .to_string()
+            .contains("unsupported template schema version"));
+    }
+
+    #[test]
+    fn recognizes_only_http_urls() {
+        assert!(is_http_url("https://example.invalid/template.json"));
+        assert!(is_http_url("http://example.invalid/template.json"));
+        assert!(!is_http_url("template.json"));
+        assert!(!is_http_url("file:///tmp/template.json"));
+    }
+
+    #[test]
+    fn writes_pretty_json_atomically() {
+        let dir = tempdir().expect("tempdir");
+        let path = PathBuf::from(dir.path()).join("nested/template.json");
+
+        write(
+            &serde_json::json!({"kind": "active_observability_template", "schema_version": 1}),
+            Some(&path),
+            false,
+        )
+        .expect("write");
+
+        assert_eq!(
+            std::fs::read_to_string(path).expect("read"),
+            "{\n  \"kind\": \"active_observability_template\",\n  \"schema_version\": 1\n}\n"
+        );
+    }
+
+    #[test]
+    fn refuses_to_overwrite_without_force() {
+        let dir = tempdir().expect("tempdir");
+        let path = dir.path().join("template.json");
+        std::fs::write(&path, "keep me").expect("write existing file");
+
+        let value = serde_json::json!({
+            "kind": "active_observability_template",
+            "schema_version": 1
+        });
+        let err = write(&value, Some(&path), false).expect_err("overwrite requires force");
+        assert!(err.to_string().contains("--force"));
+        assert_eq!(std::fs::read_to_string(&path).unwrap(), "keep me");
+
+        write(&value, Some(&path), true).expect("forced overwrite");
+        assert!(std::fs::read_to_string(path)
+            .expect("read overwritten file")
+            .contains("active_observability_template"));
+    }
+}