pub(crate) fn extract_test_fn_arg(args: &[String]) -> Option<&str> {
let mut iter = args.iter();
while let Some(a) = iter.next() {
if let Some(val) = a.strip_prefix("--ktstr-test-fn=") {
return Some(val);
}
if a == "--ktstr-test-fn" {
return iter.next().map(|s| s.as_str());
}
}
None
}
pub(crate) fn extract_probe_stack_arg(args: &[String]) -> Option<String> {
for a in args {
if let Some(val) = a.strip_prefix("--ktstr-probe-stack=")
&& !val.is_empty()
{
return Some(val.to_string());
}
}
None
}
pub(crate) fn extract_topo_arg(args: &[String]) -> Option<String> {
for a in args {
if let Some(val) = a.strip_prefix("--ktstr-topo=")
&& !val.is_empty()
{
return Some(val.to_string());
}
}
None
}
pub(crate) fn extract_work_type_arg(args: &[String]) -> Option<String> {
for a in args {
if let Some(val) = a.strip_prefix("--ktstr-work-type=")
&& !val.is_empty()
{
return Some(val.to_string());
}
}
None
}
pub(crate) fn extract_export_test_arg(args: &[String]) -> Option<&str> {
for a in args {
if let Some(val) = a.strip_prefix("--ktstr-export-test=") {
return Some(val);
}
}
None
}
pub(crate) fn extract_export_output_arg(args: &[String]) -> Option<&str> {
for a in args {
if let Some(val) = a.strip_prefix("--ktstr-export-output=")
&& !val.is_empty()
{
return Some(val);
}
}
None
}
pub(crate) const CELL_PARENT_CGROUP_FLAG: &str = "--cell-parent-cgroup";
const _: () = {
let flag = CELL_PARENT_CGROUP_FLAG.as_bytes();
let prefix = b"--cell-parent-cgroup=";
assert!(prefix.len() == flag.len() + 1);
let mut i = 0;
while i < flag.len() {
assert!(prefix[i] == flag[i]);
i += 1;
}
assert!(prefix[flag.len()] == b'=');
};
#[derive(Debug, PartialEq, Eq)]
pub(crate) enum CellParentCgroupArg<'a> {
Absent,
Value(&'a str),
MissingValue,
}
pub(crate) fn cell_parent_path_is_valid(path: &str) -> bool {
if !path.starts_with('/') {
return false;
}
let mut has_normal = false;
for component in std::path::Path::new(path).components() {
use std::path::Component;
match component {
Component::Normal(_) => has_normal = true,
Component::RootDir => {} Component::CurDir | Component::ParentDir => return false,
Component::Prefix(_) => return false, }
}
has_normal
}
pub(crate) fn parse_cell_parent_cgroup<'a>(
args: impl IntoIterator<Item = &'a str>,
) -> CellParentCgroupArg<'a> {
const COMBINED_PREFIX: &str = "--cell-parent-cgroup=";
let mut iter = args.into_iter();
while let Some(a) = iter.next() {
if a == CELL_PARENT_CGROUP_FLAG {
return match iter.next() {
Some(v) => CellParentCgroupArg::Value(v),
None => CellParentCgroupArg::MissingValue,
};
}
if let Some(rest) = a.strip_prefix(COMBINED_PREFIX) {
return CellParentCgroupArg::Value(rest);
}
}
CellParentCgroupArg::Absent
}
pub(crate) fn resolve_cgroup_root(args: &[String]) -> String {
if let Ok(raw) = std::fs::read_to_string("/workload_root_cgroup") {
let trimmed = raw.trim();
if cell_parent_path_is_valid(trimmed) {
return format!("/sys/fs/cgroup{trimmed}");
} else if !trimmed.is_empty() {
eprintln!(
"ktstr_test: ignoring malformed `/workload_root_cgroup` \
value {trimmed:?}; falling back to legacy cgroup-root \
resolution. The host-side gate in `CgroupPath::new` \
normally rejects this at compile time."
);
}
}
let sched_args = std::fs::read_to_string("/sched_args").unwrap_or_default();
if let Some(path) = absolute_cell_parent_value(
parse_cell_parent_cgroup(sched_args.split_whitespace()),
"/sched_args",
) {
return format!("/sys/fs/cgroup{path}");
}
if let Some(path) = absolute_cell_parent_value(
parse_cell_parent_cgroup(args.iter().map(String::as_str)),
"process argv",
) {
return format!("/sys/fs/cgroup{path}");
}
"/sys/fs/cgroup/ktstr".to_string()
}
fn absolute_cell_parent_value<'a>(
parsed: CellParentCgroupArg<'a>,
source: &str,
) -> Option<&'a str> {
match parsed {
CellParentCgroupArg::Value(path) if path.starts_with('/') && path != "/" => Some(path),
CellParentCgroupArg::Value(path) => {
eprintln!(
"ktstr_test: ignoring malformed `--cell-parent-cgroup` value {path:?} \
from {source}; falling back to default cgroup root. The host-side \
gate normally panics on this; reaching this branch means the gate \
was bypassed (hand-edited export script, ad-hoc argv injection).",
);
None
}
CellParentCgroupArg::MissingValue => {
eprintln!(
"ktstr_test: ignoring bare `--cell-parent-cgroup` (no following value) \
from {source}; falling back to default cgroup root.",
);
None
}
CellParentCgroupArg::Absent => None,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn extract_test_fn_arg_equals() {
let args = vec![
"ktstr".into(),
"run".into(),
"--ktstr-test-fn=my_test".into(),
];
assert_eq!(extract_test_fn_arg(&args), Some("my_test"));
}
#[test]
fn extract_test_fn_arg_space() {
let args = vec![
"ktstr".into(),
"run".into(),
"--ktstr-test-fn".into(),
"my_test".into(),
];
assert_eq!(extract_test_fn_arg(&args), Some("my_test"));
}
#[test]
fn extract_test_fn_arg_missing() {
let args = vec!["ktstr".into(), "run".into()];
assert!(extract_test_fn_arg(&args).is_none());
}
#[test]
fn extract_test_fn_arg_trailing() {
let args = vec!["ktstr".into(), "run".into(), "--ktstr-test-fn".into()];
assert!(extract_test_fn_arg(&args).is_none());
}
#[test]
fn extract_test_fn_arg_empty_value() {
let args = vec!["ktstr".into(), "run".into(), "--ktstr-test-fn=".into()];
assert_eq!(extract_test_fn_arg(&args), Some(""));
}
#[test]
fn extract_test_fn_arg_space_form_empty_args() {
let args: Vec<String> = vec![];
assert!(extract_test_fn_arg(&args).is_none());
}
#[test]
fn extract_probe_stack_arg_equals() {
let args = vec![
"ktstr".into(),
"run".into(),
"--ktstr-probe-stack=func_a,func_b".into(),
];
assert_eq!(
extract_probe_stack_arg(&args),
Some("func_a,func_b".to_string())
);
}
#[test]
fn extract_probe_stack_arg_missing() {
let args = vec!["ktstr".into(), "run".into()];
assert!(extract_probe_stack_arg(&args).is_none());
}
#[test]
fn extract_probe_stack_arg_empty_value() {
let args = vec!["ktstr".into(), "--ktstr-probe-stack=".into()];
assert!(extract_probe_stack_arg(&args).is_none());
}
#[test]
fn extract_topo_arg_equals() {
let args = vec!["bin".into(), "--ktstr-topo=1n2l4c2t".into()];
assert_eq!(extract_topo_arg(&args), Some("1n2l4c2t".to_string()));
}
#[test]
fn extract_topo_arg_missing() {
let args = vec!["bin".into(), "--ktstr-test-fn=test".into()];
assert!(extract_topo_arg(&args).is_none());
}
#[test]
fn extract_topo_arg_empty_value() {
let args = vec!["bin".into(), "--ktstr-topo=".into()];
assert!(extract_topo_arg(&args).is_none());
}
#[test]
fn extract_topo_arg_with_other_args() {
let args = vec![
"bin".into(),
"--ktstr-test-fn=my_test".into(),
"--ktstr-topo=1n1l2c1t".into(),
];
assert_eq!(extract_topo_arg(&args), Some("1n1l2c1t".to_string()));
}
#[test]
fn extract_work_type_arg_equals() {
let args = vec!["ktstr".into(), "--ktstr-work-type=SpinWait".into()];
assert_eq!(extract_work_type_arg(&args), Some("SpinWait".to_string()));
}
#[test]
fn extract_work_type_arg_missing() {
let args = vec!["ktstr".into(), "run".into()];
assert!(extract_work_type_arg(&args).is_none());
}
#[test]
fn extract_work_type_arg_empty_value() {
let args = vec!["ktstr".into(), "--ktstr-work-type=".into()];
assert!(extract_work_type_arg(&args).is_none());
}
#[test]
fn extract_export_test_arg_equals() {
let args = vec![
"test_bin".into(),
"--ktstr-export-test=preempt_regression".into(),
];
assert_eq!(extract_export_test_arg(&args), Some("preempt_regression"),);
}
#[test]
fn extract_export_test_arg_missing() {
let args = vec!["test_bin".into(), "--list".into()];
assert!(extract_export_test_arg(&args).is_none());
}
#[test]
fn extract_export_test_arg_empty_value() {
let args = vec!["test_bin".into(), "--ktstr-export-test=".into()];
assert_eq!(extract_export_test_arg(&args), Some(""));
}
#[test]
fn extract_export_output_arg_equals() {
let args = vec![
"test_bin".into(),
"--ktstr-export-output=/tmp/foo.run".into(),
];
assert_eq!(extract_export_output_arg(&args), Some("/tmp/foo.run"),);
}
#[test]
fn extract_export_output_arg_missing() {
let args = vec!["test_bin".into()];
assert!(extract_export_output_arg(&args).is_none());
}
#[test]
fn extract_export_output_arg_empty_value() {
let args = vec!["test_bin".into(), "--ktstr-export-output=".into()];
assert!(extract_export_output_arg(&args).is_none());
}
#[test]
fn parse_cell_parent_cgroup_two_token_form() {
let argv = ["--cell-parent-cgroup", "/user", "--other-flag"];
assert_eq!(
parse_cell_parent_cgroup(argv.iter().copied()),
CellParentCgroupArg::Value("/user")
);
}
#[test]
fn parse_cell_parent_cgroup_combined_form() {
let argv = ["--other-flag", "--cell-parent-cgroup=/user"];
assert_eq!(
parse_cell_parent_cgroup(argv.iter().copied()),
CellParentCgroupArg::Value("/user")
);
}
#[test]
fn parse_cell_parent_cgroup_empty_combined_value() {
let argv = ["--cell-parent-cgroup="];
assert_eq!(
parse_cell_parent_cgroup(argv.iter().copied()),
CellParentCgroupArg::Value("")
);
}
#[test]
fn parse_cell_parent_cgroup_absent() {
let argv = ["--unrelated", "--other-flag=value"];
assert_eq!(
parse_cell_parent_cgroup(argv.iter().copied()),
CellParentCgroupArg::Absent
);
}
#[test]
fn parse_cell_parent_cgroup_two_token_missing_value() {
let argv = ["--cell-parent-cgroup"];
assert_eq!(
parse_cell_parent_cgroup(argv.iter().copied()),
CellParentCgroupArg::MissingValue
);
}
#[test]
fn parse_cell_parent_cgroup_bare_flag_at_end_after_other() {
let argv = ["--other-flag", "--cell-parent-cgroup"];
assert_eq!(
parse_cell_parent_cgroup(argv.iter().copied()),
CellParentCgroupArg::MissingValue
);
}
#[test]
fn parse_cell_parent_cgroup_first_match_wins() {
let argv = [
"--cell-parent-cgroup=/first",
"--cell-parent-cgroup",
"/second",
];
assert_eq!(
parse_cell_parent_cgroup(argv.iter().copied()),
CellParentCgroupArg::Value("/first")
);
}
#[test]
fn parse_cell_parent_cgroup_first_match_wins_two_token_then_combined() {
let argv = [
"--cell-parent-cgroup",
"/first",
"--cell-parent-cgroup=/second",
];
assert_eq!(
parse_cell_parent_cgroup(argv.iter().copied()),
CellParentCgroupArg::Value("/first")
);
}
#[test]
fn parse_cell_parent_cgroup_no_match_on_sibling_long_flag() {
let argv = ["--cell-parent-cgroup-extra=val"];
assert_eq!(
parse_cell_parent_cgroup(argv.iter().copied()),
CellParentCgroupArg::Absent
);
}
#[test]
fn absolute_cell_parent_value_returns_valid_path() {
assert_eq!(
absolute_cell_parent_value(CellParentCgroupArg::Value("/ktstr"), "test"),
Some("/ktstr")
);
}
#[test]
fn absolute_cell_parent_value_rejects_empty() {
assert_eq!(
absolute_cell_parent_value(CellParentCgroupArg::Value(""), "test"),
None
);
}
#[test]
fn absolute_cell_parent_value_rejects_bare_slash() {
assert_eq!(
absolute_cell_parent_value(CellParentCgroupArg::Value("/"), "test"),
None
);
}
#[test]
fn absolute_cell_parent_value_rejects_relative() {
assert_eq!(
absolute_cell_parent_value(CellParentCgroupArg::Value("my_test"), "test"),
None
);
}
#[test]
fn absolute_cell_parent_value_rejects_missing_value() {
assert_eq!(
absolute_cell_parent_value(CellParentCgroupArg::MissingValue, "test"),
None
);
}
#[test]
fn absolute_cell_parent_value_returns_none_on_absent() {
assert_eq!(
absolute_cell_parent_value(CellParentCgroupArg::Absent, "test"),
None
);
}
}