HANDLE CNdasEventSubscriber:: WaitServer(BOOL& bStopThread) { OVERLAPPED ov = {0}; ov.hEvent = m_hDataEvent; HANDLE hPipe = CreatePipeConnection(&ov); if (INVALID_HANDLE_VALUE == hPipe) { NDAS_EVENT_INFO ei = {0}; ei.EventType = NDAS_EVENT_TYPE_CONNECTION_FAILED; CallEventProc(::GetLastError(), NULL); } // // Wait for 5 seconds and try again // while (INVALID_HANDLE_VALUE == hPipe) { { NDAS_EVENT_INFO ei = {0}; ei.EventType = NDAS_EVENT_TYPE_CONNECTION_FAILED; CallEventProc(::GetLastError(), &ei); } DWORD dwWaitResult = ::WaitForSingleObject( m_hThreadStopEvent, 5 * 1000); if (WAIT_OBJECT_0 == dwWaitResult) { bStopThread = TRUE; return INVALID_HANDLE_VALUE; } else if (WAIT_TIMEOUT == dwWaitResult) { { NDAS_EVENT_INFO ei = {0}; ei.EventType = NDAS_EVENT_TYPE_RETRYING_CONNECTION; CallEventProc(::GetLastError(), &ei); } hPipe = CreatePipeConnection(&ov); } else { DPErrorEx(_FT("Wait failed: \n")); return INVALID_HANDLE_VALUE; } } { NDAS_EVENT_INFO ei = {0}; ei.EventType = NDAS_EVENT_TYPE_CONNECTED; CallEventProc(0, &ei); } return hPipe; }
DWORD WINAPI WorkerThread(LPVOID pvVoid) { PCLIENT pClient; BOOL b; LPOVERLAPPED lpo; DWORD nBytes; HANDLE hCompletionPort = (HANDLE)pvVoid; for (BOOL fKeepLooping = TRUE; fKeepLooping;) { pClient = NULL; lpo = NULL; nBytes = 0; b = GetQueuedCompletionStatus(hCompletionPort, &nBytes, (PULONG_PTR)&pClient, &lpo, INFINITE); if (!b) { if (pClient) { if (GetLastError() == ERROR_BROKEN_PIPE) { pClient->LogMessageV("<!-- Client closed pipe. -->"); } else { pClient->LogMessageV("<!-- *** GetQueuedCompletionStatus failed %d -->", GetLastError()); } CloseConnection(pClient); } continue; } if (pClient->fAwaitingAccept) { BOOL fAgain = TRUE; while (fAgain) { LONG nClient = InterlockedIncrement(&s_nTotalClients); InterlockedIncrement(&s_nActiveClients); pClient->fAwaitingAccept = FALSE; PCLIENT pNew = CreatePipeConnection(hCompletionPort, nClient); fAgain = FALSE; if (pNew != NULL) { fAgain = !pNew->fAwaitingAccept; DoRead(pNew); } } } else { if (nBytes <= offsetof(TBLOG_MESSAGE, szMessage)) { pClient->LogMessageV("</t:Process>\n"); CloseConnection(pClient); continue; } pClient->LogMessage(&pClient->Message, nBytes); } DoRead(pClient); } return 0; }
DWORD main(int argc, char **argv) { HANDLE hCompletionPort; BOOL fNeedHelp = FALSE; WCHAR wzzDrop[1024] = L"build\0nmake\0"; GetSystemTimeAsFileTime((FILETIME *)&s_llStartTime); StringCchPrintfA(s_szPipe, ARRAYSIZE(s_szPipe), "%s.%d", TBLOG_PIPE_NAME, GetCurrentProcessId()); int arg = 1; for (; arg < argc && (argv[arg][0] == '-' || argv[arg][0] == '/'); arg++) { CHAR *argn = argv[arg] + 1; CHAR *argp = argn; while (*argp && *argp != ':' && *argp != '=') { argp++; } if (*argp == ':' || *argp == '=') { *argp++ = '\0'; } switch (argn[0]) { case 'd': // Drop Processes case 'D': if (*argp) { PWCHAR pwz = wzzDrop; while (*argp) { if (*argp == ';') { *pwz++ = '\0'; } else { *pwz++ = *argp++; } } *pwz++ = '\0'; *pwz = '\0'; } case 'o': // Output file. case 'O': StringCchCopyA(s_szLogFile, ARRAYSIZE(s_szLogFile), argp); break; case 'v': // Verbose case 'V': s_fVerbose = TRUE; break; case '?': // Help. fNeedHelp = TRUE; break; default: fNeedHelp = TRUE; printf("TRACEBLD: Bad argument: %s:%s\n", argn, argp); break; } } if (arg >= argc) { fNeedHelp = TRUE; } if (fNeedHelp) { printf("Usage:\n" " tracebld [options] command {command arguments}\n" "Options:\n" " /o:file Log all events to the output files.\n" " /? Display this help message.\n" "Summary:\n" " Runs the build commands and figures out which files have dependencies..\n" "\n"); exit(9001); } // Create the completion port. hCompletionPort = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, NULL, 0); if (hCompletionPort == NULL) { MyErrExit("CreateIoCompletionPort"); } // Create completion port worker threads. // CreateWorkers(hCompletionPort); CreatePipeConnection(hCompletionPort, 0); printf("TRACEBLD: Ready for clients. Press Ctrl-C to stop.\n"); /////////////////////////////////////////////////////////// Validate DLLs. // CHAR szTmpPath[MAX_PATH]; CHAR szExePath[MAX_PATH]; CHAR szDllPath[MAX_PATH]; PCHAR pszFilePart = NULL; if (!GetModuleFileNameA(NULL, szTmpPath, ARRAYSIZE(szTmpPath))) { printf("TRACEBLD: Couldn't retreive exe name.\n"); return 9002; } if (!GetFullPathNameA(szTmpPath, ARRAYSIZE(szExePath), szExePath, &pszFilePart) || pszFilePart == NULL) { printf("TRACEBLD: Error: %s is not a valid path name..\n", szTmpPath); return 9002; } StringCchCopyA(pszFilePart, szExePath + ARRAYSIZE(szExePath) - pszFilePart, "trcbld" DETOURS_STRINGIFY(DETOURS_BITS) ".dll"); StringCchCopyA(szDllPath, ARRAYSIZE(szDllPath), szExePath); ////////////////////////////////////////////////////////////////////////// STARTUPINFOA si; PROCESS_INFORMATION pi; CHAR szCommand[2048]; CHAR szExe[MAX_PATH]; CHAR szFullExe[MAX_PATH] = "\0"; PCHAR pszFileExe = NULL; ZeroMemory(&si, sizeof(si)); ZeroMemory(&pi, sizeof(pi)); si.cb = sizeof(si); szCommand[0] = L'\0'; StringCchCopyA(szExe, sizeof(szExe), argv[arg]); for (; arg < argc; arg++) { if (strchr(argv[arg], ' ') != NULL || strchr(argv[arg], '\t') != NULL) { StringCchCatA(szCommand, sizeof(szCommand), "\""); StringCchCatA(szCommand, sizeof(szCommand), argv[arg]); StringCchCatA(szCommand, sizeof(szCommand), "\""); } else { StringCchCatA(szCommand, sizeof(szCommand), argv[arg]); } if (arg + 1 < argc) { StringCchCatA(szCommand, sizeof(szCommand), " "); } } printf("TRACEBLD: Starting: `%s'\n", szCommand); printf("TRACEBLD: with `%s'\n", szDllPath); fflush(stdout); DWORD dwFlags = CREATE_DEFAULT_ERROR_MODE | CREATE_SUSPENDED; SetLastError(0); SearchPathA(NULL, szExe, ".exe", ARRAYSIZE(szFullExe), szFullExe, &pszFileExe); if (!DetourCreateProcessWithDllExA(szFullExe[0] ? szFullExe : NULL, szCommand, NULL, NULL, TRUE, dwFlags, NULL, NULL, &si, &pi, szDllPath, NULL)) { printf("TRACEBLD: DetourCreateProcessWithDllEx failed: %d\n", GetLastError()); ExitProcess(9007); } ZeroMemory(&s_Payload, sizeof(s_Payload)); s_Payload.nParentProcessId = GetCurrentProcessId(); s_Payload.nTraceProcessId = GetCurrentProcessId(); s_Payload.nGeneology = 1; s_Payload.rGeneology[0] = 0; StringCchCopyW(s_Payload.wzStdin, ARRAYSIZE(s_Payload.wzStdin), L"\\\\.\\CONIN$"); StringCchCopyW(s_Payload.wzStdout, ARRAYSIZE(s_Payload.wzStdout), L"\\\\.\\CONOUT$"); StringCchCopyW(s_Payload.wzStderr, ARRAYSIZE(s_Payload.wzStderr), L"\\\\.\\CONOUT$"); StringCchCopyW(s_Payload.wzParents, ARRAYSIZE(s_Payload.wzParents), L""); CopyEnvironment(s_Payload.wzzDrop, wzzDrop); LPWCH pwStrings = GetEnvironmentStringsW(); CopyEnvironment(s_Payload.wzzEnvironment, pwStrings); FreeEnvironmentStringsW(pwStrings); if (!DetourCopyPayloadToProcess(pi.hProcess, s_guidTrace, &s_Payload, sizeof(s_Payload))) { printf("TRACEBLD: DetourCopyPayloadToProcess failed: %d\n", GetLastError()); ExitProcess(9008); } ResumeThread(pi.hThread); WaitForSingleObject(pi.hProcess, INFINITE); DWORD dwResult = 0; if (!GetExitCodeProcess(pi.hProcess, &dwResult)) { printf("TRACEBLD: GetExitCodeProcess failed: %d\n", GetLastError()); return 9008; } printf("TRACEBLD: %d processes.\n", s_nTotalClients); return dwResult; }