fix: Финальный рефакторинг

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

client.h

@@ -1,3 +1,5 @@
+#pragma once
+
 #ifndef WIN32_LEAN_AND_MEAN
 #define WIN32_LEAN_AND_MEAN
 #endif

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.

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))
+    if (!CreatePipe(&g_hChildStd_OUT_Rd, &g_hChildStd_OUT_Wr, &saAttr, 0)) {
-        ErrorExit(TEXT("StdoutRd CreatePipe"));
+        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))
+    if (!SetHandleInformation(g_hChildStd_OUT_Rd, HANDLE_FLAG_INHERIT, 0)) {
-        ErrorExit(TEXT("Stdout SetHandleInformation"));
+        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))
+    if (!CreatePipe(&g_hChildStd_IN_Rd, &g_hChildStd_IN_Wr, &saAttr, 0)) {
-        ErrorExit(TEXT("Stdin CreatePipe"));
+        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))
+    if (!SetHandleInformation(g_hChildStd_IN_Wr, HANDLE_FLAG_INHERIT, 0)) {
-        ErrorExit(TEXT("Stdin SetHandleInformation"));
+        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, WriteToPipe, &WriteThreadInfo, 0, NULL);
-    PipeThreads[1] = CreateThread(NULL, 0, ReadFromPipe, &ReadThreadInfo, 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(2, PipeThreads, TRUE, INFINITE);
+    WaitForMultipleObjects(3, PipeThreads, FALSE, 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)
+    if (!bSuccess) {
-        ErrorExit(TEXT("CreateProcess"));
+        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)) {
-    if (!CloseHandle(g_hChildStd_IN_Wr))
+        printf("StdInWr CloseHandle\n");
-        ErrorExit(TEXT("StdInWr CloseHandle"));
+        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);
-}

server.h

@@ -1,3 +1,5 @@
+#pragma once
+
 #ifndef WIN32_LEAN_AND_MEAN
 #define WIN32_LEAN_AND_MEAN
 #endif

service.h

@@ -1,3 +1,5 @@
+#pragma once
+
 #include <windows.h>
 #include <tchar.h>
 #include <strsafe.h>