Improve compiler macro `dbtest`

Excluded some default ports to avoid permission errors
Also enabled skipping of functions in dbtest modules
Fixed runtime creation statement in `tdb-macros`
next
Sayan Nandan 4 years ago
parent 270dea251f
commit 64c14f70c0
No known key found for this signature in database
GPG Key ID: C31EFD7DDA12AEE0

@ -1,6 +1,6 @@
[package]
name = "tdb_macros"
version = "0.1.0"
version = "0.2.0"
authors = ["Sayan Nandan <ohsayan@outlook.com>"]
edition = "2018"

@ -22,10 +22,11 @@
//! A library containing a collection of custom derives used by TerrabaseDB
use proc_macro::TokenStream;
use proc_macro2::Span;
use quote::quote;
use rand::*;
use std::collections::HashSet;
use syn;
use syn::{self};
// TODO(@ohsayan): Write docs and also make this use the tokio runtime
@ -34,7 +35,7 @@ fn parse_dbtest(mut input: syn::ItemFn, rand: u16) -> Result<TokenStream, syn::E
let fname = sig.ident.to_string();
let body = &input.block;
let attrs = &input.attrs;
let vis = input.vis;
let vis = &input.vis;
let header = quote! {
#[::core::prelude::v1::test]
};
@ -44,24 +45,27 @@ fn parse_dbtest(mut input: syn::ItemFn, rand: u16) -> Result<TokenStream, syn::E
}
sig.asyncness = None;
let body = quote! {
let mut socket = tokio::net::TcpStream::connect(#rand).await.unwrap();
let fut1 = tokio::spawn(crate::tests::start_server(db.clone(), socket));
use crate::tests::{fresp, start_server, terrapipe, QueryVec, TcpStream};
use crate::__func__;
use tokio::prelude::*;
let addr = crate::tests::start_test_server(#rand).await;
let mut stream = tokio::net::TcpStream::connect(&addr).await.unwrap();
#body
socket.shutdown(::std::net::Shutdown::Both).unwrap();
::std::mem::drop(fut1);
stream.shutdown(::std::net::Shutdown::Write).unwrap();
};
let result = quote! {
#header
#(#attrs)*
#vis #sig {
let db = ::std::sync::Arc::new(crate::coredb::CoreDB::new_empty(2));
let runtime = tokio::runtime::Builder::new_multi_thread()
.worker_threads(4)
.thread_name(#fname)
.thread_stack_size(3 * 1024 * 1024)
.build()
.unwrap()
.block_on(async { #body });
tokio::runtime::Builder::new()
.threaded_scheduler()
.core_threads(4)
.thread_name(#fname)
.thread_stack_size(3 * 1024 * 1024)
.enable_all()
.build()
.unwrap()
.block_on(async { #body });
}
};
Ok(result.into())
@ -86,7 +90,7 @@ fn parse_test_sig(input: syn::ItemFn, rand: u16) -> TokenStream {
parse_dbtest(input, rand).unwrap_or_else(|e| e.to_compile_error().into())
}
fn parse_test_module(_args: TokenStream, item: TokenStream) -> TokenStream {
fn parse_test_module(args: TokenStream, item: TokenStream) -> TokenStream {
let input = syn::parse_macro_input!(item as syn::ItemMod);
let content = match input.content {
Some((_, c)) => c,
@ -96,24 +100,87 @@ fn parse_test_module(_args: TokenStream, item: TokenStream) -> TokenStream {
.into()
}
};
let attrs = input.attrs;
let vis = input.vis;
let mod_token = input.mod_token;
let modname = input.ident;
let args = syn::parse_macro_input!(args as syn::AttributeArgs);
let mut skips = Vec::new();
for arg in args {
match arg {
syn::NestedMeta::Meta(syn::Meta::NameValue(namevalue)) => {
let ident = namevalue.path.get_ident();
if ident.is_none() {
let msg = "Must have specified ident";
return syn::Error::new_spanned(namevalue, msg)
.to_compile_error()
.into();
}
match ident.unwrap().to_string().to_lowercase().as_str() {
"skip" => {
let skip_lit = namevalue.lit.clone();
let span = skip_lit.span();
skips = match parse_string(skip_lit, span, "skip") {
Ok(s) => s,
Err(_) => {
return syn::Error::new_spanned(
namevalue,
"Expected a value for argument `skip`",
)
.to_compile_error()
.into();
}
}
.split_whitespace()
.map(|val| val.to_string())
.collect();
}
x => {
let msg = format!("Unknown attribute {} is specified; expected `skip`", x);
return syn::Error::new_spanned(namevalue, msg)
.to_compile_error()
.into();
}
}
}
_ => (),
}
}
let attrs = &input.attrs;
let vis = &input.vis;
let mod_token = &input.mod_token;
let modname = &input.ident;
let mut rng = thread_rng();
let mut in_set = HashSet::<u16>::new();
in_set.insert(80);
in_set.insert(443);
// We will exclude a couple of the default system ports
in_set.insert(80); // HTTP
in_set.insert(443); // SSL
in_set.insert(21); // FTP
in_set.insert(389); // LDAP
in_set.insert(636); // LDAP-SSL
in_set.insert(161); // SNMP
in_set.insert(22); // SSH
in_set.insert(23); // TELNET
in_set.insert(25); // SMTP
in_set.insert(53); // DNS
in_set.insert(119); // NNTP
in_set.insert(143); // IMAP
in_set.insert(993); // IMAP-SSL
let mut result = quote! {};
for item in content {
let mut rand: u16 = rng.gen_range(0, 65535);
let mut rand: u16 = rng.gen_range(1, 65535);
while in_set.contains(&rand) {
rand = rng.gen_range(0, 65535);
rand = rng.gen_range(1, 65535);
}
in_set.insert(rand);
match item {
// We just care about functions, so parse functions and ignore everything
// else
syn::Item::Fn(function) => {
if skips.contains(&function.sig.ident.to_string()) {
result = quote! {
#result
#function
};
continue;
}
let inp = parse_test_sig(function, rand);
let __tok: syn::ItemFn = syn::parse_macro_input!(inp as syn::ItemFn);
let tok = quote! {
@ -131,7 +198,6 @@ fn parse_test_module(_args: TokenStream, item: TokenStream) -> TokenStream {
#result
};
let finalres = quote! {
#(#attrs)*
#mod_token #vis #modname {
#result
}
@ -139,6 +205,17 @@ fn parse_test_module(_args: TokenStream, item: TokenStream) -> TokenStream {
finalres.into()
}
fn parse_string(int: syn::Lit, span: Span, field: &str) -> Result<String, syn::Error> {
match int {
syn::Lit::Str(s) => Ok(s.value()),
syn::Lit::Verbatim(s) => Ok(s.to_string()),
_ => Err(syn::Error::new(
span,
format!("Failed to parse {} into a string.", field),
)),
}
}
#[proc_macro_attribute]
pub fn dbtest(args: TokenStream, item: TokenStream) -> TokenStream {
parse_test_module(args, item)

Loading…
Cancel
Save