From 778c19f4c0ae4fa90b7e90f1a11ff6dc9264c19a Mon Sep 17 00:00:00 2001 From: Leonard Hecker Date: Fri, 13 Mar 2026 23:08:41 +0100 Subject: [PATCH] wip --- src/conpty/PtyServer.clients.cpp | 134 ++++++++++++++++++++-- src/conpty/PtyServer.cpp | 20 ++-- src/conpty/PtyServer.h | 50 +++++++- src/conpty/PtyServer.io.cpp | 183 ++++++++++++++++++++++++++++++ src/conpty/conpty.vcxproj | 1 + src/conpty/conpty.vcxproj.filters | 3 + src/conpty/pch.h | 4 +- 7 files changed, 378 insertions(+), 17 deletions(-) create mode 100644 src/conpty/PtyServer.io.cpp diff --git a/src/conpty/PtyServer.clients.cpp b/src/conpty/PtyServer.clients.cpp index 0137393eae..43d4341646 100644 --- a/src/conpty/PtyServer.clients.cpp +++ b/src/conpty/PtyServer.clients.cpp @@ -1,14 +1,12 @@ #include "pch.h" #include "PtyServer.h" -#include - // Reads [part of] the input payload of the given message from the driver. // Analogous to the OG ReadMessageInput() in csrutil.cpp. // // For CONSOLE_IO_CONNECT, offset is 0 and the payload is a CONSOLE_SERVER_MSG. // For CONSOLE_IO_USER_DEFINED, offset would typically be past the message header. -__declspec(noinline) void PtyServer::readInput(const CD_IO_DESCRIPTOR& desc, ULONG offset, void* buffer, ULONG size) +void PtyServer::readInput(const CD_IO_DESCRIPTOR& desc, ULONG offset, void* buffer, ULONG size) { CD_IO_OPERATION op{}; op.Identifier = desc.Identifier; @@ -29,6 +27,21 @@ void PtyServer::completeIo(CD_IO_COMPLETE& completion) THROW_IF_NTSTATUS_FAILED(ioctl(IOCTL_CONDRV_COMPLETE_IO, &completion, sizeof(completion), nullptr, 0)); } +// Writes data back to the client's output buffer for the given message. +// Analogous to the IOCTL_CONDRV_WRITE_OUTPUT call in the OG ReleaseMessageBuffers() (csrutil.cpp). +// +// The driver matches the Identifier to the pending IO and copies data into +// the client's buffer at the specified offset. +void PtyServer::writeOutput(const CD_IO_DESCRIPTOR& desc, ULONG offset, const void* buffer, ULONG size) +{ + CD_IO_OPERATION op{}; + op.Identifier = desc.Identifier; + op.Buffer.Offset = offset; + op.Buffer.Data = const_cast(buffer); + op.Buffer.Size = size; + THROW_IF_NTSTATUS_FAILED(ioctl(IOCTL_CONDRV_WRITE_OUTPUT, &op, sizeof(op), nullptr, 0)); +} + PtyClient* PtyServer::findClient(ULONG_PTR handle) { auto ptr = reinterpret_cast(handle); @@ -81,10 +94,12 @@ void PtyServer::handleConnect(CONSOLE_API_MSG& msg) // 4. The first connection is the root process (console owner). client->rootProcess = !m_initialized; - // 5. Allocate opaque handle IDs for input and output. - // The driver echoes these back in Descriptor.Object for future IO messages. - client->inputHandle = m_nextHandleId++; - client->outputHandle = m_nextHandleId++; + // 5. Allocate IO handles for input and output. + // In the OG, AllocateIoHandle creates a CONSOLE_HANDLE_DATA pointing to + // the input buffer or screen buffer. Here we create lightweight PtyHandle + // objects. The driver echoes these back in Descriptor.Object. + client->inputHandle = allocateHandle(CONSOLE_INPUT_HANDLE, GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE); + client->outputHandle = allocateHandle(CONSOLE_OUTPUT_HANDLE, GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE); if (!m_initialized) { @@ -132,5 +147,110 @@ void PtyServer::handleDisconnect(CONSOLE_API_MSG& msg) return; } + // Cancel any pending IOs from this client. + cancelPendingIOs(msg.Descriptor.Process); + + // Free the client's IO handles, mirroring OG FreeProcessData which calls + // ConsoleCloseHandle on InputHandle and OutputHandle. + if (client->inputHandle) + { + freeHandle(client->inputHandle); + } + if (client->outputHandle) + { + freeHandle(client->outputHandle); + } + std::erase_if(m_clients, [client](const auto& c) { return c.get() == client; }); } + +// Handle management. +// +// Analogous to OG AllocateIoHandle (handle.cpp). The OG creates a CONSOLE_HANDLE_DATA +// with share/access tracking and a pointer to the underlying console object. +// We create a lightweight PtyHandle and return its pointer cast to ULONG_PTR. +ULONG_PTR PtyServer::allocateHandle(ULONG handleType, ACCESS_MASK access, ULONG shareMode) +{ + auto h = std::make_unique(); + h->handleType = handleType; + h->access = access; + h->shareMode = shareMode; + auto ptr = reinterpret_cast(h.get()); + m_handles.push_back(std::move(h)); + return ptr; +} + +// Analogous to OG ConsoleCloseHandle → FreeConsoleHandle (handle.cpp). +void PtyServer::freeHandle(ULONG_PTR handle) +{ + auto ptr = reinterpret_cast(handle); + std::erase_if(m_handles, [ptr](const auto& h) { return h.get() == ptr; }); +} + +// Handles CONSOLE_IO_CREATE_OBJECT messages. +// +// Protocol (from OG ConsoleCreateObject in srvinit.cpp): +// 1. Read CD_CREATE_OBJECT_INFORMATION from the message (already in msg.CreateObject). +// 2. Resolve CD_IO_OBJECT_TYPE_GENERIC based on DesiredAccess. +// 3. Allocate a handle of the appropriate type. +// 4. Reply via completeIo with the handle value in IoStatus.Information. +void PtyServer::handleCreateObject(CONSOLE_API_MSG& msg) +{ + auto& info = msg.CreateObject; + + // Resolve generic object type based on desired access, matching OG behavior. + if (info.ObjectType == CD_IO_OBJECT_TYPE_GENERIC) + { + if ((info.DesiredAccess & (GENERIC_READ | GENERIC_WRITE)) == GENERIC_READ) + { + info.ObjectType = CD_IO_OBJECT_TYPE_CURRENT_INPUT; + } + else if ((info.DesiredAccess & (GENERIC_READ | GENERIC_WRITE)) == GENERIC_WRITE) + { + info.ObjectType = CD_IO_OBJECT_TYPE_CURRENT_OUTPUT; + } + } + + ULONG_PTR handle = 0; + + switch (info.ObjectType) + { + case CD_IO_OBJECT_TYPE_CURRENT_INPUT: + handle = allocateHandle(CONSOLE_INPUT_HANDLE, info.DesiredAccess, info.ShareMode); + break; + + case CD_IO_OBJECT_TYPE_CURRENT_OUTPUT: + handle = allocateHandle(CONSOLE_OUTPUT_HANDLE, info.DesiredAccess, info.ShareMode); + break; + + case CD_IO_OBJECT_TYPE_NEW_OUTPUT: + // In the OG, this creates a new screen buffer via ConsoleCreateScreenBuffer. + // For now, we allocate a handle that tracks as an output handle. + handle = allocateHandle(CONSOLE_OUTPUT_HANDLE, info.DesiredAccess, info.ShareMode); + break; + + default: + THROW_NTSTATUS(STATUS_INVALID_PARAMETER); + } + + // Reply with the handle value in IoStatus.Information. + // The driver stores this and echoes it back in Descriptor.Object for future IO. + CD_IO_COMPLETE completion{}; + completion.Identifier = msg.Descriptor.Identifier; + completion.IoStatus.Status = STATUS_SUCCESS; + completion.IoStatus.Information = handle; + + completeIo(completion); +} + +// Handles CONSOLE_IO_CLOSE_OBJECT messages. +// +// Protocol (from OG SrvCloseHandle in stream.cpp): +// 1. Descriptor.Object contains the opaque handle value. +// 2. Close/free the handle. +// +// The caller replies with STATUS_SUCCESS inline. +void PtyServer::handleCloseObject(CONSOLE_API_MSG& msg) +{ + freeHandle(msg.Descriptor.Object); +} diff --git a/src/conpty/PtyServer.cpp b/src/conpty/PtyServer.cpp index 7dfd85925d..04ede71aa2 100644 --- a/src/conpty/PtyServer.cpp +++ b/src/conpty/PtyServer.cpp @@ -135,23 +135,29 @@ HRESULT PtyServer::Run() break; case CONSOLE_IO_CREATE_OBJECT: printf("Received create object request for object %llu from process %llu\n", req.Descriptor.Object, req.Descriptor.Process); - res.IoStatus.Status = STATUS_NOT_IMPLEMENTED; - hasRes = true; + handleCreateObject(req); break; case CONSOLE_IO_CLOSE_OBJECT: printf("Received close object request for object %llu from process %llu\n", req.Descriptor.Object, req.Descriptor.Process); - res.IoStatus.Status = STATUS_NOT_IMPLEMENTED; + handleCloseObject(req); + res.IoStatus.Status = STATUS_SUCCESS; hasRes = true; break; case CONSOLE_IO_RAW_WRITE: printf("Received raw write request of %lu bytes from process %llu\n", req.Descriptor.InputSize, req.Descriptor.Process); - res.IoStatus.Status = STATUS_NOT_IMPLEMENTED; - hasRes = true; + if (handleRawWrite(req)) + { + res.IoStatus.Status = STATUS_SUCCESS; + hasRes = true; + } break; case CONSOLE_IO_RAW_READ: printf("Received raw read request of %lu bytes from process %llu\n", req.Descriptor.OutputSize, req.Descriptor.Process); - res.IoStatus.Status = STATUS_NOT_IMPLEMENTED; - hasRes = true; + if (handleRawRead(req)) + { + res.IoStatus.Status = STATUS_SUCCESS; + hasRes = true; + } break; case CONSOLE_IO_USER_DEFINED: printf("Received user defined IO request: %lu\n", req.Descriptor.InputSize); diff --git a/src/conpty/PtyServer.h b/src/conpty/PtyServer.h index 40df9078da..845c507248 100644 --- a/src/conpty/PtyServer.h +++ b/src/conpty/PtyServer.h @@ -2,6 +2,7 @@ #include +#include #include #include @@ -33,6 +34,22 @@ struct CONSOLE_API_MSG }; }; +// Handle type flags, from the OG server.h. +// These are internal to the console server, not part of the condrv protocol. +#define CONSOLE_INPUT_HANDLE 0x00000001 +#define CONSOLE_OUTPUT_HANDLE 0x00000002 + +// Handle tracking data, analogous to the OG CONSOLE_HANDLE_DATA. +// In the OG, handles are raw pointers to CONSOLE_HANDLE_DATA which contain +// share mode/access tracking and a pointer to the underlying object +// (INPUT_INFORMATION or SCREEN_INFORMATION). We simplify this for now. +struct PtyHandle +{ + ULONG handleType = 0; // CONSOLE_INPUT_HANDLE or CONSOLE_OUTPUT_HANDLE + ACCESS_MASK access = 0; + ULONG shareMode = 0; +}; + // Per-client tracking data, analogous to the OG CONSOLE_PROCESS_HANDLE. struct PtyClient { @@ -43,6 +60,19 @@ struct PtyClient ULONG_PTR outputHandle = 0; }; +// A pending IO request that couldn't be completed immediately. +// For writes: the output is paused (e.g. the user hit Pause). +// For reads: the input queue is empty; we'll complete it when data arrives. +struct PendingIO +{ + LUID identifier{}; // ConDrv message identifier, needed for completeIo/writeOutput. + ULONG_PTR process = 0; // Descriptor.Process, for cleanup on disconnect. + ULONG_PTR object = 0; // Descriptor.Object, the handle this IO targets. + ULONG function = 0; // CONSOLE_IO_RAW_READ or CONSOLE_IO_RAW_WRITE. + ULONG outputSize = 0; // For reads: max bytes the client accepts. + std::vector inputData; // For writes: the data the client sent. +}; + struct PtyServer : IPtyServer { PtyServer(); @@ -79,20 +109,38 @@ private: // ConDrv communication helpers. void readInput(const CD_IO_DESCRIPTOR& desc, ULONG offset, void* buffer, ULONG size); + void writeOutput(const CD_IO_DESCRIPTOR& desc, ULONG offset, const void* buffer, ULONG size); void completeIo(CD_IO_COMPLETE& completion); // Message handlers (implemented in PtyServer.clients.cpp). + // Handlers returning bool: true = reply pending (don't reply inline), false = reply inline. void handleConnect(CONSOLE_API_MSG& msg); void handleDisconnect(CONSOLE_API_MSG& msg); + void handleCreateObject(CONSOLE_API_MSG& msg); + void handleCloseObject(CONSOLE_API_MSG& msg); + bool handleRawWrite(CONSOLE_API_MSG& msg); + bool handleRawRead(CONSOLE_API_MSG& msg); + + // Complete pending IOs (called when state changes make progress possible). + void completePendingRead(const void* data, ULONG size); + void completePendingWrites(); + void cancelPendingIOs(ULONG_PTR process); // Client lookup by opaque handle value (the raw PtyClient pointer cast to ULONG_PTR). PtyClient* findClient(ULONG_PTR handle); + // Handle management. + ULONG_PTR allocateHandle(ULONG handleType, ACCESS_MASK access, ULONG shareMode); + void freeHandle(ULONG_PTR handle); + std::atomic m_refCount{ 1 }; unique_nthandle m_server; wil::unique_event m_inputAvailableEvent; bool m_initialized = false; - ULONG_PTR m_nextHandleId = 1; + bool m_outputPaused = false; std::vector> m_clients; + std::vector> m_handles; + std::deque m_pendingReads; + std::deque m_pendingWrites; }; diff --git a/src/conpty/PtyServer.io.cpp b/src/conpty/PtyServer.io.cpp new file mode 100644 index 0000000000..ff193fcfd0 --- /dev/null +++ b/src/conpty/PtyServer.io.cpp @@ -0,0 +1,183 @@ +#include "pch.h" +#include "PtyServer.h" + +#include + +// Handles CONSOLE_IO_RAW_WRITE messages. +// +// Protocol (from OG ConsoleIoThread RAW_WRITE case + SrvWriteConsole in stream.cpp): +// 1. The client's write data is available via readInput (IOCTL_CONDRV_READ_INPUT). +// Descriptor.InputSize tells us how many bytes are available. +// 2. In the OG, SrvWriteConsole calls GetInputBuffer() to pull all client data, +// then DoSrvWriteConsole() feeds it to the output renderer. +// 3. If output is paused (e.g. Pause key), the write is deferred: we save the +// message and complete it later when output is resumed. This is the +// ReplyPending mechanism from the OG. +// +// Returns true if the reply is pending (caller must NOT reply inline). +// Returns false if the write completed immediately (caller replies inline). +bool PtyServer::handleRawWrite(CONSOLE_API_MSG& msg) +{ + const auto size = msg.Descriptor.InputSize; + + // Read the client's write payload from the driver upfront, regardless + // of whether we can process it now. The driver expects us to consume it. + std::vector buffer(size); + if (size > 0) + { + readInput(msg.Descriptor, 0, buffer.data(), size); + } + + // If output is paused, defer this write. The OG creates a wait block + // that gets signaled when output is resumed (ConsoleNotifyWait). + if (m_outputPaused) + { + PendingIO pending; + pending.identifier = msg.Descriptor.Identifier; + pending.process = msg.Descriptor.Process; + pending.object = msg.Descriptor.Object; + pending.function = CONSOLE_IO_RAW_WRITE; + pending.inputData = std::move(buffer); + m_pendingWrites.push_back(std::move(pending)); + return false; // reply pending + } + + printf(" %*s\r\n", static_cast(buffer.size()), reinterpret_cast(buffer.data())); + return true; // reply immediately +} + +// Handles CONSOLE_IO_RAW_READ messages. +// +// Protocol (from OG ConsoleIoThread RAW_READ case + SrvReadConsole in stream.cpp): +// 1. Descriptor.OutputSize tells us the max bytes the client wants to read. +// 2. In the OG, SrvReadConsole calls GetAugmentedOutputBuffer() to allocate a +// server-side buffer, then ReadChars() fills it with input data. +// 3. If the input queue is empty, the OG returns CONSOLE_STATUS_WAIT and +// creates a wait block on the input buffer's ReadWaitQueue. When input +// arrives, ConsoleNotifyWait completes the pending read. +// 4. On success, ReleaseMessageBuffers() writes the output buffer back via +// IOCTL_CONDRV_WRITE_OUTPUT, and IoStatus.Information is set to bytes read. +// +// Returns true if the reply is pending (caller must NOT reply inline). +// Returns false if the read completed immediately. +bool PtyServer::handleRawRead(CONSOLE_API_MSG& msg) +{ + const auto maxBytes = msg.Descriptor.OutputSize; + + // TODO: Try to read data from the input queue. + // For now, we always pend — there's no input source yet. + // When input data becomes available, call completePendingRead(). + + PendingIO pending; + pending.identifier = msg.Descriptor.Identifier; + pending.process = msg.Descriptor.Process; + pending.object = msg.Descriptor.Object; + pending.function = CONSOLE_IO_RAW_READ; + pending.outputSize = maxBytes; + m_pendingReads.push_back(std::move(pending)); + return false; // reply pending +} + +// Completes the oldest pending read with the given data. +// Called when input data becomes available (e.g. from the terminal's input pipeline). +// +// In the OG, this is analogous to ConsoleNotifyWait on the ReadWaitQueue, +// which re-invokes the read routine and, on success, calls ReleaseMessageBuffers +// to write the output data back to the client via IOCTL_CONDRV_WRITE_OUTPUT. +void PtyServer::completePendingRead(const void* data, ULONG size) +{ + if (m_pendingReads.empty()) + { + return; + } + + auto pending = std::move(m_pendingReads.front()); + m_pendingReads.pop_front(); + + auto bytesToWrite = std::min(size, pending.outputSize); + + // Write the data back to the client's read buffer. + if (bytesToWrite > 0) + { + CD_IO_OPERATION op{}; + op.Identifier = pending.identifier; + op.Buffer.Offset = 0; + op.Buffer.Data = const_cast(data); + op.Buffer.Size = bytesToWrite; + THROW_IF_NTSTATUS_FAILED(ioctl(IOCTL_CONDRV_WRITE_OUTPUT, &op, sizeof(op), nullptr, 0)); + } + + // Complete the read with the number of bytes returned. + CD_IO_COMPLETE completion{}; + completion.Identifier = pending.identifier; + completion.IoStatus.Status = STATUS_SUCCESS; + completion.IoStatus.Information = bytesToWrite; + + completeIo(completion); +} + +// Completes all pending writes (called when output is unpaused). +// +// In the OG, this is analogous to ConsoleNotifyWait on the OutputQueue, +// which re-invokes DoSrvWriteConsole for each deferred write. +void PtyServer::completePendingWrites() +{ + while (!m_pendingWrites.empty()) + { + auto pending = std::move(m_pendingWrites.front()); + m_pendingWrites.pop_front(); + + // TODO: Feed pending.inputData to the output/rendering pipeline. + + // Complete the write. + CD_IO_COMPLETE completion{}; + completion.Identifier = pending.identifier; + completion.IoStatus.Status = STATUS_SUCCESS; + completion.IoStatus.Information = 0; + + completeIo(completion); + } +} + +// Cancels all pending IOs for a disconnecting process. +// +// In the OG, FreeProcessData iterates the WaitBlockQueue and calls +// ConsoleNotifyWaitBlock with fThreadDying=TRUE for each pending wait. +// The wait routines then complete the IO with an error status. +void PtyServer::cancelPendingIOs(ULONG_PTR process) +{ + // Cancel pending reads from this process. + while (true) + { + auto it = std::find_if(m_pendingReads.begin(), m_pendingReads.end(), + [process](const PendingIO& p) { return p.process == process; }); + if (it == m_pendingReads.end()) + break; + + CD_IO_COMPLETE completion{}; + completion.Identifier = it->identifier; + completion.IoStatus.Status = STATUS_CANCELLED; + completion.IoStatus.Information = 0; + + // Best-effort: if the process is gone, the driver may reject this. + ioctl(IOCTL_CONDRV_COMPLETE_IO, &completion, sizeof(completion), nullptr, 0); + m_pendingReads.erase(it); + } + + // Cancel pending writes from this process. + while (true) + { + auto it = std::find_if(m_pendingWrites.begin(), m_pendingWrites.end(), + [process](const PendingIO& p) { return p.process == process; }); + if (it == m_pendingWrites.end()) + break; + + CD_IO_COMPLETE completion{}; + completion.Identifier = it->identifier; + completion.IoStatus.Status = STATUS_CANCELLED; + completion.IoStatus.Information = 0; + + ioctl(IOCTL_CONDRV_COMPLETE_IO, &completion, sizeof(completion), nullptr, 0); + m_pendingWrites.erase(it); + } +} diff --git a/src/conpty/conpty.vcxproj b/src/conpty/conpty.vcxproj index 19ae9afe64..8b0b3ab9c2 100644 --- a/src/conpty/conpty.vcxproj +++ b/src/conpty/conpty.vcxproj @@ -27,6 +27,7 @@ + diff --git a/src/conpty/conpty.vcxproj.filters b/src/conpty/conpty.vcxproj.filters index 2a40eec8e5..fb964ed601 100644 --- a/src/conpty/conpty.vcxproj.filters +++ b/src/conpty/conpty.vcxproj.filters @@ -29,6 +29,9 @@ Source Files + + Source Files + diff --git a/src/conpty/pch.h b/src/conpty/pch.h index ffc91c5816..d1b5aadd16 100644 --- a/src/conpty/pch.h +++ b/src/conpty/pch.h @@ -1,8 +1,8 @@ #pragma once -#define WIN32_LEAN_AND_MEAN -#define UMDF_USING_NTSTATUS #define NOMINMAX +#define UMDF_USING_NTSTATUS +#define WIN32_LEAN_AND_MEAN #include #include