From c847e3cdf20689879f6b5607eb7e769a1401da14 Mon Sep 17 00:00:00 2001
From: midipix <writeonce@midipix.org>
Date: Jun 08 2020 01:00:10 +0000
Subject: driver: added --tokens support.


---

diff --git a/src/driver/toks_driver_ctx.c b/src/driver/toks_driver_ctx.c
index 269b846..cbf025c 100644
--- a/src/driver/toks_driver_ctx.c
+++ b/src/driver/toks_driver_ctx.c
@@ -183,6 +183,8 @@ int toks_get_driver_ctx(
 	struct argv_entry *		uuid;
 	nt_guid				svcguid;
 	const char *			program;
+	const char *			ch;
+	int				ntokens;
 	void *				hkernel32;
 	char *				targv[TOKS_SARGV_ELEMENTS];
 
@@ -209,6 +211,7 @@ int toks_get_driver_ctx(
 		flags |= TOKS_DRIVER_MODE_SERVER;
 
 	uuid    = 0;
+	ntokens = 0;
 	program = argv_program_name(argv[0]);
 
 	memset(&cctx,0,sizeof(cctx));
@@ -246,6 +249,16 @@ int toks_get_driver_ctx(
 					uuid = entry;
 					break;
 
+				case TAG_TOKENS:
+					for (ntokens=0, ch=entry->arg; *ch && (ntokens>=0); ch++)
+						if ((*ch < '0') || (*ch >'9'))
+							ntokens = -1;
+						else if (ntokens >= 1000)
+							ntokens = -1;
+						else
+							ntokens = ntokens * 10 + (*ch - '0');
+					break;
+
 				case TAG_CONNECT:
 					cctx.drvflags &= ~(uint64_t)TOKS_DRIVER_MODE_SERVER;
 					cctx.drvflags |= TOKS_DRIVER_MODE_CLIENT;
@@ -256,6 +269,14 @@ int toks_get_driver_ctx(
 			return toks_driver_usage(program,0,optv,meta);
 	}
 
+	if ((cctx.drvflags & TOKS_DRIVER_MODE_SERVER) && (ntokens <= 0)) {
+		if (flags & TOKS_DRIVER_VERBOSITY_ERRORS)
+			toks_dprintf(STDERR_FILENO,
+				"%s: error: number of tokens not set or is invalid.",
+				program);
+		return toks_get_driver_ctx_fail(meta);
+	}
+
 	if (uuid && ntapi->tt_string_to_guid_utf8(uuid->arg,&svcguid)) {
 		if (flags & TOKS_DRIVER_VERBOSITY_ERRORS)
 			toks_dprintf(STDERR_FILENO,
@@ -300,15 +321,20 @@ int toks_get_driver_ctx(
 				&(nt_filetime){{0,0}},
 				&ctx->ticks.pcfreq);
 
+	ctx->ntokens            = ntokens;
 	ctx->ctx.program	= program;
 	ctx->ctx.cctx		= &ctx->cctx;
 	ctx->cctx.uuid          = &ctx->uuid;
 
 	toks_daemon_ctx.driver_ctx = &ctx->ctx;
 
-	if (cctx.drvflags & TOKS_DRIVER_MODE_SERVER)
+	if (cctx.drvflags & TOKS_DRIVER_MODE_SERVER) {
+		if (!(ctx->tokens = toks_calloc(ntokens,sizeof(*ctx->tokens))))
+			return toks_get_driver_ctx_fail(meta);
+
 		if (toks_daemon_init(&toks_daemon_ctx,&ctx->uuid))
 			return toks_get_driver_ctx_fail(meta);
+	}
 
 	if (cctx.drvflags & TOKS_DRIVER_MODE_CLIENT)
 		if (toks_client_connect(&ctx->ctx))
@@ -329,6 +355,9 @@ static void toks_free_driver_ctx_impl(struct toks_driver_ctx_alloc * ictx)
 	if (ictx->ctx.hsvclink)
 		ntapi->zw_close(ictx->ctx.hsvclink);
 
+	if (ictx->ctx.tokens)
+		toks_free(ictx->ctx.tokens);
+
 	argv_free(ictx->meta);
 	free(ictx);
 }
diff --git a/src/internal/toksvc_daemon_impl.h b/src/internal/toksvc_daemon_impl.h
index 41f15ae..377c391 100644
--- a/src/internal/toksvc_daemon_impl.h
+++ b/src/internal/toksvc_daemon_impl.h
@@ -25,6 +25,11 @@ struct toks_client_ctx {
 	nt_cid  cid;
 };
 
+struct toks_token {
+	struct _nt_port_keys	keys;
+	struct toks_client_ctx	client;
+};
+
 typedef int32_t __stdcall toks_daemon_routine(struct toks_daemon_ctx *);
 
 struct toks_daemon_ctx {
diff --git a/src/internal/toksvc_driver_impl.h b/src/internal/toksvc_driver_impl.h
index 921e721..c0df86f 100644
--- a/src/internal/toksvc_driver_impl.h
+++ b/src/internal/toksvc_driver_impl.h
@@ -24,6 +24,7 @@ enum app_tags {
 	TAG_SYSROOT,
 	TAG_UUID,
 	TAG_CONNECT,
+	TAG_TOKENS,
 };
 
 struct toks_ticks {
@@ -36,10 +37,12 @@ struct toks_driver_ctx_impl {
 	struct toks_common_ctx	cctx;
 	struct toks_driver_ctx	ctx;
 	struct toks_ticks	ticks;
+	struct toks_token *	tokens;
 	void *			hevent;
 	void *			hsvcdir;
 	void *			hsvclink;
 	void *			hservice;
+	int			ntokens;
 	nt_guid			uuid;
 };
 
@@ -118,6 +121,20 @@ static inline const nt_guid * toks_get_driver_uuid(const struct toks_driver_ctx 
 	return &ictx->uuid;
 }
 
+static inline struct toks_token * toks_get_driver_tokens(const struct toks_driver_ctx * dctx)
+{
+	struct toks_driver_ctx_impl * ictx;
+	ictx = toks_get_driver_ictx(dctx);
+	return ictx->tokens;
+}
+
+static inline int toks_get_driver_ntokens(const struct toks_driver_ctx * dctx)
+{
+	struct toks_driver_ctx_impl * ictx;
+	ictx = toks_get_driver_ictx(dctx);
+	return ictx->ntokens;
+}
+
 static inline void toks_query_performance_counters(const struct toks_driver_ctx * dctx, nt_filetime * ticks)
 {
 	struct toks_driver_ctx_impl * ictx;
diff --git a/src/skin/toks_skin_default.c b/src/skin/toks_skin_default.c
index 2a1d4f7..60a1c9f 100644
--- a/src/skin/toks_skin_default.c
+++ b/src/skin/toks_skin_default.c
@@ -22,6 +22,10 @@ const struct argv_option toks_default_options[] = {
 	{"uuid",		'u',TAG_UUID,ARGV_OPTARG_REQUIRED,0,0,"<uuid>",
 				"set the service identifier to %s."},
 
+	{"tokens",		't',TAG_TOKENS,ARGV_OPTARG_REQUIRED,0,0,"<ntokens>",
+				"set the number of available tokens to %s, which should be "
+				"in the range of 1..9999."},
+
 	{"connect",		'c',TAG_CONNECT,ARGV_OPTARG_NONE,0,0,0,
 				"connect to the server without placing any request"},