1 files changed, 34 insertions, 18 deletions

diff --git a/src/tty/ntapi_tty_query_information_server.c b/src/tty/ntapi_tty_query_information_server.c
index d1ebf87..a698561 100644
--- a/src/tty/ntapi_tty_query_information_server.c
+++ b/src/tty/ntapi_tty_query_information_server.c
@@ -9,32 +9,48 @@
 #include <ntapi/ntapi.h>
 #include "ntapi_impl.h"
 
-int32_t __stdcall __ntapi_tty_query_information_server(
+#define __MSGBUF_ELEMENTS (NT_LPC_MAX_MSG_DATA_SIZE/sizeof(uintptr_t)) + 1
+
+int32_t	__stdcall __ntapi_tty_query_information_server(
 	__in	void *			hport,
-	__in	nt_tty_server_info *	srvinfo)
+	__out	nt_io_status_block *	iosb,
+	__out	void *			tty_info,
+	__in	uint32_t		tty_info_length,
+	__in	int			tty_info_class)
 {
-	int32_t		  status;
-	nt_tty_server_msg msg;
-
-	hport = hport ? hport : __ntapi_internals()->hport_tty_session;
+	int32_t			status;
+	nt_tty_server_msg *	msg;
+	uintptr_t		msgbuf[__MSGBUF_ELEMENTS];
 
 	__ntapi->tt_aligned_block_memset(
-		&msg,0,sizeof(msg));
+		msgbuf,0,
+		sizeof(msgbuf));
+
+	msg				= (nt_tty_server_msg *)msgbuf;
+	msg->header.msg_type		= NT_LPC_NEW_MESSAGE;
+	msg->header.data_size		= sizeof(msg->data);
+	msg->header.msg_size		= sizeof(*msg);
+	msg->data.ttyinfo.opcode	= NT_TTY_QUERY_INFORMATION_SERVER;
+	msg->data.srvinfo.info_class	= tty_info_class;
 
-	msg.header.msg_type	= NT_LPC_NEW_MESSAGE;
-	msg.header.data_size	= sizeof(msg.data);
-	msg.header.msg_size	= sizeof(msg);
-	msg.data.ttyinfo.opcode	= NT_TTY_QUERY_INFORMATION_SERVER;
+	if (!hport)
+		hport = __ntapi_internals()->hport_tty_session;
 
-	if ((status = __ntapi->zw_request_wait_reply_port(hport,&msg,&msg)))
+	if ((status = __ntapi->zw_request_wait_reply_port(hport,msg,msg)))
 		return status;
-	else if (msg.data.ttyinfo.status)
-		return msg.data.ttyinfo.status;
+	else if (msg->data.ttyinfo.status)
+		return msg->data.ttyinfo.status;
+
+	if (tty_info_length < msg->data.srvinfo.info_length)
+		return NT_STATUS_BUFFER_TOO_SMALL;
+
+	__ntapi->tt_generic_memcpy(
+		tty_info,
+		msg->data.srvinfo.info_buffer,
+		msg->data.srvinfo.info_length);
 
-	__ntapi->tt_aligned_block_memcpy(
-		(uintptr_t *)srvinfo,
-		(uintptr_t *)&(msg.data.srvinfo),
-		sizeof(*srvinfo));
+	iosb->status = NT_STATUS_SUCCESS;
+	iosb->info   = msg->data.srvinfo.info_length;
 
 	return NT_STATUS_SUCCESS;
 }