nihonium/mipt_2024_win_shell_server
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

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

Browse source
This commit is contained in:
Nihonium 2024-12-13 04:11:18 +03:00
parent 72bf8d917a
commit 934ce494d6
6 changed files with 73 additions and 71 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
client.c
View file

@ -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, TRUE, INFINITE);
WaitForMultipleObjects(2, threads, FALSE, 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) {

2
client.h
View file

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

6
main.c
View file

@ -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.

111
server.c
View file

@ -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);
WaitForMultipleObjects(2, PipeThreads, TRUE, INFINITE);
if (!(PipeThreads[1] && PipeThreads[2])) {
printf("Failed to create pipe threads\n");
return 1;
}
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)
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);
}

2
server.h
View file

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

2
service.h
View file

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