This commit is contained in:
manojampalam 2016-04-13 19:59:06 -07:00
parent 8fb27e8fc6
commit 491769a99f
5 changed files with 162 additions and 69 deletions

View File

@ -20,8 +20,12 @@
</ProjectConfiguration>
</ItemGroup>
<ItemGroup>
<ClCompile Include="$(OpenSSH-Src-Path)\contrib\win32\win32compat\ssh-agent\ssh-agent-main.cpp" />
<ClCompile Include="$(OpenSSH-Src-Path)\contrib\win32\win32compat\ssh-agent\ssh-agent.cpp" />
<ClCompile Include="..\win32compat\ssh-agent\agent-main.c" />
<ClCompile Include="..\win32compat\ssh-agent\agent.c" />
<ClCompile Include="..\win32compat\ssh-agent\connection.c" />
</ItemGroup>
<ItemGroup>
<ClInclude Include="..\win32compat\ssh-agent\agent.h" />
</ItemGroup>
<PropertyGroup Label="Globals">
<ProjectGuid>{F6644EC5-D6B6-42A1-828C-75E2977470E0}</ProjectGuid>

View File

@ -66,7 +66,7 @@ static VOID WINAPI service_handler(DWORD dwControl)
{
case SERVICE_CONTROL_STOP: {
ReportSvcStatus(SERVICE_STOP_PENDING, NO_ERROR, 500);
GenerateConsoleCtrlEvent(CTRL_BREAK_EVENT, 0);
agent_shutdown();
ReportSvcStatus(SERVICE_STOPPED, NO_ERROR, 0);
return;
}
@ -79,11 +79,20 @@ static VOID WINAPI service_handler(DWORD dwControl)
ReportSvcStatus(service_status.dwCurrentState, NO_ERROR, 0);
}
BOOL WINAPI ctrl_c_handler(
_In_ DWORD dwCtrlType
) {
/* for any Ctrl type, shutdown agent*/
agent_shutdown();
return TRUE;
}
int main() {
if (!StartServiceCtrlDispatcher(diapatch_table)) {
if (GetLastError() == ERROR_FAILED_SERVICE_CONTROLLER_CONNECT) {
//console app
start_agent();
SetConsoleCtrlHandler(ctrl_c_handler, TRUE);
return agent_start();
}
else
return -1;
@ -97,6 +106,6 @@ int scm_start_servie(DWORD num, LPWSTR* args) {
service_status.dwServiceType = SERVICE_WIN32_OWN_PROCESS;
ReportSvcStatus(SERVICE_START_PENDING, NO_ERROR, 300);
ReportSvcStatus(SERVICE_RUNNING, NO_ERROR, 0);
return start_agent();
return agent_start();
}

View File

@ -36,17 +36,11 @@ static HANDLE ioc_port;
static volatile long action_queue;
static struct agent_connection* list;
enum agent_sm_event {
NEW_CLIENT_CONNECTION = 1,
CONNECTION_DONE = 2,
SHUTDOWN = 3
};
#define ACTION_LISTEN 0x80000000
#define ACTION_SHUTDOWN 0x40000000
void agent_sm_process_action_queue() {
long actions_remaining = 0;
do {
if (action_queue & ACTION_SHUTDOWN) {
/* go through the list and disconect each connection */
@ -58,10 +52,12 @@ void agent_sm_process_action_queue() {
/* remove unwanted queued actions */
InterlockedAnd(&action_queue, ~ACTION_LISTEN);
actions_remaining = InterlockedAnd(&action_queue, ~ACTION_SHUTDOWN);
if (InterlockedAnd(&action_queue, ~ACTION_SHUTDOWN) == ACTION_SHUTDOWN)
break;
}
else if (action_queue & ACTION_LISTEN) {
HANDLE h;
HANDLE h, temp;
long prev_queue;
struct agent_connection* con =
(struct agent_connection*)malloc(sizeof(struct agent_connection));
memset(con, 0, sizeof(struct agent_connection));
@ -78,12 +74,13 @@ void agent_sm_process_action_queue() {
NULL);
/* remove action from queue before assigning iocp port*/
con->connection = h;
con->next = list;
list = con;
actions_remaining = InterlockedAnd(&action_queue, ~ACTION_LISTEN);
CreateIoCompletionPort(h, ioc_port, con, 0);
prev_queue = InterlockedAnd(&action_queue, ~ACTION_LISTEN);
temp = CreateIoCompletionPort(h, ioc_port, (ULONG_PTR)con, 0);
if (prev_queue == ACTION_LISTEN)
break;
}
else {
/* cleanup up a done connection*/
@ -101,35 +98,30 @@ void agent_sm_process_action_queue() {
prev = tmp;
tmp = tmp->next;
}
actions_remaining = InterlockedDecrement(&action_queue);
if (InterlockedDecrement(&action_queue) == 0)
break;
}
} while (actions_remaining);
} while (1);
}
void agent_sm_raise(enum agent_sm_event event) {
long ret = 0;
switch (event) {
case NEW_CLIENT_CONNECTION:
ret = InterlockedOr(&action_queue, ACTION_LISTEN);
if (ret == 0)
agent_sm_process_action_queue();
break;
case SHUTDOWN:
ret = InterlockedOr(&action_queue, ACTION_SHUTDOWN);
if (ret == 0)
agent_sm_process_action_queue();
break;
case CONNECTION_DONE:
ret = InterlockedIncrement(&action_queue);
if (ret == 1)
agent_sm_process_action_queue();
break;
default:
DebugBreak();
}
void
agent_cleanup_connection(struct agent_connection* con) {
if (InterlockedIncrement(&action_queue) == 1)
agent_sm_process_action_queue();
}
/* is this the first action queued */
void
agent_listen() {
if (InterlockedOr(&action_queue, ACTION_LISTEN) == 0)
agent_sm_process_action_queue();
}
void agent_shutdown() {
if (InterlockedOr(&action_queue, ACTION_SHUTDOWN) == 0)
agent_sm_process_action_queue();
while (list != NULL)
Sleep(100);
}
HANDLE iocp_workers[4];
@ -139,7 +131,7 @@ DWORD WINAPI iocp_work(LPVOID lpParam) {
struct agent_connection* con;
OVERLAPPED *p_ol;
while (1) {
GetQueuedCompletionStatus(ioc_port, &bytes, &con, &p_ol, INFINITE);
GetQueuedCompletionStatus(ioc_port, &bytes, &(ULONG_PTR)con, &p_ol, INFINITE);
agent_connection_on_io(con, bytes, p_ol);
}
}
@ -148,17 +140,13 @@ int agent_start() {
int i;
action_queue = 0;
list = NULL;
ioc_port = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, NULL, 0);
ioc_port = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, (ULONG_PTR)NULL, 0);
for (i = 0; i < 4; i++)
iocp_workers[i] = CreateThread(NULL, 0, iocp_work, NULL, 0, NULL);
for (i = 0; i < 3; i++)
QueueUserWorkItem(iocp_work, NULL, 0);
action_queue = ACTION_LISTEN;
agent_sm_process_action_queue();
agent_listen();
iocp_work(NULL);
return 1;
}
void agent_shutdown() {
agent_sm_raise(SHUTDOWN);
while (list != NULL)
Sleep(100);
}

View File

@ -1,20 +1,30 @@
#include <Windows.h>
#include <stdio.h>
#define BUFSIZE 5 * 1024
int agent_start();
void agent_shutdown();
#define HEADER_SIZE 4
struct agent_connection {
OVERLAPPED ol;
HANDLE connection;
struct {
DWORD size;
DWORD read;
char buf[BUFSIZE];
} request;
enum {
LISTENING = 0,
READING_HEADER,
READING,
WRITING,
DONE
} state;
HANDLE connection;
struct agent_con* next;
struct agent_connection* next;
};
void agent_connection_on_io(struct agent_connection*, DWORD, OVERLAPPED*);
void agent_connection_disconnect(struct agent_connection*);
void agent_connection_disconnect(struct agent_connection*);
int agent_start();
void agent_shutdown();
void agent_listen();
void agent_cleanup_connection(struct agent_connection*);

View File

@ -29,18 +29,100 @@
* THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/
#include "agent.h"
#define BUFSIZE 5 * 1024
void agent_connection_on_io(struct agent_connection* con, DWORD bytes, OVERLAPPED* ol) {
switch (con->state) {
case LISTENING:
break;
case READING:
break;
case WRITING:
break;
default:
/* process error */
if ( (bytes == 0) && (GetOverlappedResult(con->connection, ol, &bytes, FALSE) == FALSE)) {
con->state = DONE;
agent_cleanup_connection(con);
return;
}
if (con->state == DONE)
DebugBreak();
while (1) {
switch (con->state) {
case WRITING:
/* Writing is done, read next request */
case LISTENING:
con->state = READING_HEADER;
if (con->state == LISTENING)
agent_listen();
ZeroMemory(&con->request, sizeof(con->request));
if (ReadFile(con->connection, con->request.buf,
HEADER_SIZE, NULL, &con->ol)) {
bytes = HEADER_SIZE;
continue;
}
if (GetLastError() != ERROR_IO_PENDING) {
con->state = DONE;
agent_cleanup_connection(con);
return;
}
break;
case READING_HEADER:
con->request.read += bytes;
if (con->request.read == HEADER_SIZE) {
con->request.size = *((DWORD*)con->request.buf);
con->state = READING;
if (ReadFile(con->connection, con->request.buf,
con->request.size, NULL, &con->ol)) {
bytes = con->request.size;
continue;
}
if (GetLastError() != ERROR_IO_PENDING) {
con->state = DONE;
agent_cleanup_connection(con);
return;
}
}
else {
if (ReadFile(con->connection, con->request.buf + con->request.read,
HEADER_SIZE - con->request.read, NULL, &con->ol)) {
bytes = HEADER_SIZE - con->request.read;
continue;
}
if (GetLastError() != ERROR_IO_PENDING) {
con->state = DONE;
agent_cleanup_connection(con);
return;
}
}
break;
case READING:
con->request.read += bytes;
if (con->request.read == con->request.size) {
/* process request and get response */
con->state = WRITING;
if (WriteFile(con->connection, con->request.buf,
con->request.size, NULL, &con->ol)) {
bytes = con->request.size;
continue;
}
if (GetLastError() != ERROR_IO_PENDING) {
con->state = DONE;
agent_cleanup_connection(con);
return;
}
}
else {
if (ReadFile(con->connection, con->request.buf + con->request.read,
con->request.size - con->request.read, NULL, &con->ol)) {
bytes = con->request.size - con->request.read;
continue;
}
if (GetLastError() != ERROR_IO_PENDING) {
con->state = DONE;
agent_cleanup_connection(con);
return;
}
}
break;
default:
DebugBreak();
}
}
}