2 files changed, 56 insertions, 21 deletions

diff --git a/include/ntapi/nt_tty.h b/include/ntapi/nt_tty.h
index 4c626c2..06dbd08 100644
--- a/include/ntapi/nt_tty.h
+++ b/include/ntapi/nt_tty.h
@@ -87,6 +87,16 @@ typedef enum _nt_tty_session_type {
 } nt_tty_session_type;
 
 
+typedef enum _nt_tty_server_info_class {
+	NT_TTY_SERVER_BASIC_INFORMATION,
+	NT_TTY_SERVER_PTM_SLOT_INFORMATION,
+	NT_TTY_SERVER_PTS_SLOT_INFORMATION,
+	NT_TTY_SERVER_DBM_SLOT_INFORMATION,
+	NT_TTY_SERVER_DBG_SLOT_INFORMATION,
+	NT_TTY_SERVER_INFORMATION_CAP
+} nt_tty_server_info_class;
+
+
 typedef enum _nt_pty_info_class {
 	NT_PTY_BASIC_INFORMATION,
 	NT_PTY_CLIENT_INFORMATION,
@@ -115,13 +125,19 @@ typedef struct __attr_ptr_size_aligned__ _nt_tty_register_info {
 } nt_tty_register_info;
 
 
-typedef struct __attr_ptr_size_aligned__ _nt_tty_server_info {
+typedef struct __attr_ptr_size_aligned__ _nt_tty_server_basic_info {
 	nt_port_attr	attr;
 	intptr_t	pid;
 	intptr_t	tid;
-} nt_tty_server_info;
+} nt_tty_server_basic_info;
 
 
+typedef struct __attr_ptr_size_aligned__ _nt_tty_server_info {
+	nt_tty_server_info_class	info_class;
+	uint32_t			info_length;
+	uintptr_t			info_buffer[];
+} nt_tty_server_info;
+
 typedef struct __attr_ptr_size_aligned__ _nt_tty_vms_info {
 	void *		hroot;
 	uint32_t	hash;
@@ -372,7 +388,10 @@ typedef int32_t __stdcall	ntapi_tty_client_process_register(
 
 typedef int32_t __stdcall ntapi_tty_query_information_server(
 	__in	void *			hport,
-	__out	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);
 
 
 /* pty api */
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;
 }