use crate::{common, DataType, Provider};
use proc_macro2::TokenStream;
use quote::{format_ident, quote};
use std::{
collections::BTreeMap,
convert::TryFrom,
io::Write,
process::{Command, Stdio},
};
pub fn compile_provider_source(
source: &str,
config: &crate::CompileProvidersConfig,
) -> Result<TokenStream, crate::Error> {
let dfile = dtrace_parser::File::try_from(source)?;
let header = build_header_from_provider(&source)?;
let provider_info = extract_providers(&header);
let providers = dfile
.providers()
.into_iter()
.map(|provider| {
let provider = Provider::from(provider);
let config = crate::CompileProvidersConfig {
provider: Some(provider.name.clone()),
probe_format: config.probe_format.clone(),
module: match &config.module {
None => Some(provider.name.clone()),
other => other.clone(),
},
};
compile_provider(&provider, &provider_info[&provider.name], &config)
})
.collect::<Vec<_>>();
Ok(quote! {
#(#providers)*
})
}
pub fn compile_provider_from_definition(
provider: &Provider,
config: &crate::CompileProvidersConfig,
) -> TokenStream {
let header = build_header_from_provider(&provider.to_d_source()).unwrap();
let provider_info = extract_providers(&header);
let provider_tokens = compile_provider(provider, &provider_info[&provider.name], config);
quote! {
#provider_tokens
}
}
fn compile_provider(
provider: &Provider,
provider_info: &ProviderInfo,
config: &crate::CompileProvidersConfig,
) -> TokenStream {
let mut probe_impls = Vec::new();
for probe in provider.probes.iter() {
probe_impls.push(compile_probe(
provider,
&probe.name,
config,
provider_info,
&probe.types,
));
}
let module = config.module_ident();
quote! {
pub(crate) mod #module {
#(#probe_impls)*
}
}
}
fn compile_probe(
provider: &Provider,
probe_name: &str,
config: &crate::CompileProvidersConfig,
provider_info: &ProviderInfo,
types: &[DataType],
) -> TokenStream {
let stability = &provider_info.stability;
let stability_fn = format_ident!("stability");
let typedefs = &provider_info.typedefs;
let typedef_fn = format_ident!("typedefs");
let is_enabled = &provider_info.is_enabled[probe_name];
let is_enabled_fn = format_ident!("{}_{}_enabled", &provider.name, probe_name);
let probe = &provider_info.probes[probe_name];
let extern_probe_fn = format_ident!("__{}", config.probe_ident(probe_name));
let ffi_param_list = types.iter().map(|typ| {
let ty = typ.to_rust_ffi_type();
syn::parse2::<syn::FnArg>(quote! { _: #ty }).unwrap()
});
let (unpacked_args, in_regs) = common::construct_probe_args(types);
#[cfg(target_arch = "x86_64")]
let call_instruction = quote! { "call {extern_probe_fn}" };
#[cfg(target_arch = "aarch64")]
let call_instruction = quote! { "bl {extern_probe_fn}" };
#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
compile_error!("USDT only supports x86_64 and AArch64 architectures");
#[cfg(usdt_stable_asm)]
let asm_macro = quote! { std::arch::asm };
#[cfg(not(usdt_stable_asm))]
let asm_macro = quote! { asm };
let impl_block = quote! {
extern "C" {
#[allow(unused)]
#[link_name = #stability]
fn stability();
#[allow(unused)]
#[link_name = #typedefs]
fn typedefs();
#[allow(unused)]
#[link_name = #is_enabled]
fn #is_enabled_fn() -> i32;
#[allow(unused)]
#[link_name = #probe]
fn #extern_probe_fn(#(#ffi_param_list,)*);
}
unsafe {
if #is_enabled_fn() != 0 {
#unpacked_args
#asm_macro!(
".reference {typedefs}",
#call_instruction,
".reference {stability}",
typedefs = sym #typedef_fn,
extern_probe_fn = sym #extern_probe_fn,
stability = sym #stability_fn,
#in_regs
options(nomem, nostack, preserves_flags)
);
}
}
};
common::build_probe_macro(config, provider, probe_name, types, impl_block)
}
#[derive(Debug, Default, Clone)]
struct ProviderInfo {
pub stability: String,
pub typedefs: String,
pub is_enabled: BTreeMap<String, String>,
pub probes: BTreeMap<String, String>,
}
fn extract_providers(header: &str) -> BTreeMap<String, ProviderInfo> {
let mut providers = BTreeMap::new();
for line in header.lines() {
if let Some((provider_name, stability)) = is_stability_line(&line) {
let mut info = ProviderInfo::default();
info.stability = stability.to_string();
providers.insert(provider_name.to_string(), info);
}
if let Some((provider_name, typedefs)) = is_typedefs_line(&line) {
providers.get_mut(provider_name).unwrap().typedefs = typedefs.to_string();
}
if let Some((provider_name, probe_name, enabled)) = is_enabled_line(&line) {
providers
.get_mut(provider_name)
.unwrap()
.is_enabled
.insert(probe_name.to_string(), enabled.to_string());
}
if let Some((provider_name, probe_name, probe)) = is_probe_line(&line) {
providers
.get_mut(provider_name)
.unwrap()
.probes
.insert(probe_name.to_string(), probe.to_string());
}
}
providers
}
fn is_stability_line(line: &str) -> Option<(&str, &str)> {
contains_needle(line, "___dtrace_stability$")
}
fn is_typedefs_line(line: &str) -> Option<(&str, &str)> {
contains_needle(line, "___dtrace_typedefs$")
}
fn contains_needle<'a>(line: &'a str, needle: &str) -> Option<(&'a str, &'a str)> {
if let Some(index) = line.find(needle) {
let rest = &line[index + needle.len()..];
let provider_end = rest.find("$").unwrap();
let provider_name = &rest[..provider_end];
let needle = &line[index + 1..line.len() - 1];
Some((provider_name, needle))
} else {
None
}
}
fn is_enabled_line(line: &str) -> Option<(&str, &str, &str)> {
contains_needle2(line, "extern int __dtrace_isenabled$")
}
fn is_probe_line(line: &str) -> Option<(&str, &str, &str)> {
contains_needle2(line, "extern void __dtrace_probe$")
}
fn contains_needle2<'a>(line: &'a str, needle: &str) -> Option<(&'a str, &'a str, &'a str)> {
if let Some(index) = line.find(needle) {
let rest = &line[index + needle.len()..];
let provider_end = rest.find("$").unwrap();
let provider_name = &rest[..provider_end];
let rest = &rest[provider_end + 1..];
let probe_end = rest.find("$").unwrap();
let probe_name = &rest[..probe_end];
let end = line.rfind("(").unwrap();
let start = line.find(line.split(" ").nth(2).unwrap()).unwrap();
let needle = &line[start..end];
Some((provider_name, probe_name, needle))
} else {
None
}
}
fn build_header_from_provider(source: &str) -> Result<String, crate::Error> {
let mut child = Command::new("dtrace")
.arg("-h")
.arg("-s")
.arg("/dev/stdin")
.arg("-o")
.arg("/dev/stdout")
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.spawn()?;
{
let stdin = child.stdin.as_mut().ok_or(crate::Error::DTraceError)?;
stdin
.write_all(source.as_bytes())
.map_err(|_| crate::Error::DTraceError)?;
}
let output = child.wait_with_output()?;
String::from_utf8(output.stdout).map_err(|_| crate::Error::DTraceError)
}
pub fn register_probes() -> Result<(), crate::Error> {
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::Probe;
#[test]
fn test_is_stability_line() {
let line = "this line is ok \"___dtrace_stability$foo$bar\"";
let result = is_stability_line(line);
assert!(result.is_some());
assert_eq!(result.unwrap().0, "foo");
assert_eq!(result.unwrap().1, "__dtrace_stability$foo$bar");
assert!(is_stability_line("bad").is_none());
}
#[test]
fn test_is_typedefs_line() {
let line = "this line is ok \"___dtrace_typedefs$foo$bar\"";
let result = is_typedefs_line(line);
assert!(result.is_some());
assert_eq!(result.unwrap().0, "foo");
assert_eq!(result.unwrap().1, "__dtrace_typedefs$foo$bar");
assert!(is_typedefs_line("bad").is_none());
}
#[test]
fn test_is_enabled_line() {
let line = "extern int __dtrace_isenabled$foo$bar$xxx(void);";
let result = is_enabled_line(line);
assert!(result.is_some());
assert_eq!(result.unwrap().0, "foo");
assert_eq!(result.unwrap().1, "bar");
assert_eq!(result.unwrap().2, "__dtrace_isenabled$foo$bar$xxx");
assert!(is_enabled_line("bad").is_none());
}
#[test]
fn test_is_probe_line() {
let line = "extern void __dtrace_probe$foo$bar$xxx(whatever);";
let result = is_probe_line(line);
assert!(result.is_some());
assert_eq!(result.unwrap().0, "foo");
assert_eq!(result.unwrap().1, "bar");
assert_eq!(result.unwrap().2, "__dtrace_probe$foo$bar$xxx");
assert!(is_enabled_line("bad").is_none());
}
#[test]
fn test_compile_probe() {
let provider_name = "foo";
let probe_name = "bar";
let extern_probe_name = "__bar";
let is_enabled = "__dtrace_isenabled$foo$bar$xxx";
let probe = "__dtrace_probe$foo$bar$xxx";
let stability = "__dtrace_probe$foo$v1$1_1_1";
let typedefs = "__dtrace_typedefs$foo$v2";
let types = vec![];
let provider = Provider {
name: provider_name.to_string(),
probes: vec![Probe {
name: probe_name.to_string(),
types: types.clone(),
}],
use_statements: vec![],
};
let mut is_enabled_map = BTreeMap::new();
is_enabled_map.insert(String::from(probe_name), String::from(is_enabled));
let mut probes_map = BTreeMap::new();
probes_map.insert(String::from(probe_name), String::from(probe));
let provider_info = ProviderInfo {
stability: String::from(stability),
typedefs: String::from(typedefs),
is_enabled: is_enabled_map,
probes: probes_map,
};
let tokens = compile_probe(
&provider,
probe_name,
&crate::CompileProvidersConfig {
provider: Some(provider_name.to_string()),
..Default::default()
},
&provider_info,
&types,
);
let output = tokens.to_string();
let needle = format!("link_name = \"{is_enabled}\"", is_enabled = is_enabled);
assert!(output.find(&needle).is_some());
let needle = format!("link_name = \"{probe}\"", probe = probe);
assert!(output.find(&needle).is_some());
let needle = format!(
"fn {provider_name}_{probe_name}",
provider_name = provider_name,
probe_name = probe_name
);
assert!(output.find(&needle).is_some());
let needles = &[
"asm ! (\".reference {typedefs}\"",
"call {extern_probe_fn}",
"\".reference {stability}",
"typedefs = sym typedefs",
&format!(
"probe_fn = sym {extern_probe_name}",
extern_probe_name = extern_probe_name
),
"stability = sym stability",
];
for needle in needles.iter() {
println!("{}", needle);
assert!(output.find(needle).is_some());
}
}
}