From 971724532fcc7c55576929053d58992dab69f949 Mon Sep 17 00:00:00 2001
From: midipix <writeonce@midipix.org>
Date: Jun 21 2020 10:13:55 +0000
Subject: daemon: keep record of the number of unassigned tokens.


---

diff --git a/src/daemon/toks_daemon_acquire.c b/src/daemon/toks_daemon_acquire.c
index 7fd02c1..44884e1 100644
--- a/src/daemon/toks_daemon_acquire.c
+++ b/src/daemon/toks_daemon_acquire.c
@@ -268,7 +268,7 @@ int32_t __stdcall toks_daemon_acquire(struct toks_daemon_ctx * dctx)
 
 	toks_query_performance_counters(dctx->driver_ctx,&pcnt);
 
-	if (dctx->utokens >= toks_get_driver_ntokens(dctx->driver_ctx)) {
+	if (dctx->ftokens == 0) {
 		status = ntapi->zw_open_process(
 			&hcaller,
 			NT_PROCESS_SYNCHRONIZE
@@ -363,6 +363,7 @@ int32_t __stdcall toks_daemon_acquire(struct toks_daemon_ctx * dctx)
 	}
 
 	dctx->utokens++;
+	dctx->ftokens--;
 
 	return NT_STATUS_SUCCESS;
 }
diff --git a/src/daemon/toks_daemon_ioctl.c b/src/daemon/toks_daemon_ioctl.c
index 221d3ed..8d60a04 100644
--- a/src/daemon/toks_daemon_ioctl.c
+++ b/src/daemon/toks_daemon_ioctl.c
@@ -34,6 +34,9 @@ static int32_t toks_daemon_ioctl_set_token_count(struct toks_daemon_ctx * dctx)
 
 	toks_set_driver_ntokens(dctx->driver_ctx,ntokens);
 
+	dctx->ftokens = (dctx->utokens < ntokens)
+		? ntokens - dctx->utokens : 0;
+
 	return toks_daemon_release(dctx);
 }
 
diff --git a/src/daemon/toks_daemon_loop.c b/src/daemon/toks_daemon_loop.c
index 505c0b5..3fd5ca9 100644
--- a/src/daemon/toks_daemon_loop.c
+++ b/src/daemon/toks_daemon_loop.c
@@ -171,6 +171,7 @@ int32_t __stdcall toks_daemon_loop(void * ctx)
 
 	dctx->ctrlpid = toks_get_driver_ctrlpid(dctx->driver_ctx);
 	dctx->csyspid = toks_get_driver_csyspid(dctx->driver_ctx);
+	dctx->ftokens = toks_get_driver_ntokens(dctx->driver_ctx);
 
 	svcvtbl = (dctx->driver_ctx->cctx->drvflags & TOKS_DRIVER_MODE_SERVER)
 		? toks_daemon_vtbl
diff --git a/src/daemon/toks_daemon_release.c b/src/daemon/toks_daemon_release.c
index d69e42e..db05dcc 100644
--- a/src/daemon/toks_daemon_release.c
+++ b/src/daemon/toks_daemon_release.c
@@ -124,7 +124,8 @@ int32_t __stdcall toks_daemon_release(struct toks_daemon_ctx * dctx)
 	if (!token)
 		return NT_STATUS_INVALID_PARAMETER;
 
-	dctx->utokens--;
+	if (--dctx->utokens < toks_get_driver_ntokens(dctx->driver_ctx))
+		dctx->ftokens++;
 
 	toks_daemon_token_reset(token);