wip

src/conpty/PtyServer.clients.cpp

@@ -1,14 +1,12 @@
 #include "pch.h"
 #include "PtyServer.h"
 
-#include <algorithm>
-
 // 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<void*>(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<PtyClient*>(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<PtyHandle>();
+    h->handleType = handleType;
+    h->access = access;
+    h->shareMode = shareMode;
+    auto ptr = reinterpret_cast<ULONG_PTR>(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<PtyHandle*>(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);
+}

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);

src/conpty/PtyServer.h

@@ -2,6 +2,7 @@
 
 #include <conpty.h>
 
+#include <deque>
 #include <memory>
 #include <vector>
 
@@ -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<uint8_t> 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<ULONG> 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<std::unique_ptr<PtyClient>> m_clients;
+    std::vector<std::unique_ptr<PtyHandle>> m_handles;
+    std::deque<PendingIO> m_pendingReads;
+    std::deque<PendingIO> m_pendingWrites;
 };

src/conpty/PtyServer.io.cpp

@@ -0,0 +1,183 @@
+#include "pch.h"
+#include "PtyServer.h"
+
+#include <algorithm>
+
+// 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<uint8_t> 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<int>(buffer.size()), reinterpret_cast<const char*>(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<void*>(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);
+    }
+}

src/conpty/conpty.vcxproj

@@ -27,6 +27,7 @@
     </ClCompile>
     <ClCompile Include="PtyServer.cpp" />
     <ClCompile Include="PtyServer.clients.cpp" />
+    <ClCompile Include="PtyServer.io.cpp" />
   </ItemGroup>
   <ItemGroup>
     <Midl Include="conpty.idl">

src/conpty/conpty.vcxproj.filters

@@ -29,6 +29,9 @@
     <ClCompile Include="pch.cpp">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="PtyServer.io.cpp">
+      <Filter>Source Files</Filter>
+    </ClCompile>
   </ItemGroup>
   <ItemGroup>
     <Midl Include="conpty.idl">

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 <ntstatus.h>
 #include <Windows.h>