atproto_devtool/common/

diagnostics.rs

//! Shared miette configuration and `NamedSource` helpers.

use std::sync::Arc;

use miette::{GraphicalTheme, MietteHandlerOpts, NamedSource};

// Re-exports so stages can import `LabeledSpan` / `SourceSpan` from a single
// path.
pub use miette::{LabeledSpan, SourceSpan};

/// Install the miette panic hook and graphical report handler.
///
/// Honours `NO_COLOR=1` by dropping to an unstyled theme. Call this exactly once
/// from `main` before any `miette::Result`-returning code runs.
pub fn install_miette_handler(no_color: bool) -> miette::Result<()> {
    // `NO_COLOR` is also respected automatically by miette when set in the
    // environment; passing an explicit theme here covers the `--no-color` flag
    // path without having to touch process-wide env vars (which is `unsafe` in
    // Rust 2024).
    miette::set_hook(Box::new(move |_| {
        let theme = if no_color {
            GraphicalTheme::unicode_nocolor()
        } else {
            GraphicalTheme::unicode()
        };
        Box::new(
            MietteHandlerOpts::new()
                .graphical_theme(theme)
                .context_lines(3)
                .build(),
        )
    }))?;

    // Install miette's panic hook so panics render through the same handler.
    miette::set_panic_hook();

    Ok(())
}

/// Build a `NamedSource` from a name and raw bytes.
///
/// The bytes are cloned into an `Arc<[u8]>` via miette's constructor,
/// so callers may drop the original slice after this returns.
pub fn named_source_from_bytes(name: impl AsRef<str>, bytes: &[u8]) -> NamedSource<Arc<[u8]>> {
    NamedSource::new(name, Arc::<[u8]>::from(bytes))
}

/// Build a `NamedSource` from a name and a UTF-8 string slice.
pub fn named_source_from_str(name: impl AsRef<str>, text: &str) -> NamedSource<String> {
    NamedSource::new(name, text.to_string())
}

/// Pretty-print `body` as JSON for display in a `NamedSource`.
///
/// Real atproto servers routinely emit JSON as a single enormous line, which
/// makes miette's source-span visualization illegible. Any diagnostic that
/// embeds a JSON payload should run it through this helper first so that
/// line-based caret rendering lands somewhere readable.
///
/// If `body` parses as JSON, returns an `Arc<[u8]>` holding the pretty-printed
/// form. If it does not parse, or re-serialization fails, returns an `Arc<[u8]>`
/// containing the original bytes unchanged so callers can still hand something
/// to `NamedSource`.
///
/// Spans computed against the returned bytes must be derived from those same
/// bytes — do not mix a pretty body with a span calculated against the raw
/// body, or vice versa.
pub fn pretty_json_for_display(body: &[u8]) -> Arc<[u8]> {
    match serde_json::from_slice::<serde_json::Value>(body) {
        Ok(value) => match serde_json::to_vec_pretty(&value) {
            Ok(pretty) => Arc::from(pretty),
            Err(_) => Arc::from(body),
        },
        Err(_) => Arc::from(body),
    }
}

/// Convert a 1-based `(line, column)` pair (as produced by `serde_json::Error`)
/// into a `SourceSpan` pointing at that byte inside `body`.
///
/// The returned span has length 1 so miette renders a caret at the exact
/// failure site. `line == 0` is the `serde_json` sentinel for "unknown
/// location" and produces a 1-byte span at the last byte of `body`. If the
/// column runs past the end of the matched line, the span is clamped to the
/// last byte of that line.
pub fn span_at_line_column(body: &[u8], line: usize, column: usize) -> SourceSpan {
    if body.is_empty() {
        return SourceSpan::new(0.into(), 0);
    }
    if line == 0 {
        let end = body.len().saturating_sub(1);
        return SourceSpan::new(end.into(), 1);
    }
    let mut current_line = 1usize;
    let mut line_start = 0usize;
    for (offset, &byte) in body.iter().enumerate() {
        if current_line == line {
            let line_end = body[line_start..]
                .iter()
                .position(|&b| b == b'\n')
                .map(|rel| line_start + rel)
                .unwrap_or(body.len());
            let column_offset = column.saturating_sub(1);
            let span_start = line_start + column_offset;
            if span_start < line_end {
                return SourceSpan::new(span_start.into(), 1);
            } else {
                let len = line_end.saturating_sub(line_start).max(1);
                return SourceSpan::new(line_start.into(), len);
            }
        }
        if byte == b'\n' {
            current_line += 1;
            line_start = offset + 1;
        }
    }
    // Requested line is past the end of the body — clamp to the last byte.
    let end = body.len().saturating_sub(1);
    SourceSpan::new(end.into(), 1)
}

/// Find the span of a JSON quoted literal (key or string value) inside `bytes`.
///
/// Scans for the literal `"<literal>"` pattern and returns the span covering
/// the entire quoted string including the quotes. This works for both JSON
/// keys and string values since both are quoted identically in JSON. Returns
/// `None` if the literal is not present.
///
/// A substring search is acceptable here because the payloads we render are
/// always small (DID documents, single records, query responses) and we only
/// invoke this to highlight an already-extracted key or value — so the risk
/// of a false match inside an unrelated string is negligible in practice.
pub fn span_for_quoted_literal(bytes: &[u8], literal: &str) -> Option<SourceSpan> {
    let search = format!("\"{literal}\"");
    bytes
        .windows(search.len())
        .position(|w| w == search.as_bytes())
        .map(|pos| SourceSpan::new(pos.into(), search.len()))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn pretty_json_for_display_wraps_compact_body() {
        let compact = br#"{"a":1,"b":[2,3]}"#;
        let pretty = pretty_json_for_display(compact);
        let text = std::str::from_utf8(&pretty).unwrap();
        assert!(
            text.contains('\n'),
            "pretty-printed body should be multi-line"
        );
        assert!(text.contains("\"a\""));
    }

    #[test]
    fn pretty_json_for_display_passes_non_json_through() {
        let garbage = b"not valid json <<<";
        let out = pretty_json_for_display(garbage);
        assert_eq!(out.as_ref(), garbage);
    }

    #[test]
    fn span_at_line_column_known_location() {
        let body = b"line1\nline2\nline3";
        let span = span_at_line_column(body, 2, 3);
        // "line2" starts at byte 6, column 3 (1-based) is byte 8 ('n').
        assert_eq!(span.offset(), 8);
        assert_eq!(span.len(), 1);
    }

    #[test]
    fn span_at_line_column_unknown_line_sentinel() {
        let body = b"abc";
        let span = span_at_line_column(body, 0, 0);
        assert_eq!(span.offset(), 2);
        assert_eq!(span.len(), 1);
    }

    #[test]
    fn span_at_line_column_empty_body() {
        let span = span_at_line_column(b"", 1, 1);
        assert_eq!(span.offset(), 0);
        assert_eq!(span.len(), 0);
    }

    #[test]
    fn span_at_line_column_column_past_end_of_line() {
        let body = b"ab\ncd\n";
        let span = span_at_line_column(body, 1, 99);
        // Clamped to the remainder of line 1.
        assert_eq!(span.offset(), 0);
        assert_eq!(span.len(), 2);
    }

    #[test]
    fn span_for_quoted_literal_finds_key() {
        let json = br#"{"service": [], "other": 123}"#;
        let span = span_for_quoted_literal(json, "service").unwrap();
        assert_eq!(
            &json[span.offset()..span.offset() + span.len()],
            b"\"service\""
        );
    }

    #[test]
    fn span_for_quoted_literal_finds_value() {
        let json = br#"{"serviceEndpoint": "https://example.com"}"#;
        let span = span_for_quoted_literal(json, "https://example.com").unwrap();
        assert_eq!(
            &json[span.offset()..span.offset() + span.len()],
            b"\"https://example.com\""
        );
    }

    #[test]
    fn span_for_quoted_literal_missing_returns_none() {
        let json = br#"{"other": 123}"#;
        assert!(span_for_quoted_literal(json, "service").is_none());
    }
}