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


---

diff --git a/src/client/toks_client_acquire.c b/src/client/toks_client_acquire.c
index 05dbab8..97e998b 100644
--- a/src/client/toks_client_acquire.c
+++ b/src/client/toks_client_acquire.c
@@ -29,6 +29,8 @@ int32_t toks_client_acquire(struct toks_driver_ctx * dctx)
 	msg.header.msg_size	= sizeof(msg);
 	msg.data.ttyinfo.opcode	= TOKS_DAEMON_ACQUIRE;
 
+	msg.data.ipcinfo.ctrlsvc.keys.key[0] = toks_get_driver_tokpid(dctx);
+
 	if ((status = ntapi->zw_request_wait_reply_port(hport,&msg,&msg)))
 		return status;
 
diff --git a/src/driver/toks_driver_ctx.c b/src/driver/toks_driver_ctx.c
index 4c748f3..c42e1ca 100644
--- a/src/driver/toks_driver_ctx.c
+++ b/src/driver/toks_driver_ctx.c
@@ -181,10 +181,12 @@ int toks_get_driver_ctx(
 	struct argv_meta *		meta;
 	struct argv_entry *		entry;
 	struct argv_entry *		uuid;
+	struct argv_entry *		pid;
 	nt_guid				svcguid;
 	const char *			program;
 	const char *			ch;
 	int				ntokens;
+	int32_t				tokpid;
 	void *				hkernel32;
 	char *				targv[TOKS_SARGV_ELEMENTS];
 
@@ -211,6 +213,7 @@ int toks_get_driver_ctx(
 		flags |= TOKS_DRIVER_MODE_SERVER;
 
 	uuid    = 0;
+	tokpid  = 0;
 	ntokens = 0;
 	program = argv_program_name(argv[0]);
 
@@ -249,6 +252,16 @@ int toks_get_driver_ctx(
 					uuid = entry;
 					break;
 
+				case TAG_PID:
+					pid = entry;
+
+					for (tokpid=0, ch=entry->arg; *ch && (tokpid>=0); ch++)
+						if ((*ch < '0') || (*ch >'9'))
+							tokpid = -1;
+						else
+							tokpid = tokpid * 10 + (*ch - '0');
+					break;
+
 				case TAG_TOKENS:
 					for (ntokens=0, ch=entry->arg; *ch && (ntokens>=0); ch++)
 						if ((*ch < '0') || (*ch >'9'))
@@ -275,6 +288,14 @@ int toks_get_driver_ctx(
 			return toks_driver_usage(program,0,optv,meta);
 	}
 
+	if ((cctx.drvflags & TOKS_DRIVER_MODE_CLIENT) && (tokpid < 0)) {
+		if (flags & TOKS_DRIVER_VERBOSITY_ERRORS)
+			toks_dprintf(STDERR_FILENO,
+				"%s: error: %s is not a valid process id.",
+				program,pid->arg);
+		return toks_get_driver_ctx_fail(meta);
+	}
+
 	if ((cctx.drvflags & TOKS_DRIVER_MODE_SERVER) && (ntokens <= 0)) {
 		if (flags & TOKS_DRIVER_VERBOSITY_ERRORS)
 			toks_dprintf(STDERR_FILENO,
@@ -327,6 +348,7 @@ int toks_get_driver_ctx(
 				&(nt_filetime){{0,0}},
 				&ctx->ticks.pcfreq);
 
+	ctx->tokpid             = tokpid;
 	ctx->ntokens            = ntokens;
 	ctx->ctx.program	= program;
 	ctx->ctx.cctx		= &ctx->cctx;
diff --git a/src/internal/toksvc_driver_impl.h b/src/internal/toksvc_driver_impl.h
index 2fec8b6..8825d47 100644
--- a/src/internal/toksvc_driver_impl.h
+++ b/src/internal/toksvc_driver_impl.h
@@ -26,6 +26,7 @@ enum app_tags {
 	TAG_CONNECT,
 	TAG_TOKENS,
 	TAG_ACQUIRE,
+	TAG_PID,
 };
 
 struct toks_ticks {
@@ -46,6 +47,7 @@ struct toks_driver_ctx_impl {
 	void *			hsvclink;
 	void *			hservice;
 	int			ntokens;
+	int			tokpid;
 	nt_guid			uuid;
 };
 
@@ -159,6 +161,13 @@ static inline int toks_get_driver_ntokens(const struct toks_driver_ctx * dctx)
 	return ictx->ntokens;
 }
 
+static inline int toks_get_driver_tokpid(const struct toks_driver_ctx * dctx)
+{
+	struct toks_driver_ctx_impl * ictx;
+	ictx = toks_get_driver_ictx(dctx);
+	return ictx->tokpid;
+}
+
 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 2a806ea..cd5c2be 100644
--- a/src/skin/toks_skin_default.c
+++ b/src/skin/toks_skin_default.c
@@ -32,5 +32,9 @@ const struct argv_option toks_default_options[] = {
 	{"acquire",		'a',TAG_ACQUIRE,ARGV_OPTARG_NONE,0,0,0,
 				"acquire the next available token"},
 
+	{"pid",			'p',TAG_PID,ARGV_OPTARG_REQUIRED,0,0,"<pid>",
+				"assign the token to the framework process "
+				"identified by %s"},
+
 	{0,0,0,0,0,0,0,0}
 };