diff --git a/dtrace-parser/src/lib.rs b/dtrace-parser/src/lib.rs index 9385a844..084a1f0f 100644 --- a/dtrace-parser/src/lib.rs +++ b/dtrace-parser/src/lib.rs @@ -156,6 +156,7 @@ impl Integer { pub enum DataType { Integer(Integer), Pointer(Integer), + CString, String, } @@ -246,7 +247,7 @@ impl DataType { match self { DataType::Integer(int) => int.to_c_type(), DataType::Pointer(int) => format!("{}*", int.to_c_type()), - DataType::String => String::from("char*"), + DataType::CString | DataType::String => String::from("char*"), } } @@ -255,7 +256,7 @@ impl DataType { match self { DataType::Integer(int) => int.to_rust_ffi_type(), DataType::Pointer(int) => format!("*const {}", int.to_rust_ffi_type()), - DataType::String => format!("*const {RUST_TYPE_PREFIX}char"), + DataType::CString | DataType::String => format!("*const {RUST_TYPE_PREFIX}char"), } } @@ -265,6 +266,7 @@ impl DataType { DataType::Integer(int) => int.to_rust_type(), DataType::Pointer(int) => format!("*const {}", int.to_rust_type()), DataType::String => String::from("&str"), + DataType::CString => String::from("&::core::ffi::CStr"), } } } diff --git a/probe-test-attr/src/main.rs b/probe-test-attr/src/main.rs index b409c97a..ec3472b2 100644 --- a/probe-test-attr/src/main.rs +++ b/probe-test-attr/src/main.rs @@ -82,6 +82,9 @@ mod test { /// Constant pointers to integer types are also supported fn work_with_pointer(_buffer: *const u8, _: u64) {} + + /// C-string types are also supported + fn cstring(_: &CStr, _: CString) {} } fn main() { @@ -106,5 +109,6 @@ fn main() { test::arg_as_tuple!(|| (arg.x, &arg.buffer[..])); test::not_json_serializable!(|| Whoops::NoBueno(0)); test::work_with_pointer!(|| (buffer.as_ptr(), buffer.len() as u64)); + test::cstring!(|| (c"hello world", c"and when owned".to_owned())); } } diff --git a/usdt-attr-macro/src/lib.rs b/usdt-attr-macro/src/lib.rs index 21d3ee5c..21f041f7 100644 --- a/usdt-attr-macro/src/lib.rs +++ b/usdt-attr-macro/src/lib.rs @@ -303,6 +303,8 @@ fn is_simple_type(ident: &syn::Ident) -> bool { | "i64" | "String" | "str" + | "CString" + | "CStr" | "usize" | "isize" ) @@ -363,6 +365,8 @@ fn data_type_from_path(path: &syn::Path, pointer: bool) -> DataType { })) } else if path.is_ident("String") || path.is_ident("str") { DataType::Native(DType::String) + } else if path.is_ident("CString") || path.is_ident("CStr") { + DataType::Native(DType::CString) } else if path.is_ident("isize") { DataType::Native(variant(Integer { sign: Sign::Signed, @@ -462,6 +466,10 @@ mod tests { #[case("String", DType::String)] #[case("&&str", DType::String)] #[case("&String", DType::String)] + #[case("&CStr", DType::CString)] + #[case("CString", DType::CString)] + #[case("&&CStr", DType::CString)] + #[case("&CString", DType::CString)] fn test_parse_probe_argument_native(#[case] name: &str, #[case] ty: dtrace_parser::DataType) { let arg = syn::parse_str(name).unwrap(); let out = parse_probe_argument(&arg, 0, 0).unwrap(); diff --git a/usdt-impl/src/common.rs b/usdt-impl/src/common.rs index e56b8d47..07e78dc0 100644 --- a/usdt-impl/src/common.rs +++ b/usdt-impl/src/common.rs @@ -67,6 +67,9 @@ pub fn construct_type_check( } } DataType::Native(dtrace_parser::DataType::String) => quote! { _: impl AsRef }, + DataType::Native(dtrace_parser::DataType::CString) => { + quote! { _: impl AsRef<::core::ffi::CStr> } + } _ => { let arg = typ.to_rust_type(); quote! { _: impl ::std::borrow::Borrow<#arg> } @@ -181,6 +184,10 @@ fn asm_type_convert(typ: &DataType, input: TokenStream) -> (TokenStream, TokenSt }, quote! { .as_ptr() as usize }, ), + DataType::Native(dtrace_parser::DataType::CString) => ( + quote! { #input.as_ref() as &::core::ffi::CStr }, + quote! { .as_ptr() as usize }, + ), DataType::Native(dtrace_parser::DataType::String) => ( quote! { [(#input.as_ref() as &str).as_bytes(), &[0_u8]].concat() @@ -397,5 +404,15 @@ mod tests { quote! { [(foo.as_ref() as &str).as_bytes(), &[0_u8]].concat() }.to_string() ); assert_eq!(post.to_string(), quote! { .as_ptr() as usize }.to_string()); + + let (out, post) = asm_type_convert( + &DataType::Native(dtrace_parser::DataType::CString), + TokenStream::from_str("foo").unwrap(), + ); + assert_eq!( + out.to_string(), + quote! { foo.as_ref() as &::core::ffi::CStr }.to_string() + ); + assert_eq!(post.to_string(), quote! { .as_ptr() as usize }.to_string()); } }