diff --git a/src/daemon/toks_daemon_acquire.c b/src/daemon/toks_daemon_acquire.c
index 60c3c10..7fd02c1 100644
--- a/src/daemon/toks_daemon_acquire.c
+++ b/src/daemon/toks_daemon_acquire.c
@@ -256,7 +256,7 @@ int32_t __stdcall toks_daemon_acquire(struct toks_daemon_ctx * dctx)
}
token = toks_get_driver_tokens(dctx->driver_ctx);
- tocap = &token[toks_get_driver_ntokens(dctx->driver_ctx)];
+ tocap = &token[toks_get_driver_atokens(dctx->driver_ctx)];
for (; token->self && (token<tocap); )
token++;
@@ -268,7 +268,7 @@ int32_t __stdcall toks_daemon_acquire(struct toks_daemon_ctx * dctx)
toks_query_performance_counters(dctx->driver_ctx,&pcnt);
- if (token == tocap) {
+ if (dctx->utokens >= toks_get_driver_ntokens(dctx->driver_ctx)) {
status = ntapi->zw_open_process(
&hcaller,
NT_PROCESS_SYNCHRONIZE
@@ -362,5 +362,7 @@ int32_t __stdcall toks_daemon_acquire(struct toks_daemon_ctx * dctx)
return status;
}
+ dctx->utokens++;
+
return NT_STATUS_SUCCESS;
}
diff --git a/src/daemon/toks_daemon_ioctl.c b/src/daemon/toks_daemon_ioctl.c
index d9d7f99..373e525 100644
--- a/src/daemon/toks_daemon_ioctl.c
+++ b/src/daemon/toks_daemon_ioctl.c
@@ -18,6 +18,22 @@ static int32_t toks_daemon_ioctl_get_token_count(struct toks_daemon_ctx * dctx)
return NT_STATUS_SUCCESS;
}
+static int32_t toks_daemon_ioctl_set_token_count(struct toks_daemon_ctx * dctx)
+{
+ nt_tty_port_msg * msg;
+ int ntokens;
+
+ msg = &dctx->reply;
+ ntokens = (int)(intptr_t)msg->ttyinfo.exarg;
+
+ if (ntokens > toks_get_driver_atokens(dctx->driver_ctx))
+ return NT_STATUS_BUFFER_TOO_SMALL;
+
+ toks_set_driver_ntokens(dctx->driver_ctx,ntokens);
+
+ return toks_daemon_release(dctx);
+}
+
int32_t __stdcall toks_daemon_ioctl(struct toks_daemon_ctx * dctx)
{
nt_tty_port_msg * msg;
@@ -30,6 +46,9 @@ int32_t __stdcall toks_daemon_ioctl(struct toks_daemon_ctx * dctx)
case TOKS_IOCTL_GET_TOKEN_COUNT:
return toks_daemon_ioctl_get_token_count(dctx);
+ case TOKS_IOCTL_SET_TOKEN_COUNT:
+ return toks_daemon_ioctl_set_token_count(dctx);
+
default:
return NT_STATUS_NOT_IMPLEMENTED;
}
diff --git a/src/daemon/toks_daemon_release.c b/src/daemon/toks_daemon_release.c
index 0c4220d..d69e42e 100644
--- a/src/daemon/toks_daemon_release.c
+++ b/src/daemon/toks_daemon_release.c
@@ -76,6 +76,24 @@ static int32_t toks_daemon_unqueue(struct toks_daemon_ctx * dctx)
return status;
}
+static int toks_daemon_unqueue_ioctl(struct toks_daemon_ctx * dctx)
+{
+ int ntokens = toks_get_driver_ntokens(dctx->driver_ctx);
+
+ dctx->opcode = TOKS_DAEMON_RELEASE;
+
+ while (dctx->utokens < ntokens) {
+ if (toks_get_driver_nwaiters(dctx->driver_ctx) == 0)
+ return NT_STATUS_SUCCESS;
+
+ toks_daemon_unqueue(dctx);
+ }
+
+ dctx->opcode = TOKS_DAEMON_IOCTL;
+
+ return NT_STATUS_SUCCESS;
+}
+
int32_t __stdcall toks_daemon_release(struct toks_daemon_ctx * dctx)
{
nt_tty_port_msg * msg;
@@ -84,11 +102,14 @@ int32_t __stdcall toks_daemon_release(struct toks_daemon_ctx * dctx)
struct toks_token * tocap;
uint32_t * keys;
+ if (dctx->opcode == TOKS_DAEMON_IOCTL)
+ return toks_daemon_unqueue_ioctl(dctx);
+
msg = &dctx->reply;
keys = msg->syncinfo.ipckeys;
toptr = toks_get_driver_tokens(dctx->driver_ctx);
- tocap = &toptr[toks_get_driver_ntokens(dctx->driver_ctx)];
+ tocap = &toptr[toks_get_driver_atokens(dctx->driver_ctx)];
for (token=0; !token && (toptr<tocap); toptr++)
if ((toptr->self)
@@ -103,6 +124,8 @@ int32_t __stdcall toks_daemon_release(struct toks_daemon_ctx * dctx)
if (!token)
return NT_STATUS_INVALID_PARAMETER;
+ dctx->utokens--;
+
toks_daemon_token_reset(token);
ntapi->zw_alert_thread(token->client.hinstance);
diff --git a/src/driver/toks_driver_ctx.c b/src/driver/toks_driver_ctx.c
index 44a77cc..5156d47 100644
--- a/src/driver/toks_driver_ctx.c
+++ b/src/driver/toks_driver_ctx.c
@@ -291,6 +291,7 @@ int toks_get_driver_ctx(
const char * program;
const char * refstr;
int ntokens;
+ int atokens;
int loglevel;
int32_t tokpid;
int32_t tsyspid;
@@ -592,12 +593,18 @@ int toks_get_driver_ctx(
&ctx->ticks.pcfreq);
+ atokens = ntokens;
+ atokens += 0x1fff;
+ atokens |= 0xfff;
+ atokens ^= 0xfff;
+
ctx->hlog = hlog;
ctx->tokpid = tokpid;
ctx->tsyspid = tsyspid;
ctx->ctrlpid = ctrlpid;
ctx->csyspid = csyspid;
ctx->ntokens = ntokens;
+ ctx->atokens = atokens;
ctx->timeout = timeout;
ctx->ctx.program = program;
ctx->ctx.cctx = &ctx->cctx;
@@ -621,7 +628,7 @@ int toks_get_driver_ctx(
toks_daemon_ctx.waiter_cap = &ctx->waiters[TOKS_MAX_WAITERS];
- if (!(ctx->tokens = toks_calloc(ntokens,sizeof(*ctx->tokens))))
+ if (!(ctx->tokens = toks_calloc(atokens,sizeof(*ctx->tokens))))
return toks_get_driver_ctx_fail(meta);
if (toks_daemon_init(&toks_daemon_ctx,&ctx->uuid))
diff --git a/src/internal/toksvc_daemon_impl.h b/src/internal/toksvc_daemon_impl.h
index 273816d..b93e1ad 100644
--- a/src/internal/toksvc_daemon_impl.h
+++ b/src/internal/toksvc_daemon_impl.h
@@ -61,6 +61,9 @@ struct toks_daemon_ctx {
nt_port_attr daemon_attr;
nt_port_name daemon_name;
+ int ftokens;
+ int utokens;
+
intptr_t ctrlpid;
intptr_t csyspid;
diff --git a/src/internal/toksvc_driver_impl.h b/src/internal/toksvc_driver_impl.h
index fdd3798..ad0ea05 100644
--- a/src/internal/toksvc_driver_impl.h
+++ b/src/internal/toksvc_driver_impl.h
@@ -69,6 +69,7 @@ struct toks_driver_ctx_impl {
void * hlog;
int nwaiters;
int ntokens;
+ int atokens;
int tokpid;
int tsyspid;
int ctrlpid;
@@ -228,6 +229,20 @@ static inline int toks_get_driver_ntokens(const struct toks_driver_ctx * dctx)
return ictx->ntokens;
}
+static inline void toks_set_driver_ntokens(const struct toks_driver_ctx * dctx, int ntokens)
+{
+ struct toks_driver_ctx_impl * ictx;
+ ictx = toks_get_driver_ictx(dctx);
+ ictx->ntokens = ntokens;
+}
+
+static inline int toks_get_driver_atokens(const struct toks_driver_ctx * dctx)
+{
+ struct toks_driver_ctx_impl * ictx;
+ ictx = toks_get_driver_ictx(dctx);
+ return ictx->atokens;
+}
+
static inline int toks_get_driver_tokpid(const struct toks_driver_ctx * dctx)
{
struct toks_driver_ctx_impl * ictx;