@ -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 )