code.rs

#![allow(dead_code)]

use std::collections::{
    hash_map::{Iter, IterMut},
    HashMap,
};

use saluki_error::GenericError;

// A raw span.
//
// Equivalent representation to a span coming over the wire from a tracer.
#[derive(Debug)]
struct RawSpan {
    parent_id: u64,
    id: u64,
    name: String,
    fields: HashMap<String, String>,
}

impl RawSpan {
    fn new<S: Into<String>>(parent_id: u64, id: u64, name: S) -> Self {
        RawSpan {
            parent_id,
            id,
            name: name.into(),
            fields: HashMap::new(),
        }
    }

    fn with_field<K, V>(mut self, key: K, value: V) -> Self
    where
        K: Into<String>,
        V: Into<String>,
    {
        self.fields.insert(key.into(), value.into());
        self
    }

    fn parent_id(&self) -> u64 {
        self.parent_id
    }

    fn id(&self) -> u64 {
        self.id
    }

    fn name(&self) -> &str {
        &self.name
    }

    fn name_mut(&mut self) -> &mut String {
        &mut self.name
    }

    fn fields(&self) -> Iter<'_, String, String> {
        self.fields.iter()
    }

    fn fields_mut(&mut self) -> IterMut<'_, String, String> {
        self.fields.iter_mut()
    }
}

// An obfuscated span.
//
// Newtype wrapper around `Span` that allows us to mark spans as obfuscated.
#[derive(Debug)]
struct ObfuscatedSpan(RawSpan);

impl ObfuscatedSpan {
    fn into_raw(self) -> RawSpan {
        self.0
    }
}

// A normalized span.
//
// Newtype wrapper around `Span` that allows us to mark spans as normalized.
#[derive(Debug)]
struct NormalizedSpan(RawSpan);

impl NormalizedSpan {
    fn into_raw(self) -> RawSpan {
        self.0
    }
}

// A span/trace processor.
trait Processor<Input> {
    type Output;

    fn process(&mut self, input: Input) -> Result<Self::Output, GenericError>;
}

// Processor layer.
trait Layer<P> {
    type Processor;

    fn layer(&self, next: P) -> Self::Processor;
}

// A span obfuscator.
struct SpanObfuscator<P> {
    next: P,
}

impl<P> Processor<RawSpan> for SpanObfuscator<P>
where
    P: Processor<ObfuscatedSpan>,
{
    type Output = P::Output;

    fn process(&mut self, mut input: RawSpan) -> Result<Self::Output, GenericError> {
        // Pretend we're doing some obfuscation here.
        for (_, field_value) in input.fields_mut() {
            // Replace any sensitive information in the span fields.
            if field_value.contains("supersecretvalue") {
                *field_value = field_value.replace("supersecretvalue", "suXXXXXXXXXXXXue");
            }
        }

        // Our span is now obfuscated, so we wrap it up as `ObfuscatedSpan` to indicate that at the type level.
        let obfuscated_span = ObfuscatedSpan(input);

        self.next.process(obfuscated_span)
    }
}

// Span obfuscator layer.
struct SpanObfuscatorLayer;

impl<P> Layer<P> for SpanObfuscatorLayer {
    type Processor = SpanObfuscator<P>;

    fn layer(&self, next: P) -> Self::Processor {
        SpanObfuscator { next }
    }
}

// A span normalizer.
struct SpanNormalizer<P> {
    next: P,
}

impl<P> Processor<ObfuscatedSpan> for SpanNormalizer<P>
where
    P: Processor<NormalizedSpan>,
{
    type Output = P::Output;

    fn process(&mut self, input: ObfuscatedSpan) -> Result<Self::Output, GenericError> {
        let mut raw_span = input.into_raw();

        // Normalize the span name if it's too long or has hyphens.
        let span_name = raw_span.name_mut();
        span_name.truncate(16);

        if span_name.contains("-") {
            *span_name = span_name.replace("-", "_");
        }

        // Our span is now normalized, so we wrap it up as `NormalizedSpan` to indicate that at the type level.
        let normalized_span = NormalizedSpan(raw_span);

        self.next.process(normalized_span)
    }
}

// Span normalizer layer.
struct SpanNormalizerLayer;

impl<P> Layer<P> for SpanNormalizerLayer {
    type Processor = SpanNormalizer<P>;

    fn layer(&self, next: P) -> Self::Processor {
        SpanNormalizer { next }
    }
}

// A no-op processor.
//
// Returns the input unchanged.
struct Identity;

impl<T> Processor<T> for Identity {
    type Output = T;

    fn process(&mut self, input: T) -> Result<Self::Output, GenericError> {
        Ok(input)
    }
}

// No-op processor/layer.
struct IdentityLayer;

impl<P> Layer<P> for IdentityLayer {
    type Processor = P;

    fn layer(&self, next: P) -> Self::Processor {
        next
    }
}

fn build_span_processor_stack() -> impl Processor<RawSpan, Output = NormalizedSpan> {
    StackBuilder::new()
        .layer(SpanObfuscatorLayer)
        .layer(SpanNormalizerLayer)
        .build()
}

struct StackBuilder<L> {
    layer: L,
}

impl StackBuilder<IdentityLayer> {
    fn new() -> Self {
        StackBuilder { layer: IdentityLayer }
    }
}

impl<L> StackBuilder<L> {
    fn layer<T>(self, layer: T) -> StackBuilder<Stack<T, L>> {
        StackBuilder {
            layer: Stack::new(layer, self.layer),
        }
    }

    fn build(self) -> L::Processor
    where
        L: Layer<Identity>,
    {
        self.layer.layer(Identity)
    }
}

// Two layers chained together.
pub struct Stack<Inner, Outer> {
    inner: Inner,
    outer: Outer,
}

impl<Inner, Outer> Stack<Inner, Outer> {
    /// Create a new `Stack`.
    pub const fn new(inner: Inner, outer: Outer) -> Self {
        Stack { inner, outer }
    }
}

impl<P, Inner, Outer> Layer<P> for Stack<Inner, Outer>
where
    Inner: Layer<P>,
    Outer: Layer<Inner::Processor>,
{
    type Processor = Outer::Processor;

    fn layer(&self, processor: P) -> Self::Processor {
        let inner = self.inner.layer(processor);
        self.outer.layer(inner)
    }
}

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

    #[test]
    fn obfuscates() {
        let mut processor_stack = build_span_processor_stack();

        let raw_span = RawSpan::new(1, 2, "db.operation").with_field(
            "db.operation.query",
            "UPDATE user_account SET password='supersecretvalue' WHERE user_id = 1",
        );

        println!("raw span: {:?}", raw_span);

        let processed_span = processor_stack.process(raw_span).unwrap();

        println!("processed span: {:?}", processed_span);
    }

    #[test]
    fn normalizes() {
        let mut processor_stack = build_span_processor_stack();

        let raw_span = RawSpan::new(1, 2, "sort-of-long-span-name");

        println!("raw span: {:?}", raw_span);

        let processed_span = processor_stack.process(raw_span).unwrap();

        println!("processed span: {:?}", processed_span);
    }
}

output.txt

running 2 tests
test destinations::datadog::common::svc_stack::tests::obfuscates ... ok
test destinations::datadog::common::svc_stack::tests::normalizes ... ok

successes:

---- destinations::datadog::common::svc_stack::tests::obfuscates stdout ----
raw span: RawSpan { parent_id: 1, id: 2, name: "db.operation", fields: {"db.operation.query": "UPDATE user_account SET password='supersecretvalue' WHERE user_id = 1"} }
processed span: NormalizedSpan(RawSpan { parent_id: 1, id: 2, name: "db.operation", fields: {"db.operation.query": "UPDATE user_account SET password='suXXXXXXXXXXXXue' WHERE user_id = 1"} })

---- destinations::datadog::common::svc_stack::tests::normalizes stdout ----
raw span: RawSpan { parent_id: 1, id: 2, name: "sort-of-long-span-name", fields: {} }
processed span: NormalizedSpan(RawSpan { parent_id: 1, id: 2, name: "sort_of_long_spa", fields: {} })