spftrace 0.3.0

// spftrace – utility for tracing SPF queries
// Copyright © 2022–2023 David Bürgin <dbuergin@gluet.ch>
//
// This program is free software: you can redistribute it and/or modify it under
// the terms of the GNU General Public License as published by the Free Software
// Foundation, either version 3 of the License, or (at your option) any later
// version.
//
// This program is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
// FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
// details.
//
// You should have received a copy of the GNU General Public License along with
// this program. If not, see <https://www.gnu.org/licenses/>.

use crate::{
    config::Config,
    model::{LookupTarget, TimedLookup},
    result::TraceResult,
};
use std::{
    collections::VecDeque,
    fmt::{self, Debug, Formatter},
    mem,
    time::Duration,
};
use viaspf::{
    lookup::Name,
    record::SpfRecord,
    trace::{Trace, TraceEvent, Tracepoint},
    SpfResult,
};

pub struct TimedQueryTrace {
    pub query_trace: QueryTrace,
    pub duration: Duration,
    pub lookup_times: Vec<TimedLookup>,
}

pub struct QueryTrace {
    pub domain: Name,
    pub result: QueryTraceResult,
    pub query_result: SpfResult,
}

impl Debug for QueryTrace {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        f.debug_struct("QueryTrace")
            .field("domain", &self.domain.to_string())
            .field("result", &self.result)
            .field("query_result", &self.query_result.to_string())
            .finish()
    }
}

pub enum QueryTraceResult {
    NoRecord(Vec<Tracepoint>),
    RecordTrace {
        spf_record: SpfRecord,
        result: RecordTraceResult,
    },
}

impl Debug for QueryTraceResult {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match self {
            QueryTraceResult::NoRecord(ts) => {
                f.debug_tuple("NoRecord")
                    .field(&TracepointSlice(&ts[..]))
                    .finish()
            }
            QueryTraceResult::RecordTrace { spf_record, result } => {
                f.debug_struct("RecordTrace")
                    .field("spf_record", &spf_record.to_string())
                    .field("result", result)
                    .finish()
            }
        }
    }
}

pub enum RecordTraceResult {
    NoTerms(Vec<Tracepoint>),
    TermsTrace {
        terms: Vec<TermTrace>,
    }
}

impl Debug for RecordTraceResult {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match self {
            RecordTraceResult::NoTerms(ts) => {
                f.debug_tuple("NoTerms")
                    .field(&TracepointSlice(&ts[..]))
                    .finish()
            }
            RecordTraceResult::TermsTrace { terms } => {
                f.debug_struct("TermsTrace")
                    .field("terms", terms)
                    .finish()
            }
        }
    }
}

pub enum TermTrace {
    Flat(Vec<Tracepoint>),
    Nested {
        term: Vec<Tracepoint>,
        query: QueryTrace,
        after: Vec<Tracepoint>,
    },
}

impl Debug for TermTrace {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match self {
            TermTrace::Flat(ts) => {
                f.debug_tuple("Flat")
                    .field(&TracepointSlice(&ts[..]))
                    .finish()
            }
            TermTrace::Nested { term, query, after } => {
                f.debug_struct("Nested")
                    .field("term", &TracepointSlice(&term[..]))
                    .field("query", query)
                    .field("after", &TracepointSlice(&after[..]))
                    .finish()
            }
        }
    }
}

struct TracepointSlice<'a>(&'a [Tracepoint]);

impl Debug for TracepointSlice<'_> {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        let mut d = f.debug_list();
        for t in self.0 {
            d.entry(&t.to_string());
        }
        d.finish()
    }
}

pub fn build_trace_tree(trace: Trace, config: &Config) -> TraceResult<TimedQueryTrace> {
    // Extract timings based on `TraceEvent` timestamps. After preparation of
    // the trace tree below, timestamps are no longer available.
    let (duration, lookup_times) = if config.time {
        let duration = get_query_duration(&trace);
        let lookup_times = get_lookup_durations(&trace);
        (duration, lookup_times)
    } else {
        Default::default()
    };

    let mut ts = trace.into_iter()
        .map(|e| e.tracepoint)
        .collect::<VecDeque<_>>();

    let query_trace = match ts.pop_front() {
        Some(Tracepoint::ExecuteQuery(domain)) => {
            let query_trace = build_query_trace(&mut ts, domain)?;
            if !ts.is_empty() {
                return Err("extra tracepoints");
            }
            query_trace
        }
        _ => return Err("no query found"),
    };

    Ok(TimedQueryTrace {
        query_trace,
        duration,
        lookup_times,
    })
}

fn get_query_duration(trace: &Trace) -> Duration {
    use Tracepoint::*;

    let start = trace.events.iter().find_map(|e| match e {
        TraceEvent { timestamp, tracepoint: ExecuteQuery(_) } => Some(*timestamp),
        _ => None,
    });
    let end = trace.events.iter().rev().find_map(|e| match e {
        TraceEvent { timestamp, tracepoint: QueryResult(_) } => Some(*timestamp),
        _ => None,
    });

    start.zip(end)
        .and_then(|(start, end)| end.duration_since(start).ok())
        .unwrap_or(Duration::ZERO)
}

fn get_lookup_durations(trace: &Trace) -> Vec<TimedLookup> {
    let mut lookups = vec![];

    // For lookup timings simply compare the lookup timestamp with the one
    // following the lookup.
    for pair in trace.events.windows(2) {
        match pair {
            [e1, e2] => {
                let lookup = match &e1.tracepoint {
                    Tracepoint::LookupA(name) => LookupTarget::A(name.clone()),
                    Tracepoint::LookupAaaa(name) => LookupTarget::Aaaa(name.clone()),
                    Tracepoint::LookupMx(name) => LookupTarget::Mx(name.clone()),
                    Tracepoint::LookupTxt(name) => LookupTarget::Txt(name.clone()),
                    Tracepoint::LookupPtr(ip) => LookupTarget::Ptr(*ip),
                    _ => continue,
                };

                let duration = e2.timestamp
                    .duration_since(e1.timestamp)
                    .unwrap_or(Duration::ZERO);

                lookups.push(TimedLookup { lookup, duration });
            }
            _ => unreachable!(),
        }
    }

    lookups
}

fn build_query_trace(ts: &mut VecDeque<Tracepoint>, domain: Name) -> TraceResult<QueryTrace> {
    use Tracepoint::*;

    fn advance_until_query_result(ts: &mut VecDeque<Tracepoint>) -> TraceResult<SpfResult> {
        while let Some(t) = ts.pop_front() {
            if let QueryResult(spf_result) = t {
                return Ok(spf_result);
            }
        }
        Err("incomplete query")
    }

    while let Some(t) = ts.pop_front() {
        match t {
            SpfRecordLookupError(_)
            | NoSpfRecord
            | MultipleSpfRecords(_)
            | InvalidSpfRecordSyntax(_) => {
                let result = QueryTraceResult::NoRecord(vec![t]);
                let query_result = advance_until_query_result(ts)?;
                return Ok(QueryTrace { domain, result, query_result });
            }
            EvaluateSpfRecord(spf_record) => {
                let result = build_record_trace(ts, spf_record)?;
                let query_result = advance_until_query_result(ts)?;
                return Ok(QueryTrace { domain, result, query_result });
            }
            _ => {}
        }
    }

    Err("incomplete query")
}

fn build_record_trace(
    ts: &mut VecDeque<Tracepoint>,
    spf_record: SpfRecord,
) -> TraceResult<QueryTraceResult> {
    use Tracepoint::*;

    while let Some(t) = ts.pop_front() {
        match t {
            MultipleRedirectModifiers(_) | MultipleExpModifiers(_) => {
                let result = RecordTraceResult::NoTerms(vec![t]);
                return Ok(QueryTraceResult::RecordTrace { spf_record, result });
            }
            NeutralResult => {
                let term = TermTrace::Flat(vec![t]);
                let result = RecordTraceResult::TermsTrace { terms: vec![term] };
                return Ok(QueryTraceResult::RecordTrace { spf_record, result });
            }
            EvaluateDirective(_) => {
                let result = build_terms_trace(ts, vec![t])?;
                return Ok(QueryTraceResult::RecordTrace { spf_record, result });
            }
            EvaluateRedirect(_) => {
                let term = build_redirect_term(ts, vec![t])?;
                let result = RecordTraceResult::TermsTrace { terms: vec![term] };
                return Ok(QueryTraceResult::RecordTrace { spf_record, result });
            }
            _ => {}
        }
    }

    Err("incomplete record")
}

fn build_terms_trace(
    ts: &mut VecDeque<Tracepoint>,
    mut term: Vec<Tracepoint>,
) -> TraceResult<RecordTraceResult> {
    use Tracepoint::*;

    enum State {
        Directive,
        NoMatch,
        Nested { query: QueryTrace, after: Vec<Tracepoint> },
        NestedNoMatch { query: QueryTrace, after: Vec<Tracepoint> },
    }

    let mut state = State::Directive;
    let mut terms = vec![];

    while let Some(t) = ts.pop_front() {
        match state {
            State::Directive => match t {
                DirectiveResult(_) | MechanismErrorResult(..) => {
                    term.push(t);
                    terms.push(TermTrace::Flat(term));
                    return Ok(RecordTraceResult::TermsTrace { terms });
                }
                MechanismNoMatch => {
                    term.push(t);
                    state = State::NoMatch;
                }
                ExecuteQuery(domain) => {
                    let query = build_query_trace(ts, domain)?;
                    state = State::Nested { query, after: vec![] };
                }
                t => term.push(t),
            },
            State::NoMatch => match t {
                EvaluateDirective(_) => {
                    let term = mem::replace(&mut term, vec![t]);
                    terms.push(TermTrace::Flat(term));
                    state = State::Directive;
                }
                EvaluateRedirect(_) => {
                    terms.push(TermTrace::Flat(term));
                    let t = build_redirect_term(ts, vec![t])?;
                    terms.push(t);
                    return Ok(RecordTraceResult::TermsTrace { terms });
                }
                NeutralResult => {
                    terms.push(TermTrace::Flat(term));
                    terms.push(TermTrace::Flat(vec![t]));
                    return Ok(RecordTraceResult::TermsTrace { terms });
                }
                t => term.push(t),
            },
            State::Nested { query, mut after } => match t {
                DirectiveResult(_) => {
                    after.push(t);
                    let t = TermTrace::Nested { term, query, after };
                    terms.push(t);
                    return Ok(RecordTraceResult::TermsTrace { terms });
                }
                MechanismErrorResult(..) => {
                    after.push(t);
                    terms.push(TermTrace::Nested { term, query, after });
                    return Ok(RecordTraceResult::TermsTrace { terms });
                }
                MechanismNoMatch => {
                    after.push(t);
                    state = State::NestedNoMatch { query, after };
                }
                t => {
                    after.push(t);
                    state = State::Nested { query, after };
                }
            },
            State::NestedNoMatch { query, mut after } => match t {
                EvaluateDirective(_) => {
                    let term = mem::replace(&mut term, vec![t]);
                    terms.push(TermTrace::Nested { term, query, after });
                    state = State::Directive;
                }
                EvaluateRedirect(_) => {
                    terms.push(TermTrace::Nested { term, query, after });
                    let t = build_redirect_term(ts, vec![t])?;
                    terms.push(t);
                    return Ok(RecordTraceResult::TermsTrace { terms });
                }
                NeutralResult => {
                    terms.push(TermTrace::Nested { term, query, after });
                    terms.push(TermTrace::Flat(vec![t]));
                    return Ok(RecordTraceResult::TermsTrace { terms });
                }
                t => {
                    after.push(t);
                    state = State::NestedNoMatch { query, after };
                }
            },
        }
    }

    Err("incomplete record")
}

fn build_redirect_term(
    ts: &mut VecDeque<Tracepoint>,
    mut term: Vec<Tracepoint>,
) -> TraceResult<TermTrace> {
    use Tracepoint::*;

    fn collect_until_redirect_result(
        ts: &mut VecDeque<Tracepoint>,
        term: &mut Vec<Tracepoint>,
    ) -> TraceResult<()> {
        while let Some(t) = ts.pop_front() {
            match t {
                RedirectResult(_) => {
                    term.push(t);
                    return Ok(());
                }
                t => term.push(t),
            }
        }
        Err("incomplete redirect")
    }

    while let Some(t) = ts.pop_front() {
        match t {
            RedirectLookupLimitExceeded | InvalidRedirectTargetName => {
                term.push(t);
                collect_until_redirect_result(ts, &mut term)?;
                return Ok(TermTrace::Flat(term));
            }
            ExecuteQuery(domain) => {
                let query = build_query_trace(ts, domain)?;

                let mut after = vec![];
                collect_until_redirect_result(ts, &mut after)?;

                return Ok(TermTrace::Nested { term, query, after });
            }
            t => term.push(t),
        }
    }

    Err("incomplete redirect")
}