diff --git a/src/driver/toks_driver_ctx.c b/src/driver/toks_driver_ctx.c index 098a426..a000f59 100644 --- a/src/driver/toks_driver_ctx.c +++ b/src/driver/toks_driver_ctx.c @@ -96,6 +96,62 @@ static int32_t toks_query_performance_counters_failover(nt_filetime * ticks) return 0; } +static int32_t toks_daemon_library_once; + +static int32_t toks_daemon_library_init(void) +{ + int32_t status; + int argc; + char ** argv; + char ** envp; + nt_rtdata * rtdata; + nt_timeout timeout; + + switch (at_locked_cas_32(&toks_daemon_library_once,0,1)) { + case 0: + status = ntapi->tt_get_argv_envp_utf8( + &argc,&argv,&envp, + 0,0,0); + + if (status) { + at_store_32(&toks_daemon_library_once,3); + return status; + } + + if ((status = ntapi->tt_get_runtime_data(&rtdata,0))) { + at_store_32(&toks_daemon_library_once,3); + return status; + } + + ntapi->tt_guid_copy( + &rtdata->srv_guid, + &toks_daemon_default_guid); + + at_locked_inc_32(&toks_daemon_library_once); + return 0; + + case 1: + timeout.quad = -10; + + for (; (at_locked_cas_32(&toks_daemon_library_once,0,1) == 1); ) + ntapi->zw_delay_execution( + NT_SYNC_ALERTABLE, + &timeout); + + return (toks_daemon_library_once == 2) + ? 0 : -1; + + case 2: + return 0; + + case 3: + default: + return -1; + } + + return NT_STATUS_INTERNAL_ERROR; +} + static struct toks_driver_ctx_impl * toks_driver_ctx_alloc( struct argv_meta * meta, const struct toks_common_ctx * cctx) @@ -104,8 +160,19 @@ static struct toks_driver_ctx_impl * toks_driver_ctx_alloc( size_t size; nt_runtime_data * rtdata; - if (ntapi->tt_get_runtime_data(&rtdata,0)) - return 0; + switch (ntapi->tt_get_runtime_data(&rtdata,0)) { + case NT_STATUS_SUCCESS: + break; + + case NT_STATUS_MORE_PROCESSING_REQUIRED: + toks_daemon_library_init(); + + if (ntapi->tt_get_runtime_data(&rtdata,0)) + return 0; + + default: + return 0; + } size = sizeof(struct toks_driver_ctx_alloc);