From patchwork Tue Jun 29 12:16:28 2010 Content-Type: text/plain; charset="utf-8" MIME-Version: 1.0 Content-Transfer-Encoding: 7bit X-Patchwork-Submitter: Michael Goldish X-Patchwork-Id: 108583 Received: from vger.kernel.org (vger.kernel.org [209.132.180.67]) by demeter.kernel.org (8.14.4/8.14.3) with ESMTP id o5TDIm4H012137 for ; Tue, 29 Jun 2010 13:20:56 GMT Received: (majordomo@vger.kernel.org) by vger.kernel.org via listexpand id S1751625Ab0F2MQ7 (ORCPT ); Tue, 29 Jun 2010 08:16:59 -0400 Received: from mx1.redhat.com ([209.132.183.28]:32362 "EHLO mx1.redhat.com" rhost-flags-OK-OK-OK-OK) by vger.kernel.org with ESMTP id S1751370Ab0F2MQ4 (ORCPT ); Tue, 29 Jun 2010 08:16:56 -0400 Received: from int-mx02.intmail.prod.int.phx2.redhat.com (int-mx02.intmail.prod.int.phx2.redhat.com [10.5.11.12]) by mx1.redhat.com (8.13.8/8.13.8) with ESMTP id o5TCGrgH017230 (version=TLSv1/SSLv3 cipher=DHE-RSA-AES256-SHA bits=256 verify=OK); Tue, 29 Jun 2010 08:16:53 -0400 Received: from ns3.rdu.redhat.com (ns3.rdu.redhat.com [10.11.255.199]) by int-mx02.intmail.prod.int.phx2.redhat.com (8.13.8/8.13.8) with ESMTP id o5TCGqfV020864; Tue, 29 Jun 2010 08:16:53 -0400 Received: from localhost.localdomain (dhcp-1-188.tlv.redhat.com [10.35.1.188]) by ns3.rdu.redhat.com (8.13.8/8.13.8) with ESMTP id o5TCGo7r004809; Tue, 29 Jun 2010 08:16:50 -0400 From: Michael Goldish To: autotest@test.kernel.org, kvm@vger.kernel.org Cc: Michael Goldish Subject: [KVM-AUTOTEST PATCH v3] [RFC] KVM test: rss.cpp: add file transfer support Date: Tue, 29 Jun 2010 15:16:28 +0300 Message-Id: <1277813789-13530-1-git-send-email-mgoldish@redhat.com> X-Scanned-By: MIMEDefang 2.67 on 10.5.11.12 Sender: kvm-owner@vger.kernel.org Precedence: bulk List-ID: X-Mailing-List: kvm@vger.kernel.org X-Greylist: IP, sender and recipient auto-whitelisted, not delayed by milter-greylist-4.2.3 (demeter.kernel.org [140.211.167.41]); Tue, 29 Jun 2010 13:20:56 +0000 (UTC) diff --git a/client/tests/kvm/deps/rss.cpp b/client/tests/kvm/deps/rss.cpp index 66d9a5b..b95a45d 100644 --- a/client/tests/kvm/deps/rss.cpp +++ b/client/tests/kvm/deps/rss.cpp @@ -1,459 +1,989 @@ -// Simple remote shell server -// Author: Michael Goldish -// Much of the code here was adapted from Microsoft code samples. - -// Usage: rss.exe [port] -// If no port is specified the default is 22. - -#define _WIN32_WINNT 0x0500 - -#include -#include -#include - -#pragma comment(lib, "ws2_32.lib") - -int port = 22; - -HWND hMainWindow = NULL; -HWND hTextBox = NULL; - -struct client_info { - SOCKET socket; - sockaddr_in addr; - int pid; - HWND hwnd; - HANDLE hJob; - HANDLE hChildOutputRead; - HANDLE hThreadChildToSocket; -}; - -void ExitOnError(char *message, BOOL winsock = 0) -{ - LPVOID system_message; - char buffer[512]; - - int error_code; - if (winsock) - error_code = WSAGetLastError(); - else - error_code = GetLastError(); - - WSACleanup(); - - FormatMessage( - FORMAT_MESSAGE_ALLOCATE_BUFFER|FORMAT_MESSAGE_FROM_SYSTEM, - NULL, error_code, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), - (LPTSTR)&system_message, 0, NULL); - - sprintf(buffer, - "%s!\n" - "Error code = %d\n" - "Error message = %s", - message, error_code, (char *)system_message); - - MessageBox(hMainWindow, buffer, "Error", MB_OK | MB_ICONERROR); - - LocalFree(system_message); - ExitProcess(1); -} - -void AppendMessage(char *message) -{ - int length = GetWindowTextLength(hTextBox); - SendMessage(hTextBox, EM_SETSEL, (WPARAM)length, (LPARAM)length); - SendMessage(hTextBox, EM_REPLACESEL, (WPARAM)FALSE, (LPARAM)message); -} - -void FormatStringForPrinting(char *dst, char *src, int size) -{ - int j = 0; - - for (int i = 0; i < size && src[i]; i++) { - if (src[i] == '\n') { - dst[j++] = '\\'; - dst[j++] = 'n'; - } else if (src[i] == '\r') { - dst[j++] = '\\'; - dst[j++] = 'r'; - } else if (src[i] == '\t') { - dst[j++] = '\\'; - dst[j++] = 't'; - } else if (src[i] == '\\') { - dst[j++] = '\\'; - dst[j++] = '\\'; - } else dst[j++] = src[i]; - } - dst[j] = 0; -} - -char* GetClientIPAddress(client_info *ci) -{ - char *address = inet_ntoa(ci->addr.sin_addr); - if (address) - return address; - else - return "unknown"; -} - -DWORD WINAPI ChildToSocket(LPVOID client_info_ptr) -{ - char buffer[1024], message[1024]; - client_info ci; - DWORD bytes_read; - int bytes_sent; - - memcpy(&ci, client_info_ptr, sizeof(ci)); - - while (1) { - // Read data from the child's STDOUT/STDERR pipes - if (!ReadFile(ci.hChildOutputRead, - buffer, sizeof(buffer), - &bytes_read, NULL) || !bytes_read) { - if (GetLastError() == ERROR_BROKEN_PIPE) - break; // Pipe done -- normal exit path - else - ExitOnError("ReadFile failed"); // Something bad happened - } - // Send data to the client - bytes_sent = send(ci.socket, buffer, bytes_read, 0); - /* - // Make sure all the data was sent - if (bytes_sent != bytes_read) { - sprintf(message, - "ChildToSocket: bytes read (%d) != bytes sent (%d)", - bytes_read, bytes_sent); - ExitOnError(message, 1); - } - */ - } - - AppendMessage("Child exited\r\n"); - shutdown(ci.socket, SD_BOTH); - - return 0; -} - -DWORD WINAPI SocketToChild(LPVOID client_info_ptr) -{ - char buffer[256], formatted_buffer[768]; - char message[1024], client_info_str[256]; - client_info ci; - DWORD bytes_written; - int bytes_received; - - memcpy(&ci, client_info_ptr, sizeof(ci)); - - sprintf(client_info_str, "address %s, port %d", - GetClientIPAddress(&ci), ci.addr.sin_port); - - sprintf(message, "New client connected (%s)\r\n", client_info_str); - AppendMessage(message); - - while (1) { - // Receive data from the socket - ZeroMemory(buffer, sizeof(buffer)); - bytes_received = recv(ci.socket, buffer, sizeof(buffer), 0); - if (bytes_received <= 0) - break; - // Report the data received - FormatStringForPrinting(formatted_buffer, buffer, sizeof(buffer)); - sprintf(message, "Client (%s) entered text: \"%s\"\r\n", - client_info_str, formatted_buffer); - AppendMessage(message); - // Send the data as a series of WM_CHAR messages to the console window - for (int i=0; ipid = pi.dwProcessId; - // Assign the process to a newly created JobObject - ci->hJob = CreateJobObject(NULL, NULL); - AssignProcessToJobObject(ci->hJob, pi.hProcess); - // Keep the console window's handle - ci->hwnd = GetConsoleWindow(); - - // Detach from the child's console - FreeConsole(); -} - -void SpawnSession(client_info *ci) -{ - HANDLE hOutputReadTmp, hOutputRead, hOutputWrite; - HANDLE hErrorWrite; - SECURITY_ATTRIBUTES sa; - - // Set up the security attributes struct. - sa.nLength = sizeof(SECURITY_ATTRIBUTES); - sa.lpSecurityDescriptor = NULL; - sa.bInheritHandle = TRUE; - - // Create the child output pipe. - if (!CreatePipe(&hOutputReadTmp, &hOutputWrite, &sa, 0)) - ExitOnError("CreatePipe failed"); - - // Create a duplicate of the output write handle for the std error - // write handle. This is necessary in case the child application - // closes one of its std output handles. - if (!DuplicateHandle(GetCurrentProcess(), hOutputWrite, - GetCurrentProcess(), &hErrorWrite, 0, - TRUE, DUPLICATE_SAME_ACCESS)) - ExitOnError("DuplicateHandle failed"); - - // Create new output read handle and the input write handles. Set - // the Properties to FALSE. Otherwise, the child inherits the - // properties and, as a result, non-closeable handles to the pipes - // are created. - if (!DuplicateHandle(GetCurrentProcess(), hOutputReadTmp, - GetCurrentProcess(), - &hOutputRead, // Address of new handle. - 0, FALSE, // Make it uninheritable. - DUPLICATE_SAME_ACCESS)) - ExitOnError("DuplicateHandle failed"); - - // Close inheritable copies of the handles you do not want to be - // inherited. - if (!CloseHandle(hOutputReadTmp)) - ExitOnError("CloseHandle failed"); - - PrepAndLaunchRedirectedChild(ci, hOutputWrite, hErrorWrite); - - ci->hChildOutputRead = hOutputRead; - - // Close pipe handles (do not continue to modify the parent). - // You need to make sure that no handles to the write end of the - // output pipe are maintained in this process or else the pipe will - // not close when the child process exits and the ReadFile will hang. - if (!CloseHandle(hOutputWrite)) ExitOnError("CloseHandle failed"); - if (!CloseHandle(hErrorWrite)) ExitOnError("CloseHandle failed"); -} - -DWORD WINAPI ListenThread(LPVOID param) -{ - WSADATA wsaData; - SOCKET ListenSocket = INVALID_SOCKET; - sockaddr_in addr; - int result, addrlen; - client_info ci; - HANDLE hThread; - - // Initialize Winsock - result = WSAStartup(MAKEWORD(2,2), &wsaData); - if (result) - ExitOnError("Winsock initialization failed"); - - // Create socket - ListenSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP); - if (ListenSocket == INVALID_SOCKET) - ExitOnError("Socket creation failed", 1); - - // Bind the socket - addr.sin_family = AF_INET; - addr.sin_addr.s_addr = htonl(INADDR_ANY); - addr.sin_port = htons(port); - - result = bind(ListenSocket, (sockaddr *)&addr, sizeof(addr)); - if (result == SOCKET_ERROR) - ExitOnError("bind failed", 1); - - // Start listening for incoming connections - result = listen(ListenSocket, SOMAXCONN); - if (result == SOCKET_ERROR) - ExitOnError("listen failed", 1); - - // Inform the user - AppendMessage("Waiting for clients to connect...\r\n"); - - while (1) { - addrlen = sizeof(ci.addr); - ci.socket = accept(ListenSocket, (sockaddr *)&ci.addr, &addrlen); - if (ci.socket == INVALID_SOCKET) { - if (WSAGetLastError() == WSAEINTR) - break; - else - ExitOnError("accept failed", 1); - } - - // Under heavy load, spawning cmd.exe might take a while, so tell the - // client to be patient - char *message = "Please wait...\r\n"; - send(ci.socket, message, strlen(message), 0); - // Spawn a new redirected cmd.exe process - SpawnSession(&ci); - // Start transferring data from the child process to the client - hThread = CreateThread(NULL, 0, ChildToSocket, (LPVOID)&ci, 0, NULL); - if (!hThread) - ExitOnError("Could not create ChildToSocket thread"); - ci.hThreadChildToSocket = hThread; - // ... and from the client to the child process - hThread = CreateThread(NULL, 0, SocketToChild, (LPVOID)&ci, 0, NULL); - if (!hThread) - ExitOnError("Could not create SocketToChild thread"); - } - - return 0; -} - -LRESULT CALLBACK WndProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam) -{ - RECT rect; - HANDLE hListenThread; - - switch (msg) { - case WM_CREATE: - // Create text box - GetClientRect(hwnd, &rect); - hTextBox = CreateWindowEx(WS_EX_CLIENTEDGE, - "EDIT", "", - WS_CHILD|WS_VISIBLE|WS_VSCROLL| - ES_MULTILINE|ES_AUTOVSCROLL, - 20, 20, - rect.right - 40, - rect.bottom - 40, - hwnd, - NULL, - GetModuleHandle(NULL), - NULL); - if (!hTextBox) - ExitOnError("Could not create text box"); - - // Set the font - SendMessage(hTextBox, WM_SETFONT, - (WPARAM)GetStockObject(DEFAULT_GUI_FONT), - MAKELPARAM(FALSE, 0)); - - // Start the listening thread - hListenThread = - CreateThread(NULL, 0, ListenThread, NULL, 0, NULL); - if (!hListenThread) - ExitOnError("Could not create server thread"); - break; - - case WM_DESTROY: - WSACleanup(); - PostQuitMessage(0); - break; - - default: - return DefWindowProc(hwnd, msg, wParam, lParam); - } - - return 0; -} - -int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, - LPSTR lpCmdLine, int nShowCmd) -{ - WNDCLASSEX wc; - MSG msg; - - if (strlen(lpCmdLine)) - sscanf(lpCmdLine, "%d", &port); - - // Make sure the firewall is disabled - system("netsh firewall set opmode disable"); - - // Create the window class - wc.cbSize = sizeof(WNDCLASSEX); - wc.style = CS_HREDRAW | CS_VREDRAW; - wc.lpfnWndProc = WndProc; - wc.cbClsExtra = 0; - wc.cbWndExtra = 0; - wc.hInstance = hInstance; - wc.hIcon = LoadIcon(NULL, IDI_APPLICATION); - wc.hIconSm = LoadIcon(NULL, IDI_APPLICATION); - wc.hbrBackground = (HBRUSH)(COLOR_BTNFACE + 1); - wc.lpszMenuName = NULL; - wc.lpszClassName = "RemoteShellServerWindowClass"; - wc.hCursor = LoadCursor(NULL, IDC_ARROW); - - if (!RegisterClassEx(&wc)) - ExitOnError("Could not register window class"); - - // Create the main window - hMainWindow = - CreateWindow("RemoteShellServerWindowClass", - "Remote Shell Server", - WS_OVERLAPPED|WS_CAPTION|WS_SYSMENU|WS_MINIMIZEBOX, - 20, 20, 500, 300, - NULL, NULL, hInstance, NULL); - if (!hMainWindow) - ExitOnError("Could not create window"); - - ShowWindow(hMainWindow, SW_SHOWMINNOACTIVE); - UpdateWindow(hMainWindow); - - // Main message loop - while (GetMessage(&msg, NULL, 0, 0)) { - TranslateMessage(&msg); - DispatchMessage(&msg); - } - - ExitProcess(0); -} +// Simple remote shell server (and file transfer server) +// Author: Michael Goldish +// Much of the code here was adapted from Microsoft code samples. + +// Usage: rss.exe [shell port] [file transfer port] +// If no shell port is specified the default is 10022. +// If no file transfer port is specified the default is 10023. + +// Definitions: +// A 'msg' is a 32 bit integer. +// A 'packet' is a 64 bit unsigned integer followed by a string of bytes. +// The 64 bit integer indicates the length of the string. + +// Protocol for file transfers: +// +// When uploading files to the server: +// 1. The client connects. +// 2. The server sends RSS_MAGIC. +// 3. The client sends RSS_SET_PATH, followed by a packet (as defined above) +// containing the path (in the server's filesystem) where files are to be +// stored. The 'current path' is set to the path given. +// Uploading a file (optional, can be repeated many times): +// 4. The client sends RSS_CREATE_FILE, followed by a packet containing the +// filename (filename only, without a path), followed by a packet +// containing the file's contents. The file is stored in the current +// path. +// Uploading a directory (optional, can be repeated many times): +// 5. The client sends RSS_CREATE_DIR, followed by a packet containing the +// name of the directory to be created (directory name only, without a +// path). The directory is created in the current path. Then, the +// 'current path' becomes the path of that new directory. +// 6. Steps 4-5 may be repeated to upload files and directories to the new +// directory. +// 7. The client sends RSS_LEAVE_DIR, and the current path is set to the path +// of its parent directory (e.g. C:\Foobar -> C:\). +// 8. The client sends RSS_DONE and waits for a response. +// 9. The server sends RSS_OK to indicate that it's still listening. +// 10. Steps 3-9 are repeated as many times as necessary. +// If a critical error occurs at any time, the server may send RSS_ERROR +// followed by a packet containing an error message, and the connection is +// closed. +// +// When downloading files from the server: +// 1. The client connects. +// 2. The server sends RSS_MAGIC. +// 3. The client sends RSS_SET_PATH, followed by a packet (as defined above) +// containing a path (in the server's filesystem) or a wildcard pattern +// indicating the files/directories the client wants to download. +// The server then searches the given path. For every file found: +// 4. The server sends RSS_CREATE_FILE, followed by a packet containing the +// filename (filename only, without a path), followed by a packet +// containing the file's contents. +// For every directory found: +// 5. The server sends RSS_CREATE_DIR, followed by a packet containing the +// name of the directory to be created (directory name only, without a +// path). +// 6. Steps 4-5 are repeated as many times as necessary. +// 7. The server sends RSS_LEAVE_DIR. +// 8. The server sends RSS_DONE. +// 9. Steps 3-8 are repeated as many times as necessary. +// If a critical error occurs, the server may send RSS_ERROR followed by a +// packet containing an error message, and the connection is closed. +// RSS_ERROR may only be sent when the client expects a msg. + +#define _WIN32_WINNT 0x0500 + +#include +#include +#include +#include +#include + +#pragma comment(lib, "ws2_32.lib") +#pragma comment(lib, "shlwapi.lib") + +#define TEXTBOX_LIMIT 16384 + +// Constants for file transfer server +#define RSS_MAGIC 0x525353 +#define RSS_OK 1 +#define RSS_ERROR 2 +#define RSS_UPLOAD 3 +#define RSS_DOWNLOAD 4 +#define RSS_SET_PATH 5 +#define RSS_CREATE_FILE 6 +#define RSS_CREATE_DIR 7 +#define RSS_LEAVE_DIR 8 +#define RSS_DONE 9 + +// Globals +int shell_port = 10022; +int file_transfer_port = 10023; + +HWND hMainWindow = NULL; +HWND hTextBox = NULL; + +struct client_info { + SOCKET socket; + sockaddr_in addr; + char addr_str[256]; + int pid; + HWND hwnd; + HANDLE hJob; + HANDLE hChildOutputRead; + HANDLE hThreadChildToSocket; +}; + +void ExitOnError(const char *message, BOOL winsock = FALSE) +{ + LPVOID system_message; + char buffer[512]; + int error_code; + + if (winsock) + error_code = WSAGetLastError(); + else + error_code = GetLastError(); + WSACleanup(); + + FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | + FORMAT_MESSAGE_FROM_SYSTEM, + NULL, + error_code, + MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), + (LPTSTR)&system_message, + 0, + NULL); + sprintf(buffer, + "%s!\n" + "Error code = %d\n" + "Error message = %s", + message, error_code, (char *)system_message); + MessageBox(hMainWindow, buffer, "Error", MB_OK | MB_ICONERROR); + + LocalFree(system_message); + ExitProcess(1); +} + +void _vAppendMessage(const char *message, va_list args) +{ + char str[512] = {0}; + + vsnprintf(str, sizeof(str) - 1, message, args); + strncat(str, "\r\n", sizeof(str) - 1 - strlen(str)); + + int length = GetWindowTextLength(hTextBox); + SendMessage(hTextBox, EM_SETSEL, (WPARAM)length, (LPARAM)length); + SendMessage(hTextBox, EM_REPLACESEL, (WPARAM)FALSE, (LPARAM)str); +} + +// Append text to the textbox +// (if the textbox is full, do nothing) +void _AppendMessage(const char *message, ...) +{ + va_list args; + + va_start(args, message); + _vAppendMessage(message, args); + va_end(args); +} + +// Append text to the textbox +// (if the textbox is full or nearly full, remove old text first) +void AppendMessage(const char *message, ...) +{ + va_list args; + + int length = GetWindowTextLength(hTextBox); + if (length > TEXTBOX_LIMIT - 512) { + SendMessage(hTextBox, EM_SETSEL, (WPARAM)0, + (LPARAM)(TEXTBOX_LIMIT * 3 / 4)); + SendMessage(hTextBox, EM_REPLACESEL, (WPARAM)FALSE, (LPARAM)"..."); + if (length == TEXTBOX_LIMIT) _AppendMessage("..."); + } + va_start(args, message); + _vAppendMessage(message, args); + va_end(args); +} + +void FormatStringForPrinting(char *dst, const char *src, int size) +{ + int j = 0; + + for (int i = 0; i < size && src[i]; i++) { + if (src[i] == '\n') { + dst[j++] = '\\'; + dst[j++] = 'n'; + } else if (src[i] == '\r') { + dst[j++] = '\\'; + dst[j++] = 'r'; + } else if (src[i] == '\t') { + dst[j++] = '\\'; + dst[j++] = 't'; + } else if (src[i] == '\\') { + dst[j++] = '\\'; + dst[j++] = '\\'; + } else dst[j++] = src[i]; + } + dst[j] = 0; +} + +const char* GetClientIPAddress(client_info *ci) +{ + char *address = inet_ntoa(ci->addr.sin_addr); + if (address) + return address; + else + return "unknown"; +} + +// Read a small chunk of data into a buffer +BOOL Receive(SOCKET socket, char *buffer, int len) +{ + while (len > 0) { + int bytes_received = recv(socket, buffer, len, 0); + if (bytes_received <= 0) + return FALSE; + buffer += bytes_received; + len -= bytes_received; + } + return TRUE; +} + +// Send data from a buffer +BOOL Send(SOCKET socket, const char *buffer, int len) +{ + while (len > 0) { + int bytes_sent = send(socket, buffer, len, 0); + if (bytes_sent <= 0) + return FALSE; + buffer += bytes_sent; + len -= bytes_sent; + } + return TRUE; +} + +/*------------- + * Shell server + *-------------*/ + +DWORD WINAPI ChildToSocket(LPVOID client_info_ptr) +{ + client_info *ci = (client_info *)client_info_ptr; + char buffer[1024]; + DWORD bytes_read; + + while (1) { + // Read data from the child's STDOUT/STDERR pipes + if (!ReadFile(ci->hChildOutputRead, + buffer, sizeof(buffer), + &bytes_read, NULL) || !bytes_read) { + if (GetLastError() == ERROR_BROKEN_PIPE) + break; // Pipe done -- normal exit path + else + ExitOnError("ReadFile failed"); // Something bad happened + } + // Send data to the client + Send(ci->socket, buffer, bytes_read); + } + + AppendMessage("Child exited"); + closesocket(ci->socket); + return 0; +} + +DWORD WINAPI SocketToChild(LPVOID client_info_ptr) +{ + client_info *ci = (client_info *)client_info_ptr; + char buffer[256], formatted_buffer[768]; + int bytes_received; + + AppendMessage("Shell server: new client connected (%s)", ci->addr_str); + + while (1) { + // Receive data from the socket + ZeroMemory(buffer, sizeof(buffer)); + bytes_received = recv(ci->socket, buffer, sizeof(buffer), 0); + if (bytes_received <= 0) + break; + // Report the data received + FormatStringForPrinting(formatted_buffer, buffer, sizeof(buffer)); + _AppendMessage("Client (%s) entered text: \"%s\"", + ci->addr_str, formatted_buffer); + // Send the data as a series of WM_CHAR messages to the console window + for (int i = 0; i < bytes_received; i++) { + SendMessage(ci->hwnd, WM_CHAR, (WPARAM)buffer[i], 0); + SendMessage(ci->hwnd, WM_SETFOCUS, 0, 0); + } + } + + AppendMessage("Shell server: client disconnected (%s)", ci->addr_str); + + // Attempt to terminate the child's process tree: + // Using taskkill (where available) + sprintf(buffer, "taskkill /PID %d /T /F", ci->pid); + system(buffer); + // .. and using TerminateJobObject() + TerminateJobObject(ci->hJob, 0); + // Wait for the ChildToSocket thread to terminate + WaitForSingleObject(ci->hThreadChildToSocket, 10000); + // In case the thread refuses to exit, terminate it + TerminateThread(ci->hThreadChildToSocket, 0); + // Close the socket + closesocket(ci->socket); + + // Free resources + CloseHandle(ci->hJob); + CloseHandle(ci->hThreadChildToSocket); + CloseHandle(ci->hChildOutputRead); + free(ci); + + AppendMessage("SocketToChild thread exited"); + return 0; +} + +void PrepAndLaunchRedirectedChild(client_info *ci, + HANDLE hChildStdOut, + HANDLE hChildStdErr) +{ + PROCESS_INFORMATION pi; + STARTUPINFO si; + + // Allocate a new console for the child + HWND hwnd = GetForegroundWindow(); + FreeConsole(); + AllocConsole(); + ShowWindow(GetConsoleWindow(), SW_HIDE); + if (hwnd) + SetForegroundWindow(hwnd); + + // Set up the start up info struct. + ZeroMemory(&si, sizeof(STARTUPINFO)); + si.cb = sizeof(STARTUPINFO); + si.dwFlags = STARTF_USESTDHANDLES | STARTF_USESHOWWINDOW; + si.hStdOutput = hChildStdOut; + si.hStdInput = GetStdHandle(STD_INPUT_HANDLE); + si.hStdError = hChildStdErr; + // Use this if you want to hide the child: + si.wShowWindow = SW_HIDE; + // Note that dwFlags must include STARTF_USESHOWWINDOW if you want to + // use the wShowWindow flags. + + // Launch the process that you want to redirect. + if (!CreateProcess(NULL, "cmd.exe", NULL, NULL, TRUE, + 0, NULL, "C:\\", &si, &pi)) + ExitOnError("CreateProcess failed"); + + // Close any unnecessary handles. + if (!CloseHandle(pi.hThread)) + ExitOnError("CloseHandle failed"); + + // Keep the process ID + ci->pid = pi.dwProcessId; + // Assign the process to a newly created JobObject + ci->hJob = CreateJobObject(NULL, NULL); + AssignProcessToJobObject(ci->hJob, pi.hProcess); + // Keep the console window's handle + ci->hwnd = GetConsoleWindow(); + + // Detach from the child's console + FreeConsole(); +} + +void SpawnSession(client_info *ci) +{ + HANDLE hOutputReadTmp, hOutputRead, hOutputWrite; + HANDLE hErrorWrite; + SECURITY_ATTRIBUTES sa; + + // Set up the security attributes struct. + sa.nLength = sizeof(SECURITY_ATTRIBUTES); + sa.lpSecurityDescriptor = NULL; + sa.bInheritHandle = TRUE; + + // Create the child output pipe. + if (!CreatePipe(&hOutputReadTmp, &hOutputWrite, &sa, 0)) + ExitOnError("CreatePipe failed"); + + // Create a duplicate of the output write handle for the std error + // write handle. This is necessary in case the child application + // closes one of its std output handles. + if (!DuplicateHandle(GetCurrentProcess(), hOutputWrite, + GetCurrentProcess(), &hErrorWrite, 0, + TRUE, DUPLICATE_SAME_ACCESS)) + ExitOnError("DuplicateHandle failed"); + + // Create new output read handle and the input write handles. Set + // the Properties to FALSE. Otherwise, the child inherits the + // properties and, as a result, non-closeable handles to the pipes + // are created. + if (!DuplicateHandle(GetCurrentProcess(), hOutputReadTmp, + GetCurrentProcess(), + &hOutputRead, // Address of new handle. + 0, FALSE, // Make it uninheritable. + DUPLICATE_SAME_ACCESS)) + ExitOnError("DuplicateHandle failed"); + + // Close inheritable copies of the handles you do not want to be + // inherited. + if (!CloseHandle(hOutputReadTmp)) + ExitOnError("CloseHandle failed"); + + PrepAndLaunchRedirectedChild(ci, hOutputWrite, hErrorWrite); + + ci->hChildOutputRead = hOutputRead; + + // Close pipe handles (do not continue to modify the parent). + // You need to make sure that no handles to the write end of the + // output pipe are maintained in this process or else the pipe will + // not close when the child process exits and the ReadFile will hang. + if (!CloseHandle(hOutputWrite)) ExitOnError("CloseHandle failed"); + if (!CloseHandle(hErrorWrite)) ExitOnError("CloseHandle failed"); +} + +SOCKET PrepareListenSocket(int port) +{ + sockaddr_in addr; + linger l; + int result; + + // Create socket + SOCKET ListenSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP); + if (ListenSocket == INVALID_SOCKET) + ExitOnError("Socket creation failed", TRUE); + + // Enable lingering + l.l_linger = 10; + l.l_onoff = 1; + setsockopt(ListenSocket, SOL_SOCKET, SO_LINGER, (char *)&l, sizeof(l)); + + // Bind the socket + addr.sin_family = AF_INET; + addr.sin_addr.s_addr = htonl(INADDR_ANY); + addr.sin_port = htons(port); + + result = bind(ListenSocket, (sockaddr *)&addr, sizeof(addr)); + if (result == SOCKET_ERROR) + ExitOnError("bind failed", TRUE); + + // Start listening for incoming connections + result = listen(ListenSocket, SOMAXCONN); + if (result == SOCKET_ERROR) + ExitOnError("listen failed", TRUE); + + return ListenSocket; +} + +DWORD WINAPI ShellListenThread(LPVOID param) +{ + client_info _ci, *ci; + HANDLE hThread; + + SOCKET ListenSocket = PrepareListenSocket(shell_port); + + // Inform the user + AppendMessage("Shell server: waiting for clients to connect..."); + + while (1) { + int addrlen = sizeof(_ci.addr); + _ci.socket = accept(ListenSocket, (sockaddr *)&_ci.addr, &addrlen); + if (_ci.socket == INVALID_SOCKET) { + if (WSAGetLastError() == WSAEINTR) + break; + else + ExitOnError("accept failed", TRUE); + } + + if (!(ci = (client_info *)malloc(sizeof(client_info)))) + ExitOnError("Could not allocate client_info struct"); + memcpy(ci, &_ci, sizeof(client_info)); + sprintf(ci->addr_str, "%s:%d", GetClientIPAddress(ci), + ci->addr.sin_port); + + // Under heavy load, spawning cmd.exe might take a while, so tell the + // client to be patient + const char *message = "Please wait...\r\n"; + Send(ci->socket, message, strlen(message)); + // Spawn a new redirected cmd.exe process + SpawnSession(ci); + // Start transferring data from the child process to the client + hThread = CreateThread(NULL, 0, ChildToSocket, (LPVOID)ci, 0, NULL); + if (!hThread) + ExitOnError("Could not create ChildToSocket thread"); + ci->hThreadChildToSocket = hThread; + // ... and from the client to the child process + hThread = CreateThread(NULL, 0, SocketToChild, (LPVOID)ci, 0, NULL); + if (!hThread) + ExitOnError("Could not create SocketToChild thread"); + } + + return 0; +} + +/*--------------------- + * File transfer server + *---------------------*/ + +#define BUFSIZE 65536 + +// Receive up to 4GB into a file +BOOL ReceiveIntoFile(SOCKET socket, const char *filename, DWORD len) +{ + char *buffer = (char *)malloc(BUFSIZE); + if (!buffer) return FALSE; + + FILE *fp = fopen(filename, "wb"); + if (!fp) { + free(buffer); + return FALSE; + } + + while (len > 0) { + int bytes_received = recv(socket, buffer, min(BUFSIZE, len), 0); + if (bytes_received <= 0) + break; + if (fwrite(buffer, bytes_received, 1, fp) < 1) + break; + len -= bytes_received; + } + + fclose(fp); + free(buffer); + return len == 0; +} + +// Send data from a file (unlimited size) +BOOL SendFromFile(SOCKET socket, const char *filename) +{ + char *buffer = (char *)malloc(BUFSIZE); + if (!buffer) return FALSE; + + FILE *fp = fopen(filename, "rb"); + if (!fp) { + free(buffer); + return FALSE; + } + + while (!feof(fp)) { + int bytes_read = fread(buffer, 1, BUFSIZE, fp); + if (bytes_read <= 0) + break; + if (!Send(socket, buffer, bytes_read)) { + fclose(fp); + free(buffer); + return FALSE; + } + } + + BOOL success = feof(fp); + fclose(fp); + free(buffer); + return success; +} + +BOOL ReceivePacket(SOCKET socket, char *buffer, DWORD max_size) +{ + DWORD packet_size = 0, padding; + + if (!Receive(socket, (char *)&packet_size, 4)) return FALSE; + if (!Receive(socket, (char *)&padding, 4)) return FALSE; + if (packet_size > max_size) return FALSE; + return Receive(socket, buffer, packet_size); +} + +BOOL ReceiveStrPacket(SOCKET socket, char *buffer, DWORD max_size) +{ + memset(buffer, 0, max_size); + return ReceivePacket(socket, buffer, max_size - 1); +} + +BOOL SendPacket(SOCKET socket, const char *buffer, DWORD len) +{ + DWORD padding = 0; + + if (!Send(socket, (char *)&len, 4)) return FALSE; + if (!Send(socket, (char *)&padding, 4)) return FALSE; + return Send(socket, buffer, len); +} + +BOOL ReceivePacketIntoFile(SOCKET socket, const char *filename) +{ + DWORD packet_size = 0, padding; + + if (!Receive(socket, (char *)&packet_size, 4)) return FALSE; + if (!Receive(socket, (char *)&padding, 4)) return FALSE; + return ReceiveIntoFile(socket, filename, packet_size); +} + +BOOL SendPacketFromFile(SOCKET socket, const char *filename, + DWORD size_low, DWORD size_high) +{ + if (!Send(socket, (char *)&size_low, 4)) return FALSE; + if (!Send(socket, (char *)&size_high, 4)) return FALSE; + return SendFromFile(socket, filename); +} + +BOOL SendMsg(SOCKET socket, int msg) +{ + return Send(socket, (char *)&msg, 4); +} + +BOOL ExpandPath(char *path, int size) +{ + char temp[512]; + int result; + + PathRemoveBackslash(path); + result = ExpandEnvironmentStrings(path, temp, sizeof(temp)); + if (result == 0 || result > sizeof(temp)) + return FALSE; + strncpy(path, temp, size - 1); + return TRUE; +} + +int TerminateTransfer(client_info *ci, const char *message) +{ + AppendMessage(message); + AppendMessage("File transfer server: client disconnected (%s)", + ci->addr_str); + closesocket(ci->socket); + free(ci); + return 0; +} + +int TerminateWithError(client_info *ci, const char *message) +{ + SendMsg(ci->socket, RSS_ERROR); + SendPacket(ci->socket, message, strlen(message)); + return TerminateTransfer(ci, message); +} + +int ReceiveThread(client_info *ci) +{ + char path[512], filename[512]; + int msg = 0; + + AppendMessage("Client (%s) wants to upload files", ci->addr_str); + + while (1) { + if (!Receive(ci->socket, (char *)&msg, 4)) + return TerminateTransfer(ci, "Could not receive further msgs"); + + switch (msg) { + case RSS_SET_PATH: + if (!ReceiveStrPacket(ci->socket, path, sizeof(filename))) + return TerminateWithError(ci, + "RSS_SET_PATH: could not receive path, or path too long"); + _AppendMessage("Client (%s) set path to %s", ci->addr_str, path); + if (!ExpandPath(path, sizeof(path))) + return TerminateWithError(ci, + "RSS_SET_PATH: error expanding environment strings"); + break; + + case RSS_CREATE_FILE: + if (!ReceiveStrPacket(ci->socket, filename, sizeof(filename))) + return TerminateWithError(ci, + "RSS_CREATE_FILE: could not receive filename"); + if (PathIsDirectory(path)) + PathAppend(path, filename); + _AppendMessage("Client (%s) is uploading %s", ci->addr_str, path); + if (!ReceivePacketIntoFile(ci->socket, path)) + return TerminateWithError(ci, + "RSS_CREATE_FILE: error receiving or writing file " + "contents"); + PathAppend(path, ".."); + break; + + case RSS_CREATE_DIR: + if (!ReceiveStrPacket(ci->socket, filename, sizeof(filename))) + return TerminateWithError(ci, + "RSS_CREATE_DIR: could not receive dirname"); + if (PathIsDirectory(path)) + PathAppend(path, filename); + _AppendMessage("Entering dir %s", path); + if (PathFileExists(path)) { + if (!PathIsDirectory(path)) + return TerminateWithError(ci, + "RSS_CREATE_DIR: path exists and is not a directory"); + } else { + if (!CreateDirectory(path, NULL)) + return TerminateWithError(ci, + "RSS_CREATE_DIR: could not create directory"); + } + break; + + case RSS_LEAVE_DIR: + PathAppend(path, ".."); + _AppendMessage("Returning to dir %s", path); + break; + + case RSS_DONE: + if (!SendMsg(ci->socket, RSS_OK)) + return TerminateTransfer(ci, + "RSS_DONE: could not send OK msg"); + break; + + default: + return TerminateWithError(ci, "Received unexpected msg"); + } + } +} + +// Given a path or a pattern with wildcards, send files or directory trees to +// the client +int SendFiles(client_info *ci, const char *pattern) +{ + char path[512]; + WIN32_FIND_DATA ffd; + FILE *fp; + + HANDLE hFind = FindFirstFile(pattern, &ffd); + if (hFind == INVALID_HANDLE_VALUE) { + // If a weird error occurred (like failure to list directory contents + // due to insufficient permissions) print a warning and continue. + if (GetLastError() != ERROR_FILE_NOT_FOUND) + AppendMessage("WARNING: FindFirstFile failed on pattern %s", + pattern); + return 1; + } + + strncpy(path, pattern, sizeof(path) - 1); + PathAppend(path, ".."); + + do { + if (ffd.dwFileAttributes & FILE_ATTRIBUTE_REPARSE_POINT) + continue; + if (ffd.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) { + // Directory + if (!strcmp(ffd.cFileName, ".") || !strcmp(ffd.cFileName, "..")) + continue; + PathAppend(path, ffd.cFileName); + _AppendMessage("Entering dir %s", path); + PathAppend(path, "*"); + if (!SendMsg(ci->socket, RSS_CREATE_DIR)) { + FindClose(hFind); + return TerminateTransfer(ci, + "Could not send RSS_CREATE_DIR msg"); + } + if (!SendPacket(ci->socket, ffd.cFileName, + strlen(ffd.cFileName))) { + FindClose(hFind); + return TerminateTransfer(ci, "Could not send dirname"); + } + if (!SendFiles(ci, path)) { + FindClose(hFind); + return 0; + } + if (!SendMsg(ci->socket, RSS_LEAVE_DIR)) { + FindClose(hFind); + return TerminateTransfer(ci, + "Could not send RSS_LEAVE_DIR msg"); + } + PathAppend(path, ".."); + PathAppend(path, ".."); + _AppendMessage("Returning to dir %s", path); + } else { + // File + PathAppend(path, ffd.cFileName); + _AppendMessage("Client (%s) is downloading %s", + ci->addr_str, path); + // Make sure the file is readable + fp = fopen(path, "rb"); + if (fp) fclose(fp); + else { + AppendMessage("WARNING: could not read file %s", path); + PathAppend(path, ".."); + continue; + } + if (!SendMsg(ci->socket, RSS_CREATE_FILE)) { + FindClose(hFind); + return TerminateTransfer(ci, + "Could not send RSS_CREATE_FILE msg"); + } + if (!SendPacket(ci->socket, ffd.cFileName, + strlen(ffd.cFileName))) { + FindClose(hFind); + return TerminateTransfer(ci, "Could not send filename"); + } + if (!SendPacketFromFile(ci->socket, path, + ffd.nFileSizeLow, ffd.nFileSizeHigh)) { + FindClose(hFind); + return TerminateTransfer(ci, "Could not send file contents"); + } + PathAppend(path, ".."); + } + } while (FindNextFile(hFind, &ffd)); + + if (GetLastError() == ERROR_NO_MORE_FILES) { + FindClose(hFind); + return 1; + } else { + FindClose(hFind); + return TerminateWithError(ci, "FindNextFile failed"); + } +} + +int SendThread(client_info *ci) +{ + char pattern[512]; + int msg = 0; + + AppendMessage("Client (%s) wants to download files", ci->addr_str); + + while (1) { + if (!Receive(ci->socket, (char *)&msg, 4)) + return TerminateTransfer(ci, "Could not receive further msgs"); + + switch (msg) { + case RSS_SET_PATH: + if (!ReceiveStrPacket(ci->socket, pattern, sizeof(pattern))) + return TerminateWithError(ci, + "RSS_SET_PATH: could not receive path, or path too long"); + AppendMessage("Client (%s) asked for %s", ci->addr_str, pattern); + if (!ExpandPath(pattern, sizeof(pattern))) + return TerminateWithError(ci, + "RSS_SET_PATH: error expanding environment strings"); + + if (!SendFiles(ci, pattern)) + return 0; + if (!SendMsg(ci->socket, RSS_DONE)) + return TerminateTransfer(ci, + "RSS_SET_PATH: could not send RSS_DONE msg"); + break; + + default: + return TerminateWithError(ci, "Received unexpected msg"); + } + } +} + +DWORD WINAPI TransferThreadEntry(LPVOID client_info_ptr) +{ + client_info *ci = (client_info *)client_info_ptr; + int msg = 0; + + AppendMessage("File transfer server: new client connected (%s)", + ci->addr_str); + + if (!SendMsg(ci->socket, RSS_MAGIC)) + return TerminateTransfer(ci, "Could not send greeting message"); + if (!Receive(ci->socket, (char *)&msg, 4)) + return TerminateTransfer(ci, "Error receiving msg"); + + if (msg == RSS_UPLOAD) + return ReceiveThread(ci); + else if (msg == RSS_DOWNLOAD) + return SendThread(ci); + return TerminateWithError(ci, "Received unexpected msg"); +} + +DWORD WINAPI FileTransferListenThread(LPVOID param) +{ + client_info _ci, *ci; + HANDLE hThread; + + SOCKET ListenSocket = PrepareListenSocket(file_transfer_port); + + // Inform the user + AppendMessage("File transfer server: waiting for clients to connect..."); + + while (1) { + int addrlen = sizeof(_ci.addr); + _ci.socket = accept(ListenSocket, (sockaddr *)&_ci.addr, &addrlen); + if (_ci.socket == INVALID_SOCKET) { + if (WSAGetLastError() == WSAEINTR) + break; + else + ExitOnError("accept failed", TRUE); + } + + if (!(ci = (client_info *)malloc(sizeof(client_info)))) + ExitOnError("Could not allocate client_info struct"); + memcpy(ci, &_ci, sizeof(client_info)); + sprintf(ci->addr_str, "%s:%d", GetClientIPAddress(ci), + ci->addr.sin_port); + hThread = CreateThread(NULL, 0, TransferThreadEntry, (LPVOID)ci, 0, + NULL); + if (!hThread) + ExitOnError("Could not create MainTransferThread"); + } + + return 0; +} + +/*-------------------- + * WndProc and WinMain + *--------------------*/ + +LRESULT CALLBACK WndProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam) +{ + RECT rect; + WSADATA wsaData; + + switch (msg) { + case WM_CREATE: + // Create text box + GetClientRect(hwnd, &rect); + hTextBox = CreateWindowEx(WS_EX_CLIENTEDGE, + "EDIT", "", + WS_CHILD | WS_VISIBLE | WS_VSCROLL | + ES_MULTILINE | ES_AUTOVSCROLL, + 20, 20, + rect.right - 40, + rect.bottom - 40, + hwnd, + NULL, + GetModuleHandle(NULL), + NULL); + if (!hTextBox) + ExitOnError("Could not create text box"); + // Set font + SendMessage(hTextBox, WM_SETFONT, + (WPARAM)GetStockObject(DEFAULT_GUI_FONT), + MAKELPARAM(FALSE, 0)); + // Set size limit + SendMessage(hTextBox, EM_LIMITTEXT, (WPARAM)TEXTBOX_LIMIT, (LPARAM)0); + // Initialize Winsock + if (WSAStartup(MAKEWORD(2,2), &wsaData)) + ExitOnError("Winsock initialization failed"); + // Start the listening threads + if (!CreateThread(NULL, 0, ShellListenThread, NULL, 0, NULL)) + ExitOnError("Could not create shell server thread"); + if (!CreateThread(NULL, 0, FileTransferListenThread, NULL, 0, NULL)) + ExitOnError("Could not create file transfer server thread"); + break; + + case WM_DESTROY: + if (WSACleanup()) + ExitOnError("WSACleanup failed"); + PostQuitMessage(0); + break; + + default: + return DefWindowProc(hwnd, msg, wParam, lParam); + } + + return 0; +} + +int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, + LPSTR lpCmdLine, int nShowCmd) +{ + WNDCLASSEX wc; + MSG msg; + + if (strlen(lpCmdLine)) + sscanf(lpCmdLine, "%d %d", &shell_port, &file_transfer_port); + + // Make sure the firewall is disabled + system("netsh firewall set opmode disable"); + + // Create the window class + wc.cbSize = sizeof(WNDCLASSEX); + wc.style = CS_HREDRAW | CS_VREDRAW; + wc.lpfnWndProc = WndProc; + wc.cbClsExtra = 0; + wc.cbWndExtra = 0; + wc.hInstance = hInstance; + wc.hIcon = LoadIcon(NULL, IDI_APPLICATION); + wc.hIconSm = LoadIcon(NULL, IDI_APPLICATION); + wc.hbrBackground = (HBRUSH)(COLOR_BTNFACE + 1); + wc.lpszMenuName = NULL; + wc.lpszClassName = "RemoteShellServerWindowClass"; + wc.hCursor = LoadCursor(NULL, IDC_ARROW); + + if (!RegisterClassEx(&wc)) + ExitOnError("Could not register window class"); + + // Create the main window + hMainWindow = + CreateWindow("RemoteShellServerWindowClass", + "Remote Shell Server", + WS_OVERLAPPED | WS_CAPTION | WS_SYSMENU | WS_MINIMIZEBOX, + 20, 20, 600, 400, + NULL, NULL, hInstance, NULL); + if (!hMainWindow) + ExitOnError("Could not create window"); + + ShowWindow(hMainWindow, SW_SHOWMINNOACTIVE); + UpdateWindow(hMainWindow); + + // Main message loop + while (GetMessage(&msg, NULL, 0, 0)) { + TranslateMessage(&msg); + DispatchMessage(&msg); + } + + ExitProcess(0); +}