beeper 0.1.4

#![allow(unused_imports)]
use crate::{
    Dfa, Error, MatchId, MessageBuffer,
    dfa::{ANY_STATE, INIT_STATE, fmt_input},
    http1::action::Action,
    pseudo_header::{METHOD, PATH, STATUS},
};
use http::HeaderName;
use std::{collections::HashMap, mem::MaybeUninit};
use tracing::{Level, debug, trace, warn};
use types::*;
use xbpf::libbpf::{
    self as libbpf_rs, Link, MapCore, OpenObject,
    skel::{OpenSkel, Skel, SkelBuilder},
};

const CR: &str = "\r";
const LF: &str = "\n";

/// The number of ranges a parser can be configured to capture. Must stay in
/// sync with `MAX_MATCHES` of beeper/beeper.h.
const MAX_MATCHES: u8 = 32;

/// A parser for HTTP/1.1 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<Action>,

    /// The number of matches occuring in the patterns.
    num_matches: u8,

    /// The parse function name for each message buffer.
    parse_fns: HashMap<MessageBuffer, String>,

    /// The matched function name. It is agnostic to the message buffer.
    matched_fn: Option<String>,

    /// The extract function name for each message buffer.
    extract_fns: HashMap<MessageBuffer, String>,

    /// The match id of every header captured so far.
    captures: HashMap<String, MatchId>,
}

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

#[allow(dead_code)]
impl Parser {
    /// Creates a new HTTP/1.1 parser.
    pub fn new() -> Parser {
        Parser {
            dfa: Dfa::new(),
            num_matches: 0,
            parse_fns: HashMap::new(),
            matched_fn: None,
            extract_fns: HashMap::new(),
            captures: HashMap::new(),
        }
    }

    /// Specifies the name of the stub function defined with `BEEPER_HTTP1_PARSE_*`
    /// in the target program. When calling [`Parser::attach`], Beeper will insert its logic
    /// into the given template using the
    /// [BPF_PROG_TYPE_EXT](https://docs.ebpf.io/linux/program-type/BPF_PROG_TYPE_EXT/)
    /// program type.
    ///
    /// # Examples
    ///
    /// ```
    /// use beeper::{MessageBuffer, http1};
    ///
    /// let parser = http1::Parser::new().parse_fn("parse_http1", MessageBuffer::Msg);
    /// ```
    pub fn parse_fn<S: ToString>(mut self, parse_fn: S, msg_buf: MessageBuffer) -> Parser {
        self.parse_fns.insert(msg_buf, parse_fn.to_string());
        self
    }

    /// Specifies the name of the stub function defined with `BEEPER_MATCHED`
    /// in the target program. When calling [`Parser::attach`], Beeper will insert its logic
    /// into the given template using the
    /// [BPF_PROG_TYPE_EXT](https://docs.ebpf.io/linux/program-type/BPF_PROG_TYPE_EXT/)
    /// program type.
    ///
    /// # Examples
    ///
    /// ```
    /// use beeper::http1;
    ///
    /// let parser = http1::Parser::new().matched_fn("http1_matched");
    /// ```
    pub fn matched_fn<S: ToString>(mut self, matched_fn: S) -> Parser {
        self.matched_fn = Some(matched_fn.to_string());
        self
    }

    /// Specifies the name of the stub function defined with `BEEPER_EXTRACT_MATCH_*`
    /// in the target program. When calling [`Parser::attach`], Beeper will insert its logic
    /// into the given template using the
    /// [BPF_PROG_TYPE_EXT](https://docs.ebpf.io/linux/program-type/BPF_PROG_TYPE_EXT/)
    /// program type.
    ///
    /// # Examples
    ///
    /// ```
    /// use beeper::{MessageBuffer, http1};
    ///
    /// let parser = http1::Parser::new().extract_fn("extract_http1_match", MessageBuffer::Msg);
    /// ```
    pub fn extract_fn<S: ToString>(mut self, extract_fn: S, msg_buf: MessageBuffer) -> Parser {
        self.extract_fns.insert(msg_buf, extract_fn.to_string());
        self
    }

    /// Returns an unused match id.
    ///
    /// # Errors
    ///
    /// Returns an error if the parser is already configured with
    /// [`MAX_MATCHES`] matches.
    fn new_match(&mut self) -> Result<MatchId, Error> {
        if self.num_matches >= MAX_MATCHES {
            return Err(Error::MatchLimitExceeded(MAX_MATCHES as usize));
        }

        let id = MatchId(self.num_matches);
        self.num_matches += 1;
        Ok(id)
    }

    /// 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. This function also accepts pseudo header fields (see
    /// [`crate::PseudoHeaderName`]).
    ///
    /// # Errors
    ///
    /// Returns an error if the parser already captures the maximum number of
    /// fields.
    ///
    /// # Examples
    ///
    /// ```
    /// # fn main() -> Result<(), beeper::Error> {
    /// use beeper::{http1, pseudo_header::PATH};
    /// use http::header::USER_AGENT;
    ///
    /// let mut parser = http1::Parser::new();
    /// let path = parser.capture_hdr(&PATH)?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn capture_hdr<H: AsRef<str>>(&mut self, name: H) -> Result<MatchId, Error> {
        let name = name.as_ref().to_lowercase();
        if let Some(&mid) = self.captures.get(&name) {
            return Ok(mid);
        }

        let mid = self.capture_new_hdr(&name)?;
        self.captures.insert(name, mid);

        Ok(mid)
    }

    /// Configures the parser to capture the value of a header field it does
    /// not capture yet, see [`Parser::capture_hdr`].
    fn capture_new_hdr(&mut self, name: &str) -> Result<MatchId, Error> {
        if name == METHOD.as_str() || name == PATH.as_str() {
            return self.capture_status_line(name);
        } else if name == STATUS.as_str() {
            return self.capture_status_code();
        }

        let mid = self.new_match()?;
        let mut pattern = self.dfa.start_pattern(ANY_STATE);
        pattern
            .push(LF)
            .push_ci(name)
            .push_optional("\t", true)
            .push_optional(" ", true)
            .push_ci(":")
            .push_optional("\t", true)
            .push_optional(" ", true)
            .with(Action::StartCapture(mid));

        let value = pattern.state();
        pattern
            .push_any(1..)
            .with(Action::EndCapture(mid))
            .push_optional(CR, false)
            .restart_with(LF);

        self.dfa
            .start_pattern(value)
            .push_optional(CR, false)
            .restart_with(LF);

        Ok(mid)
    }

    /// 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,
    /// 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.
    ///
    /// # Errors
    ///
    /// Returns an error if the parser already captures the maximum number of
    /// fields.
    ///
    /// # Examples
    ///
    /// ```
    /// # fn main() -> Result<(), beeper::Error> {
    /// use beeper::http1;
    ///
    /// let mut parser = http1::Parser::new();
    /// let preface = parser.match_http2_preface()?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn match_http2_preface(&mut self) -> Result<MatchId, Error> {
        let mid = self.new_match()?;
        self.dfa
            .start_pattern(INIT_STATE)
            .with(Action::StartCapture(mid))
            .push(&format!(
                "PRI * HTTP/2.0{}{}{}{}SM{}{}{}{}",
                CR, LF, CR, LF, CR, LF, CR, LF
            ))
            .with(Action::EndCaptureAndDone(mid));

        Ok(mid)
    }

    /// Configures the parser to match pseudo headers in the status line of
    /// the request or response.
    ///
    /// # Panics
    ///
    /// Panics if `name` is neither [`METHOD`] nor [`PATH`].
    ///
    /// # Errors
    ///
    /// Returns an error if the parser already captures the maximum number of
    /// fields.
    fn capture_status_line(&mut self, name: &str) -> Result<MatchId, Error> {
        let methods = [
            "POST", "GET", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS", "TRACE",
        ];

        let mid = self.new_match()?;
        if name == METHOD.as_str() {
            self.dfa
                .start_pattern(INIT_STATE)
                .with(Action::StartCapture(mid))
                .push_options_ci(&methods)
                .with(Action::EndCapture(mid))
                .push(" ")
                .push_any(1..)
                .push_ci(" HTTP/1.1")
                .push_optional(CR, false)
                .restart_with(LF);
        } else if name == PATH.as_str() {
            self.dfa
                .start_pattern(INIT_STATE)
                .push_options_ci(&methods)
                .push(" ")
                .with(Action::StartCapture(mid))
                .push_any(1..)
                .with(Action::EndCapture(mid))
                .push_ci(" HTTP/1.1")
                .push_optional(CR, false)
                .restart_with(LF);
        } else {
            // TODO: this method should also accept AUTHORITY, SCHEME
            panic!(
                "capture_status_line called with unsupported header name: {}",
                name
            );
        }

        Ok(mid)
    }

    /// Configures the parser to match the status code of a response.
    ///
    /// # Errors
    ///
    /// Returns an error if the parser already captures the maximum number of
    /// fields.
    fn capture_status_code(&mut self) -> Result<MatchId, Error> {
        let mid = self.new_match()?;
        self.dfa
            .start_pattern(INIT_STATE)
            .push_ci("HTTP/1.1 ")
            .with(Action::StartCapture(mid))
            .push_any(3..=3)
            .with(Action::EndCapture(mid))
            .push_any(1..)
            .push_optional(CR, false)
            .restart_with(LF);

        Ok(mid)
    }

    /// 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(ANY_STATE)
            .push_optional(CR, false)
            .push(LF)
            .push_optional(CR, false)
            .push(LF)
            .with(Action::Done);

        self
    }

    /// Loads the configured parser and attaches it to the target program.
    ///
    /// Every function configured with [`Parser::parse_fn`],
    /// [`Parser::matched_fn`] or [`Parser::extract_fn`] is replaced in the
    /// target program, the remaining parser programs are left
    /// unloaded.
    ///
    /// # 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.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # fn main() -> Result<(), beeper::Error> {
    /// # let prog_fd = 0;
    /// use beeper::{MessageBuffer, http1, pseudo_header::PATH};
    ///
    /// let mut parser = http1::Parser::new();
    /// let path = parser.capture_hdr(&PATH)?;
    ///
    /// let parser = parser
    ///     .parse_fn("parse_http1", MessageBuffer::Msg)
    ///     .extract_fn("extract_http1_match", MessageBuffer::Msg)
    ///     .attach(prog_fd)?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn attach(self, target: i32) -> Result<AttachedParser, Error> {
        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!(target: "bpf", Level::TRACE) {
            open_skel.progs.parse_msg.set_log_level(1);
            open_skel.progs.parse_skb.set_log_level(1);
        }

        // only the programs the parser was configured with are loaded
        for mut prog in open_skel.open_object_mut().progs_mut() {
            prog.set_autoload(false);
        }

        for (msg_buf, func) in &parser.parse_fns {
            let prog = match msg_buf {
                MessageBuffer::Msg => &mut open_skel.progs.parse_msg,
                MessageBuffer::Skb => &mut open_skel.progs.parse_skb,
            };
            prog.set_autoload(true);
            prog.set_attach_target(target, Some(func.clone()))?;
        }

        if let Some(func) = &parser.matched_fn {
            let prog = &mut open_skel.progs.matched;
            prog.set_autoload(true);
            prog.set_attach_target(target, Some(func.clone()))?;
        }

        for (msg_buf, func) in &parser.extract_fns {
            let prog = match msg_buf {
                MessageBuffer::Msg => &mut open_skel.progs.extract_match_msg,
                MessageBuffer::Skb => &mut open_skel.progs.extract_match_skb,
            };
            prog.set_autoload(true);
            prog.set_attach_target(target, Some(func.clone()))?;
        }

        parser.inject(&mut open_skel)?;

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

        let mut links = Vec::new();

        for msg_buf in parser.parse_fns.keys() {
            links.push(match msg_buf {
                MessageBuffer::Msg => skel.progs.parse_msg.attach()?,
                MessageBuffer::Skb => skel.progs.parse_skb.attach()?,
            });
        }

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

        for msg_buf in parser.extract_fns.keys() {
            links.push(match msg_buf {
                MessageBuffer::Msg => skel.progs.extract_match_msg.attach()?,
                MessageBuffer::Skb => skel.progs.extract_match_skb.attach()?,
            });
        }

        debug!("Beeper HTTP/1.1 attached");

        Ok(AttachedParser { links })
    }

    /// Writes the DFA into the read-only data of the parser program.
    ///
    /// # Errors
    ///
    /// Returns an error if the DFA exceeds the eBPF limits on the number
    /// of states, edges, or actions.
    fn inject(&self, skel: &mut OpenParserSkel) -> Result<(), Error> {
        let Some(data) = skel.maps.rodata_data.as_mut() else {
            panic!("the parser program has no read-only data to inject into");
        };

        let num_states = self.dfa.num_states() as usize;
        if num_states > data.s2ts.len() {
            warn!(
                "the patterns take {num_states} states, the parser holds {}",
                data.s2ts.len()
            );
            return Err(Error::ParserExceedsStateLimit);
        }

        // action index 0 is reserved for the noop action
        let mut action_idx = HashMap::new();
        action_idx.insert(None, 0usize);

        for (from, input, to, action) in self.dfa.iter_transitions() {
            let new_action_idx = action_idx.len();
            let action = *action_idx.entry(action).or_insert(new_action_idx);
            if action >= data.a2as.len() {
                warn!(
                    "the patterns take more actions than the {} the parser holds",
                    data.a2as.len()
                );
                return Err(Error::ParserExceedsStateLimit);
            }

            let action = action as u16;
            let input = input as usize;
            if input >= data.s2ts[0].len() {
                warn!("the patterns read inputs the parser has no column for: {input}");
                return Err(Error::ParserExceedsStateLimit);
            }

            trace!(
                "inject; from={} to={} input={} action={}",
                from.0,
                to.0,
                fmt_input(input as u16),
                action
            );

            data.s2ts[from.0 as usize][input] = trans {
                state: to.0,
                action,
            };
        }

        for (action, i) in action_idx {
            let Some(action) = action else { continue };
            data.a2as[i] = action.into();
        }

        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>,
}

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

    fn hdr(i: u8) -> HeaderName {
        HeaderName::from_bytes(format!("x-{i}").as_bytes()).unwrap()
    }

    #[test]
    fn a_parser_captures_at_most_max_matches_ranges() {
        let mut parser = Parser::new();
        for i in 0..MAX_MATCHES {
            let mid = parser.capture_hdr(&hdr(i)).expect("capture header");
            assert_eq!(u8::from(mid), i);
        }

        assert!(matches!(
            parser.capture_hdr(&hdr(MAX_MATCHES)),
            Err(Error::MatchLimitExceeded(limit)) if limit == MAX_MATCHES as usize
        ));
    }

    #[test]
    fn the_same_header_is_captured_under_one_match_id() {
        let names: [&dyn AsRef<str>; 4] = [&hdr(0), &METHOD, &PATH, &STATUS];
        for name in names {
            let name = name.as_ref();
            let mut parser = Parser::new();
            let first = parser.capture_hdr(name).expect("capture header");
            let second = parser.capture_hdr(name).expect("capture header again");

            assert_eq!(
                first, second,
                "capturing {name} twice handed out two ids for one range"
            );
        }
    }
}