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::{
    model::*,
    result::TraceResult,
    trace::{QueryTrace, QueryTraceResult, RecordTraceResult, TermTrace, TimedQueryTrace},
};
use std::collections::VecDeque;
use viaspf::{
    lookup::Name,
    record::{Exists, Include, Ip4, Ip6, Mechanism, Mx, Ptr, Redirect, A},
    trace::Tracepoint,
    ExplanationString,
};

pub fn parse_evaluated_query(trace: TimedQueryTrace) -> TraceResult<TimedEvaluatedQuery> {
    let TimedQueryTrace { query_trace, duration, lookup_times } = trace;

    let query = parse_query(query_trace)?;

    Ok(TimedEvaluatedQuery {
        query,
        duration,
        lookup_times,
    })
}

fn parse_query(trace: QueryTrace) -> TraceResult<EvaluatedQuery> {
    let domain = trace.domain;
    let query_result = trace.query_result;

    let result = match trace.result {
        QueryTraceResult::NoRecord(ts) => {
            // Don’t treat absence of an SPF record as an error. All other
            // causes of `NoRecord` are errors.
            match ts.first() {
                Some(Tracepoint::NoSpfRecord) => EvaluatedQueryResult::NoRecord,
                Some(t) => EvaluatedQueryResult::Error(t.to_string()),
                None => return Err("no tracepoints in query"),
            }
        }
        QueryTraceResult::RecordTrace { spf_record, result } => {
            // `NoTerms` always due to an error to do with multiple modifiers.
            match result {
                RecordTraceResult::NoTerms(_) => {
                    let msg = format!("invalid modifiers in SPF record \"{spf_record}\"");
                    EvaluatedQueryResult::Error(msg)
                }
                RecordTraceResult::TermsTrace { terms } => {
                    let terms = parse_evaluated_terms(terms)?;
                    EvaluatedQueryResult::Record(EvaluatedRecord { spf_record, terms })
                }
            }
        }
    };

    Ok(EvaluatedQuery {
        domain,
        result,
        query_result,
    })
}

fn parse_evaluated_terms(traces: Vec<TermTrace>) -> TraceResult<Vec<EvaluatedTerm>> {
    let mut result = vec![];
    for t in traces {
        let terms = parse_evaluated_term(t)?;
        result.extend(terms);
    }
    Ok(result)
}

fn parse_evaluated_term(trace: TermTrace) -> TraceResult<Vec<EvaluatedTerm>> {
    use Tracepoint::*;

    match trace {
        TermTrace::Flat(ts) => {
            let mut ts = VecDeque::from(ts);
            while let Some(t) = ts.pop_front() {
                match t {
                    EvaluateMechanism(Mechanism::All) => return parse_all(ts),
                    EvaluateMechanism(Mechanism::Ip4(ip4)) => return parse_ip4(ip4, ts),
                    EvaluateMechanism(Mechanism::Ip6(ip6)) => return parse_ip6(ip6, ts),
                    EvaluateMechanism(Mechanism::A(a)) => return parse_a(a, ts),
                    EvaluateMechanism(Mechanism::Mx(mx)) => return parse_mx(mx, ts),
                    EvaluateMechanism(Mechanism::Ptr(ptr)) => return parse_ptr(ptr, ts),
                    EvaluateMechanism(Mechanism::Exists(exists)) => return parse_exists(exists, ts),
                    EvaluateMechanism(Mechanism::Include(include)) => {
                        return parse_include_flat(include, ts);
                    }
                    EvaluateRedirect(redirect) => {
                        return parse_redirect_flat(redirect, ts);
                    }
                    NeutralResult => return Ok(vec![EvaluatedTerm::NeutralDefault]),
                    _ => {}
                }
            }
        }
        TermTrace::Nested { term, query, after } => {
            let mut ts = VecDeque::from(term);
            while let Some(t) = ts.pop_front() {
                match t {
                    EvaluateMechanism(Mechanism::Include(include)) => {
                        return parse_include(include, ts, query, after);
                    }
                    EvaluateRedirect(redirect) => {
                        return parse_redirect(redirect, ts, query, after);
                    }
                    _ => {}
                }
            }
        }
    }

    Err("unexpected tracepoints in term trace")
}

fn parse_all(mut ts: VecDeque<Tracepoint>) -> TraceResult<Vec<EvaluatedTerm>> {
    let (result, exp_modifier) = split_off_directive_result(&mut ts)?;

    let mut terms = vec![EvaluatedTerm::All(AllDirective { result })];

    terms.extend(exp_modifier);

    Ok(terms)
}

fn parse_ip4(mechanism: Ip4, mut ts: VecDeque<Tracepoint>) -> TraceResult<Vec<EvaluatedTerm>> {
    let (result, exp_modifier) = split_off_directive_result(&mut ts)?;

    let mut terms = vec![EvaluatedTerm::Ip4(Ip4Directive { mechanism, result })];

    terms.extend(exp_modifier);

    Ok(terms)
}

fn parse_ip6(mechanism: Ip6, mut ts: VecDeque<Tracepoint>) -> TraceResult<Vec<EvaluatedTerm>> {
    let (result, exp_modifier) = split_off_directive_result(&mut ts)?;

    let mut terms = vec![EvaluatedTerm::Ip6(Ip6Directive { mechanism, result })];

    terms.extend(exp_modifier);

    Ok(terms)
}

fn parse_a(mechanism: A, mut ts: VecDeque<Tracepoint>) -> TraceResult<Vec<EvaluatedTerm>> {
    let lookup_counts = get_lookup_counts(&ts);

    let (result, exp_modifier) = split_off_directive_result(&mut ts)?;

    let mut iter = ts.into_iter();

    let target_domain = find_target_domain(&mut iter);

    let ips = iter
        .filter_map(|t| match t {
            Tracepoint::TryIpAddr(addr) => Some(addr),
            _ => None,
        })
        .collect();

    let mut terms = vec![EvaluatedTerm::A(ADirective {
        mechanism,
        lookup_counts,
        target_domain,
        ips,
        result,
    })];

    terms.extend(exp_modifier);

    Ok(terms)
}

fn parse_mx(mechanism: Mx, mut ts: VecDeque<Tracepoint>) -> TraceResult<Vec<EvaluatedTerm>> {
    let lookup_counts = get_lookup_counts(&ts);

    let (result, exp_modifier) = split_off_directive_result(&mut ts)?;

    let mut iter = ts.into_iter();

    let target_domain = find_target_domain(&mut iter);

    let mxs = if let Some(target_domain) = iter.find_map(|t| match t {
        Tracepoint::TryMxName(d) => Some(d),
        _ => None,
    }) {
        let mut mxs = vec![];
        let mut mx_name = MxName { target_domain, ips: vec![] };

        for t in iter {
            match t {
                Tracepoint::TryMxName(target_domain) => {
                    mxs.push(mx_name);
                    mx_name = MxName { target_domain, ips: vec![] };
                }
                Tracepoint::TryIpAddr(addr) => {
                    mx_name.ips.push(addr);
                }
                _ => {}
            }
        }

        mxs.push(mx_name);

        mxs
    } else {
        vec![]
    };

    let mut terms = vec![EvaluatedTerm::Mx(MxDirective {
        mechanism,
        lookup_counts,
        target_domain,
        mxs,
        result,
    })];

    terms.extend(exp_modifier);

    Ok(terms)
}

fn parse_ptr(mechanism: Ptr, mut ts: VecDeque<Tracepoint>) -> TraceResult<Vec<EvaluatedTerm>> {
    let lookup_counts = get_lookup_counts(&ts);

    let (result, exp_modifier) = split_off_directive_result(&mut ts)?;

    let mut iter = ts.into_iter();

    let target_domain = find_target_domain(&mut iter);

    let mut ts = iter.collect::<VecDeque<_>>();

    // This must be after `TargetName` has been found: otherwise when
    // domain-spec contains %{p} there is a `ValidatePtrName` before!
    let ptrs = split_off_ptr_names(&mut ts)?;

    let mut iter = ts.into_iter();

    let ip = iter.find_map(|t| match t {
        Tracepoint::LookupPtr(ip) => Some(ip),
        _ => None,
    });

    let lookup_error = iter.find_map(|t| match t {
        Tracepoint::ReverseLookupError(e) => Some(e),
        _ => None,
    });

    let mut terms = vec![EvaluatedTerm::Ptr(PtrDirective {
        mechanism,
        lookup_counts,
        target_domain,
        ip,
        lookup_error,
        ptrs,
        result,
    })];

    terms.extend(exp_modifier);

    Ok(terms)
}

fn split_off_ptr_names(ts: &mut VecDeque<Tracepoint>) -> TraceResult<Vec<PtrName>> {
    let ptrs = if let Some(mut ts) = ts.iter()
        .position(|t| matches!(t, Tracepoint::ValidatePtrName(_)))
        .map(|i| ts.split_off(i))
    {
        let mut ptrs = vec![];

        while let Some(tsnext) = ts.iter()
            .skip(1)
            .position(|t| matches!(t, Tracepoint::ValidatePtrName(_)))
            .map(|i| ts.split_off(i + 1))
        {
            let ptr = parse_ptr_name(ts)?;
            ptrs.push(ptr);

            ts = tsnext;
        }

        let ptr = parse_ptr_name(ts)?;
        ptrs.push(ptr);

        ptrs
    } else {
        vec![]
    };

    Ok(ptrs)
}

fn parse_ptr_name(ts: VecDeque<Tracepoint>) -> TraceResult<PtrName> {
    let mut iter = ts.into_iter();

    let target_domain = iter
        .find_map(|t| match t {
            Tracepoint::ValidatePtrName(d) => Some(d),
            _ => None,
        })
        .ok_or("no PTR name found")?;

    let mut error = None;
    let mut ips = vec![];
    let mut validated = false;

    for t in iter {
        match t {
            Tracepoint::PtrAddressLookupLimitExceeded => {
                error = Some("PTR address lookup limit exceeded".into());
            }
            Tracepoint::PtrAddressLookupError(e) => {
                error = Some(e.into());
            }
            Tracepoint::TryIpAddr(addr) => {
                ips.push(addr);
            }
            Tracepoint::PtrNameValidated => {
                validated = true;
            }
            _ => {}
        }
    }

    Ok(PtrName {
        target_domain,
        error,
        ips,
        validated,
    })
}

fn parse_exists(
    mechanism: Exists,
    mut ts: VecDeque<Tracepoint>,
) -> TraceResult<Vec<EvaluatedTerm>> {
    let lookup_counts = get_lookup_counts(&ts);

    let (result, exp_modifier) = split_off_directive_result(&mut ts)?;

    let target_domain = find_target_domain(&mut ts.into_iter());

    let mut terms = vec![EvaluatedTerm::Exists(ExistsDirective {
        mechanism,
        lookup_counts,
        target_domain,
        result,
    })];

    terms.extend(exp_modifier);

    Ok(terms)
}

fn parse_include_flat(
    mechanism: Include,
    mut ts: VecDeque<Tracepoint>,
) -> TraceResult<Vec<EvaluatedTerm>> {
    let lookup_counts = get_lookup_counts(&ts);

    // This is always an error result, no target domain or exp used.
    let (result, _none) = split_off_directive_result(&mut ts)?;

    let include = EvaluatedTerm::Include(IncludeDirective {
        mechanism,
        lookup_counts,
        target_domain: None,
        query: None,
        result,
    });

    Ok(vec![include])
}

fn parse_include(
    mechanism: Include,
    ts: VecDeque<Tracepoint>,
    query: QueryTrace,
    after: Vec<Tracepoint>,
) -> TraceResult<Vec<EvaluatedTerm>> {
    let lookup_counts = get_lookup_counts(&ts);

    let target_domain = find_target_domain(&mut ts.into_iter());

    let query = parse_query(query)?;

    let (result, exp_modifier) = split_off_directive_result(&mut after.into())?;

    let mut terms = vec![EvaluatedTerm::Include(IncludeDirective {
        mechanism,
        lookup_counts,
        target_domain,
        query: Some(query),
        result,
    })];

    terms.extend(exp_modifier);

    Ok(terms)
}

fn parse_redirect_flat(
    modifier: Redirect,
    mut ts: VecDeque<Tracepoint>,
) -> TraceResult<Vec<EvaluatedTerm>> {
    let lookup_counts = get_lookup_counts(&ts);

    let mut error = None;

    let result = loop {
        match ts.pop_front() {
            Some(Tracepoint::RedirectLookupLimitExceeded) => {
                error = Some("lookup limit exceeded".into());
            }
            Some(Tracepoint::InvalidRedirectTargetName) => {
                error = Some("invalid redirect target name".into());
            }
            Some(Tracepoint::RedirectResult(spf_result)) => break spf_result,
            Some(_) => {}
            None => return Err("no redirect result found"),
        }
    };

    let redirect = EvaluatedTerm::Redirect(RedirectModifier {
        modifier,
        lookup_counts,
        error,
        target_domain: None,
        query: None,
        result,
    });

    Ok(vec![redirect])
}

fn parse_redirect(
    modifier: Redirect,
    ts: VecDeque<Tracepoint>,
    query: QueryTrace,
    after: Vec<Tracepoint>,
) -> TraceResult<Vec<EvaluatedTerm>> {
    let lookup_counts = get_lookup_counts(&ts);

    let target_domain = find_target_domain(&mut ts.into_iter());

    let query = parse_query(query)?;

    let mut ts = VecDeque::from(after);
    let mut error = None;

    let result = loop {
        match ts.pop_front() {
            Some(Tracepoint::RedirectNoSpfRecord) => {
                error = Some("no SPF record for redirect".into());
            }
            Some(Tracepoint::RedirectResult(spf_result)) => break spf_result,
            Some(_) => {}
            None => return Err("no redirect result found"),
        }
    };

    let redirect = EvaluatedTerm::Redirect(RedirectModifier {
        modifier,
        lookup_counts,
        target_domain,
        error,
        query: Some(query),
        result,
    });

    Ok(vec![redirect])
}

fn split_off_directive_result(
    ts: &mut VecDeque<Tracepoint>,
) -> TraceResult<(DirectiveResult, Option<EvaluatedTerm>)> {
    let mut ts = ts.iter()
        .position(|t| {
            use Tracepoint::*;
            matches!(t, MechanismErrorResult(..) | MechanismMatch | MechanismNoMatch)
        })
        .map(|i| ts.split_off(i))
        .ok_or("no match result found")?;

    match ts.pop_front().unwrap() {
        Tracepoint::MechanismErrorResult(e, spf_result) => {
            Ok((DirectiveResult::Error(e, spf_result), None))
        }
        Tracepoint::MechanismMatch => {
            let result = loop {
                match ts.pop_back() {
                    Some(Tracepoint::DirectiveResult(spf_result)) => {
                        break DirectiveResult::Match(spf_result);
                    }
                    Some(_) => {}
                    None => return Err("no directive result found"),
                }
            };

            let exp_modifier = parse_exp(ts)?;

            Ok((result, exp_modifier))
        }
        Tracepoint::MechanismNoMatch => Ok((DirectiveResult::NotMatch, None)),
        _ => unreachable!(),
    }
}

fn parse_exp(mut ts: VecDeque<Tracepoint>) -> TraceResult<Option<EvaluatedTerm>> {
    let mut error = None;

    let explanation = loop {
        match ts.pop_back() {
            Some(Tracepoint::ExplanationStringErrorResult(e)) => {
                error = Some(e.into());
                break ExplanationString::Default;
            }
            Some(Tracepoint::ExplanationStringResult(explanation_string)) => {
                break explanation_string;
            }
            Some(_) => {}
            None => return Ok(None),
        }
    };

    let mut iter = ts.into_iter();

    let modifier = iter
        .find_map(|t| match t {
            Tracepoint::EvaluateExplanation(exp) => Some(exp),
            _ => None,
        })
        .ok_or("no exp modifier found")?;

    let target_domain = find_target_domain(&mut iter);

    let mut explain_string = None;

    for t in iter {
        match t {
            Tracepoint::ExplainStringLookupError(e) => {
                error = Some(e.into());
            }
            Tracepoint::MultipleExplainStrings(_)
            | Tracepoint::NoExplainString
            | Tracepoint::InvalidExplainStringSyntax(_) => {
                error = Some(t.to_string().into());
            }
            Tracepoint::ExpandExplainString(es) => {
                explain_string = Some(es);
            }
            _ => {}
        }
    }

    let exp = EvaluatedTerm::Exp(ExpModifier {
        modifier,
        target_domain,
        error,
        explain_string,
        explanation,
    });

    Ok(Some(exp))
}

fn get_lookup_counts(ts: &VecDeque<Tracepoint>) -> LookupCounts {
    let mut lookup_counts = LookupCounts::default();
    for t in ts {
        match t {
            Tracepoint::IncrementLookupCount => lookup_counts.lookups += 1,
            Tracepoint::IncrementPerMechanismLookupCount => lookup_counts.nested += 1,
            Tracepoint::IncrementVoidLookupCount => lookup_counts.void += 1,
            _ => {}
        }
    }
    lookup_counts
}

fn find_target_domain(iter: &mut impl Iterator<Item = Tracepoint>) -> Option<Name> {
    // Note that only after tracepoint `TargetName` are macro-processing
    // tracepoints like `%{p}` out of the way.
    iter.find_map(|t| match t {
        Tracepoint::TargetName(d) => Some(d),
        _ => None,
    })
}