diff --git a/terminkey.h b/terminkey.h index 81e35a4..99945a2 100644 --- a/terminkey.h +++ b/terminkey.h @@ -112,19 +112,32 @@ int ipc_check_message(uintptr_t buffer_ptr, int max_len) { char* buffer = (char*)buffer_ptr; DWORD bytesRead = 0; + + // Check if a client is connected (or already connected) BOOL connected = ConnectNamedPipe(hServerPipe, NULL) ? TRUE : (GetLastError() == ERROR_PIPE_CONNECTED); if (connected) { - BOOL success = ReadFile(hServerPipe, buffer, max_len - 1, &bytesRead, NULL); - if (success && bytesRead > 0) { - buffer[bytesRead] = '\0'; - DisconnectNamedPipe(hServerPipe); - ConnectNamedPipe(hServerPipe, NULL); - return (int)bytesRead; + DWORD bytesAvail = 0; + // Sneak peek to see if data has safely arrived in the buffer yet + if (PeekNamedPipe(hServerPipe, NULL, 0, NULL, &bytesAvail, NULL)) { + if (bytesAvail > 0) { + // Data is ready, perform the read safely + BOOL success = ReadFile(hServerPipe, buffer, max_len - 1, &bytesRead, NULL); + if (success && bytesRead > 0) { + buffer[bytesRead] = '\0'; + DisconnectNamedPipe(hServerPipe); + ConnectNamedPipe(hServerPipe, NULL); // Re-arm for next connection + return (int)bytesRead; + } + } + } else { + // If Peek fails because the client dropped out without sending anything + if (GetLastError() == ERROR_BROKEN_PIPE) { + DisconnectNamedPipe(hServerPipe); + ConnectNamedPipe(hServerPipe, NULL); + } } - DisconnectNamedPipe(hServerPipe); - ConnectNamedPipe(hServerPipe, NULL); } return 0; }