Skip to content

Commit

Permalink
Clean up kdl parsing
Browse files Browse the repository at this point in the history
  • Loading branch information
@jamesmunns
jamesmunns committed Jun 20, 2024
1 parent 2a6b846 commit 92cf86e
Show file tree
Hide file tree
Showing 6 changed files with 311 additions and 164 deletions.
4 changes: 1 addition & 3 deletions source/river/assets/test-config.kdl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@ services {
listeners {
"0.0.0.0:8080"
"0.0.0.0:4443" cert-path="./assets/test.crt" key-path="./assets/test.key"
"0.0.0.0:8443" cert-path="./assets/test.crt" key-path="./assets/test.key"
}

connectors {
Expand All @@ -21,8 +20,7 @@ services {

path-control {
upstream-request {
filter kind="remove-header-key-regex" pattern=".*SECRET.*"
filter kind="remove-header-key-regex" pattern=".*secret.*"
filter kind="remove-header-key-regex" pattern=".*(secret|SECRET).*"
filter kind="upsert-header" key="x-proxy-friend" value="river"
}
upstream-response {
Expand Down
2 changes: 1 addition & 1 deletion source/river/src/config/internal.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ impl Config {
///
/// Note that we use `BTreeMap` and NOT `HashMap`, as we want to maintain the
/// ordering from the configuration file.
#[derive(Debug, Clone, Default)]
#[derive(Debug, Clone, Default, PartialEq)]
pub struct PathControl {
pub(crate) upstream_request_filters: Vec<BTreeMap<String, String>>,
pub(crate) upstream_response_filters: Vec<BTreeMap<String, String>>,
Expand Down
196 changes: 37 additions & 159 deletions source/river/src/config/kdl.rs → source/river/src/config/kdl/mod.rs
View file
Original file line number Diff line number Diff line change
@@ -1,24 +1,24 @@
use std::{
collections::{BTreeMap, HashMap},
net::SocketAddr,
path::PathBuf,
};
use std::{collections::BTreeMap, net::SocketAddr, path::PathBuf};

use kdl::{KdlDocument, KdlEntry, KdlNode};
use miette::{bail, Diagnostic, SourceSpan};
use pingora::upstreams::peer::HttpPeer;

use crate::{
config::internal::{DiscoveryKind, HealthCheckKind, SelectionKind},
config::internal::{
Config, DiscoveryKind, HealthCheckKind, ListenerConfig, ListenerKind, PathControl,
ProxyConfig, SelectionKind, TlsConfig, UpstreamOptions,
},
proxy::request_selector::{
null_selector, source_addr_and_uri_path_selector, uri_path_selector, RequestSelector,
},
};

use super::internal::{
Config, ListenerConfig, ListenerKind, PathControl, ProxyConfig, TlsConfig, UpstreamOptions,
};
#[cfg(test)]
mod test;
mod utils;

/// This is the primary interface for parsing the document.
impl TryFrom<KdlDocument> for Config {
type Error = miette::Error;

Expand All @@ -34,73 +34,10 @@ impl TryFrom<KdlDocument> for Config {
}
}

fn required_child_doc<'a>(
doc: &KdlDocument,
here: &'a KdlDocument,
name: &str,
) -> miette::Result<&'a KdlDocument> {
let node = here
.get(name)
.or_bail(&format!("'{name}' is required!"), doc, here.span())?;

node.children()
.or_bail("expected a nested node", doc, node.span())
}

fn optional_child_doc<'a>(
_doc: &KdlDocument,
here: &'a KdlDocument,
name: &str,
) -> Option<&'a KdlDocument> {
let node = here.get(name)?;

node.children()
}

fn wildcard_argless_child_docs<'a>(
doc: &KdlDocument,
here: &'a KdlDocument,
) -> miette::Result<Vec<(&'a str, &'a KdlDocument)>> {
// TODO: assert no args?
let mut children = vec![];
for node in here.nodes() {
let name = node.name().value();
let child = node.children().or_bail(
&format!("'{name}' should be a nested block"),
doc,
node.span(),
)?;
children.push((name, child));
}
Ok(children)
}

/// Intended to be used with the internal nodes of a section, for example:
///
/// ```kdl
/// listeners {
/// // vvvvvvvvvvvvvv <-------------------------------------- These are the &'str name parts
/// "0.0.0.0:8080" // <\
/// "0.0.0.0:4443" // <----- These are the data nodes
/// "0.0.0.0:8443" cert-path="./assets/test.crt" // </
/// // ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ <-------- These are the KdlEntry parts
/// }
/// ```
fn data_nodes<'a>(
_doc: &KdlDocument,
here: &'a KdlDocument,
) -> miette::Result<Vec<(&'a KdlNode, &'a str, &'a [KdlEntry])>> {
let mut out = vec![];
for node in here.nodes() {
out.push((node, node.name().value(), node.entries()));
}
Ok(out)
}

/// Extract all services from the top level document
fn extract_services(doc: &KdlDocument) -> miette::Result<Vec<ProxyConfig>> {
let service_node = required_child_doc(doc, doc, "services")?;
let services = wildcard_argless_child_docs(doc, service_node)?;
let service_node = utils::required_child_doc(doc, doc, "services")?;
let services = utils::wildcard_argless_child_docs(doc, service_node)?;

let mut proxies = vec![];
for (name, service) in services {
Expand Down Expand Up @@ -133,13 +70,13 @@ fn collect_filters(
doc: &KdlDocument,
node: &KdlDocument,
) -> miette::Result<Vec<BTreeMap<String, String>>> {
let filters = data_nodes(doc, node)?;
let filters = utils::data_nodes(doc, node)?;
let mut fout = vec![];
for (_node, name, args) in filters {
if name != "filter" {
bail!("Invalid Filter Rule");
}
let args = str_str_args(doc, args)?;
let args = utils::str_str_args(doc, args)?;
fout.push(
args.iter()
.map(|(k, v)| (k.to_string(), v.to_string()))
Expand All @@ -157,8 +94,8 @@ fn extract_service(
) -> miette::Result<ProxyConfig> {
// Listeners
//
let listener_node = required_child_doc(doc, node, "listeners")?;
let listeners = data_nodes(doc, listener_node)?;
let listener_node = utils::required_child_doc(doc, node, "listeners")?;
let listeners = utils::data_nodes(doc, listener_node)?;
if listeners.is_empty() {
return Err(Bad::docspan("nonzero listeners required", doc, listener_node.span()).into());
}
Expand All @@ -170,14 +107,14 @@ fn extract_service(

// Connectors
//
let conn_node = required_child_doc(doc, node, "connectors")?;
let conns = data_nodes(doc, conn_node)?;
let conn_node = utils::required_child_doc(doc, node, "connectors")?;
let conns = utils::data_nodes(doc, conn_node)?;
let mut conn_cfgs = vec![];
let mut load_balance: Option<UpstreamOptions> = None;
for (node, name, args) in conns {
if name == "load-balance" {
if load_balance.is_some() {
todo!();
panic!("Don't have two 'load-balance' sections");
}
load_balance = Some(extract_load_balance(doc, node)?);
continue;
Expand All @@ -189,14 +126,14 @@ fn extract_service(
// Path Control (optional)
//
let mut pc = PathControl::default();
if let Some(pc_node) = optional_child_doc(doc, node, "path-control") {
if let Some(pc_node) = utils::optional_child_doc(doc, node, "path-control") {
// upstream-request (optional)
if let Some(ureq_node) = optional_child_doc(doc, pc_node, "upstream-request") {
if let Some(ureq_node) = utils::optional_child_doc(doc, pc_node, "upstream-request") {
pc.upstream_request_filters = collect_filters(doc, ureq_node)?;
}

// upstream-response (optional)
if let Some(uresp_node) = optional_child_doc(doc, pc_node, "upstream-response") {
if let Some(uresp_node) = utils::optional_child_doc(doc, pc_node, "upstream-response") {
pc.upstream_response_filters = collect_filters(doc, uresp_node)?
}
}
Expand All @@ -210,28 +147,9 @@ fn extract_service(
})
}

/// Useful for collecting all arguments as str:str key pairs
fn str_str_args<'a>(
doc: &KdlDocument,
args: &'a [KdlEntry],
) -> miette::Result<Vec<(&'a str, &'a str)>> {
let mut out = vec![];
for arg in args {
let name =
arg.name()
.map(|a| a.value())
.or_bail("arguments should be named", doc, arg.span())?;
let val =
arg.value()
.as_string()
.or_bail("arg values should be a string", doc, arg.span())?;
out.push((name, val));
}
Ok(out)
}

/// Extracts the `load-balance` structure from the `connectors` section
fn extract_load_balance(doc: &KdlDocument, node: &KdlNode) -> miette::Result<UpstreamOptions> {
let items = data_nodes(
let items = utils::data_nodes(
doc,
node.children()
.or_bail("'load-balance' should have children", doc, node.span())?,
Expand All @@ -245,14 +163,19 @@ fn extract_load_balance(doc: &KdlDocument, node: &KdlNode) -> miette::Result<Ups
for (node, name, args) in items {
match name {
"selection" => {
let (sel, args) =
extract_one_str_arg_with_kv_args(doc, node, name, args, |val| match val {
let (sel, args) = utils::extract_one_str_arg_with_kv_args(
doc,
node,
name,
args,
|val| match val {
"RoundRobin" => Some(SelectionKind::RoundRobin),
"Random" => Some(SelectionKind::Random),
"FNV" => Some(SelectionKind::Fnv),
"Ketama" => Some(SelectionKind::Ketama),
_ => None,
})?;
},
)?;
match sel {
SelectionKind::RoundRobin | SelectionKind::Random => {
// No key required, selection is random
Expand Down Expand Up @@ -282,7 +205,7 @@ fn extract_load_balance(doc: &KdlDocument, node: &KdlNode) -> miette::Result<Ups
selection = Some(sel);
}
"health-check" => {
health = Some(extract_one_str_arg(
health = Some(utils::extract_one_str_arg(
doc,
node,
name,
Expand All @@ -294,7 +217,7 @@ fn extract_load_balance(doc: &KdlDocument, node: &KdlNode) -> miette::Result<Ups
)?);
}
"discovery" => {
discover = Some(extract_one_str_arg(
discover = Some(utils::extract_one_str_arg(
doc,
node,
name,
Expand All @@ -320,52 +243,6 @@ fn extract_load_balance(doc: &KdlDocument, node: &KdlNode) -> miette::Result<Ups
})
}

fn extract_one_str_arg<T, F: FnOnce(&str) -> Option<T>>(
doc: &KdlDocument,
node: &KdlNode,
name: &str,
args: &[KdlEntry],
f: F,
) -> miette::Result<T> {
match args {
[one] => one.value().as_string().and_then(f),
_ => None,
}
.or_bail(format!("Incorrect argument for '{name}'"), doc, node.span())
}

fn extract_one_str_arg_with_kv_args<T, F: FnOnce(&str) -> Option<T>>(
doc: &KdlDocument,
node: &KdlNode,
name: &str,
args: &[KdlEntry],
f: F,
) -> miette::Result<(T, HashMap<String, String>)> {
let (first, rest) =
args.split_first()
.or_bail(format!("Missing arguments for '{name}'"), doc, node.span())?;
let first = first.value().as_string().and_then(f).or_bail(
format!("Incorrect argument for '{name}'"),
doc,
node.span(),
)?;
let mut kvs = HashMap::new();
rest.iter().try_for_each(|arg| -> miette::Result<()> {
let key = arg
.name()
.or_bail("Should be a named argument", doc, arg.span())?
.value();
let val = arg
.value()
.as_string()
.or_bail("Should be a string value", doc, arg.span())?;
kvs.insert(key.to_string(), val.to_string());
Ok(())
})?;

Ok((first, kvs))
}

/// Extracts a single connector from the `connectors` section
fn extract_connector(
doc: &KdlDocument,
Expand All @@ -377,7 +254,7 @@ fn extract_connector(
return Err(Bad::docspan("Not a valid socket address", doc, node.span()).into());
};

let args = str_str_args(doc, args)?;
let args = utils::str_str_args(doc, args)?;
let (tls, sni) = match args.as_slice() {
[] => (false, String::new()),
[("tls-sni", sni)] => (true, sni.to_string()),
Expand All @@ -401,7 +278,7 @@ fn extract_listener(
name: &str,
args: &[KdlEntry],
) -> miette::Result<ListenerConfig> {
let mut args = str_str_args(doc, args)?;
let mut args = utils::str_str_args(doc, args)?;
args.sort_by_key(|a| a.0);

// Is this a bindable name?
Expand Down Expand Up @@ -448,7 +325,7 @@ fn extract_listener(
// system { threads-per-service N }
fn extract_threads_per_service(doc: &KdlDocument) -> miette::Result<usize> {
let Some(tps) =
optional_child_doc(doc, doc, "system").and_then(|n| n.get("threads-per-service"))
utils::optional_child_doc(doc, doc, "system").and_then(|n| n.get("threads-per-service"))
else {
// Not present, go ahead and return the default
return Ok(8);
Expand Down Expand Up @@ -516,6 +393,7 @@ impl<T> OptExtParse for Option<T> {
}

impl Bad {
/// Helper function for creating a miette span from a given error
fn docspan(msg: impl Into<String>, doc: &KdlDocument, span: &SourceSpan) -> Self {
Self {
error: msg.into(),
Expand Down
Loading

0 comments on commit 92cf86e

Please sign in to comment.