beeper 0.1.0

#![allow(unused_imports)]
use crate::{
    autoload_and_attach,
    h1::{dfa::Dfa, Action, CaptureId, MatchId, StateId},
    header::{METHOD, PATH, STATUS},
};
use anyhow::Result;
use http::HeaderName;
use std::{collections::HashMap, mem::MaybeUninit};
use tracing::{debug, trace, warn, Level};
use types::*;
use xbpf::libbpf::{
    self as libbpf_rs,
    skel::{OpenSkel, Skel, SkelBuilder},
    Link, MapCore, OpenObject,
};

/// The sequence terminating the lines of a message.
const CRLF: &str = "\r\n";

/// A parser for HTTP/1.x messages.
///
/// The builder methods configure which fields the parser captures and which
/// functions of the target program it replaces. Nothing is loaded into the
/// kernel until [`Parser::attach`] is called.
pub struct Parser {
    /// The patterns configured so far, compiled into a DFA.
    dfa: Dfa,

    parse_msg_fn: Option<String>,
    parse_buf_fn: Option<String>,
    parse_skb_fn: Option<String>,
    extract_fn: Option<String>,
    matched_fn: Option<String>,
}

xbpf::include_bpf!("h1/parser");

/// Encodes a transition the way the BPF parser reads it out of its transition
/// table.
///
/// An action is a bit field: the flags identifying the action occupy the high
/// bits, the capture and match ids the low ones. [`Action::EndCapture`] needs
/// both ids, so it packs the capture id above the match id.
fn new_transition(state: StateId, action: Option<Action>, rodata: &rodata) -> trans {
    fn start(cid: CaptureId, rodata: &rodata) -> u16 {
        rodata.a_start_capture | (cid.0 as u16) & rodata.a_id_mask
    }
    fn end(cid: CaptureId, mid: MatchId, rodata: &rodata) -> u16 {
        let id = (cid.0 as u16) << 6 | (mid.0 as u16);
        rodata.a_end_capture | id & rodata.a_id_mask
    }

    let action = match action {
        Some(Action::StartCapture(cid)) => start(cid, rodata),
        Some(Action::EndCapture(cid, mid)) => end(cid, mid, rodata),
        Some(Action::Done) => rodata.a_done,
        Some(Action::StartCaptureAndDone(cid)) => start(cid, rodata) | rodata.a_done,
        Some(Action::EndCaptureAndDone(cid, mid)) => end(cid, mid, rodata) | rodata.a_done,
        None => 0,
    };

    trans {
        state: state.0,
        action,
    }
}

#[allow(dead_code)]
impl Parser {
    /// Creates a new HTTP/1.1 parser.
    ///
    /// Additional configuration must be done through the builder methods before calling `attach`.
    pub fn new() -> Parser {
        Parser {
            dfa: Dfa::new(),
            parse_msg_fn: None,
            parse_buf_fn: None,
            parse_skb_fn: None,
            extract_fn: None,
            matched_fn: None,
        }
    }

    /// Specifies the function template in the target program to be replaced with an HTTP/1.1
    /// parser. The function will not be replaced until `attach` is called.
    ///
    /// # Arguments
    ///
    /// * `parse_fn` - The name of the function to replace in the target program
    pub fn replace_parse_msg<S: ToString>(mut self, parse_fn: S) -> Parser {
        self.parse_msg_fn = Some(parse_fn.to_string());
        self
    }

    /// Specifies the function template in the target program to be replaced with a parser
    /// reading from a `sk_buff`. The function will not be replaced until `attach` is called.
    ///
    /// # Arguments
    ///
    /// * `parse_fn` - The name of the function to replace in the target program
    pub fn replace_parse_skb<S: ToString>(mut self, parse_fn: S) -> Parser {
        self.parse_skb_fn = Some(parse_fn.to_string());
        self
    }

    /// Specifies the function template in the target program to be replaced with a parser
    /// reading from a dynptr. The function will not be replaced until `attach` is called.
    ///
    /// # Arguments
    ///
    /// * `parse_fn` - The name of the function to replace in the target program
    pub fn replace_parse_buf<S: ToString>(mut self, parse_fn: S) -> Parser {
        self.parse_buf_fn = Some(parse_fn.to_string());
        self
    }

    /// Specifies the function template in the target program to be called when a pattern match
    /// is completed. The function will not be replaced until `attach` is called.
    ///
    /// # Arguments
    ///
    /// * `matched_fn` - The name of the matched callback function in the target program
    pub fn replace_matched<S: ToString>(mut self, matched_fn: S) -> Parser {
        self.matched_fn = Some(matched_fn.to_string());
        self
    }

    /// Specifies the function template in the target program to be called when extracting
    /// matched content. The function will not be replaced until `attach` is called.
    ///
    /// # Arguments
    ///
    /// * `extract_fn` - The name of the extract callback function in the target program
    pub fn replace_extract<S: ToString>(mut self, extract_fn: S) -> Parser {
        self.extract_fn = Some(extract_fn.to_string());
        self
    }

    /// Configures the parser to capture the value of a header field.
    ///
    /// The field is matched case insensitively and its value is captured up to
    /// the end of the line, without the optional whitespace that may follow the
    /// colon. [`METHOD`], [`PATH`] and [`STATUS`] are not header fields in
    /// HTTP/1.x and are captured from the request or status line instead.
    ///
    /// # Arguments
    ///
    /// * `name` - The header name whose value to capture
    pub fn capture_hdr(mut self, name: &HeaderName) -> Parser {
        if name == &METHOD || name == &PATH {
            return self.capture_status_line_hdr(name);
        } else if name == &STATUS {
            return self.capture_status_code();
        }

        self.dfa
            .start_pattern(false)
            .push(CRLF)
            .push(name.as_str())
            .push_optional("\t")
            .push_optional(" ")
            .push(":")
            .push_optional("\t")
            .push_optional(" ")
            .start_capturing()
            .push_any(1..)
            .end_capturing()
            .restart_with(CRLF);

        self
    }

    /// Configures the parser to match an HTTP/2 preface in an HTTP/1.1 connection.
    ///
    /// This method sets up pattern matching for the HTTP/2 connection preface
    /// (`PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n`), which is used to upgrade from HTTP/1.1 to HTTP/2.
    ///
    /// The preface is captured as a match, so the target program can detect the
    /// upgrade and switch to an HTTP/2 parser for the rest of the connection.
    pub fn match_h2_preface(mut self) -> Parser {
        self.dfa
            .start_pattern(true)
            .start_capturing()
            .push(&format!("PRI * HTTP/2.0{}{}SM{}{}", CRLF, CRLF, CRLF, CRLF))
            .end_capturing()
            .done();

        self
    }

    /// Configures the parser to stop at the empty line that ends the header
    /// block, so that it never walks into the body of a message.
    fn done_on_hdr_end(mut self) -> Parser {
        self.dfa.start_pattern(false).push(CRLF).push(CRLF).done();

        self
    }

    /// Configures the parser to match the request line and capture the field
    /// `name` addresses.
    ///
    /// # Panics
    ///
    /// Panics if `name` is neither [`METHOD`] nor [`PATH`].
    fn capture_status_line_hdr(mut self, name: &HeaderName) -> Parser {
        let methods = [
            "POST", "GET", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS", "TRACE",
        ];

        if name == &METHOD {
            self.dfa
                .start_pattern(true)
                .start_capturing()
                .push_options(&methods)
                .end_capturing()
                .push(" ")
                .push_any(1..)
                .push(" HTTP/1.1")
                .restart_with(CRLF);
        } else if name == &PATH {
            self.dfa
                .start_pattern(true)
                .push_options(&methods)
                .push(" ")
                .start_capturing()
                .push_any(1..)
                .end_capturing()
                .push(" HTTP/1.1")
                .restart_with(CRLF);
        } else {
            panic!(
                "capture_status_line_hdr called with unsupported header name: {}",
                name
            );
        }

        self
    }

    /// Configures the parser to match the status line of a response and capture
    /// its status code.
    fn capture_status_code(mut self) -> Parser {
        self.dfa
            .start_pattern(true)
            .push("HTTP/1.1 ")
            .start_capturing()
            .push_any(3..=3)
            .end_capturing()
            // the reason phrase is matched but not captured
            .push_any(1..)
            .restart_with(CRLF);

        self
    }

    /// Loads the configured parser and attaches it to the target program.
    ///
    /// Every function configured with one of the `replace_*` methods is
    /// replaced in the target program, the remaining parser programs are left
    /// unloaded. The parser always stops at the end of the header block, no
    /// matter which patterns were configured.
    ///
    /// # Arguments
    ///
    /// * `target` - The file descriptor of the target program to attach to
    ///
    /// # Errors
    ///
    /// Returns an error if the parser cannot be loaded, or if one of the
    /// functions it should replace does not exist in the target program with a
    /// matching signature.
    pub fn attach<'obj>(self, target: i32) -> Result<AttachedParser> {
        let parser = self.done_on_hdr_end();

        let skel_builder = ParserSkelBuilder::default();
        let mut open_obj: MaybeUninit<OpenObject> = MaybeUninit::uninit();
        let mut open_skel = skel_builder.open(&mut open_obj)?;
        if tracing::event_enabled!(Level::TRACE) {
            open_skel.progs.parse_msg.set_log_level(1);
            open_skel.progs.parse_buf.set_log_level(1);
            open_skel.progs.parse_buf.set_log_level(1);
        }

        let progs = vec![
            (&mut open_skel.progs.parse_msg, parser.parse_msg_fn.clone()),
            (&mut open_skel.progs.parse_skb, parser.parse_skb_fn.clone()),
            (&mut open_skel.progs.parse_buf, parser.parse_buf_fn.clone()),
            (&mut open_skel.progs.matched, parser.matched_fn.clone()),
            (
                &mut open_skel.progs.extract_match,
                parser.extract_fn.clone(),
            ),
        ];

        for (prog, func) in progs {
            autoload_and_attach(prog, target, func)?;
        }

        parser.inject(&mut open_skel)?;

        let skel = open_skel.load()?;
        xbpf::tracing::try_init(skel.object())?;

        let mut links = Vec::new();
        if parser.parse_msg_fn.is_some() {
            links.push(skel.progs.parse_msg.attach()?);
        }
        if parser.parse_skb_fn.is_some() {
            links.push(skel.progs.parse_skb.attach()?);
        }
        if parser.parse_buf_fn.is_some() {
            links.push(skel.progs.parse_buf.attach()?);
        }

        if parser.matched_fn.is_some() {
            links.push(skel.progs.matched.attach()?);
        }

        if parser.extract_fn.is_some() {
            links.push(skel.progs.extract_match.attach()?);
        }

        debug!("Beeper http/1 attached");

        anyhow::Ok(AttachedParser { links })
    }

    /// Writes the transition table of the DFA into the read-only data of the
    /// parser program. This has to happen before the program is loaded, as the
    /// kernel freezes the section afterwards.
    fn inject(&self, skel: &mut OpenParserSkel) -> Result<()> {
        for (from, to, input, action) in self.dfa.iter_transitions() {
            let s = from.0 as usize;
            let data = skel.maps.rodata_data.as_mut().unwrap();
            let t = new_transition(*to, action, data);
            trace!(
                "inject; from={} to={} input={} action={:?}",
                from.0,
                to.0,
                *input as u8 as char,
                action
            );
            data.s2ts[s][*input as usize] = t;
        }

        Ok(())
    }
}

/// A [`Parser`] attached to a target program.
///
/// It owns the links of the attached programs, so the target program keeps its
/// parser for as long as this value is alive.
pub struct AttachedParser {
    #[allow(dead_code)]
    links: Vec<Link>,
}