This commit is contained in:
Manoj Ampalam 2016-05-07 21:04:54 -07:00
parent b22f7cddb6
commit 4c0e2c2078
4 changed files with 206 additions and 117 deletions

View File

@ -84,6 +84,7 @@ BOOL WINAPI ctrl_c_handler(
_In_ DWORD dwCtrlType _In_ DWORD dwCtrlType
) { ) {
/* for any Ctrl type, shutdown agent*/ /* for any Ctrl type, shutdown agent*/
debug("Ctrl+C received");
agent_shutdown(); agent_shutdown();
return TRUE; return TRUE;
} }
@ -98,11 +99,11 @@ int main(int argc, char **argv) {
/* console app - start in debug mode*/ /* console app - start in debug mode*/
SetConsoleCtrlHandler(ctrl_c_handler, TRUE); SetConsoleCtrlHandler(ctrl_c_handler, TRUE);
log_init("ssh-agent", 7, 1, 1); log_init("ssh-agent", 7, 1, 1);
return agent_start(TRUE, FALSE, 0); return agent_start(TRUE, FALSE, 0, 0);
} }
else { else {
log_init("ssh-agent", config_log_level(), 1, 0); log_init("ssh-agent", config_log_level(), 1, 0);
return agent_start(FALSE, TRUE, (HANDLE)atoi(*argv)); return agent_start(FALSE, TRUE, (HANDLE)atoi(*(argv+1)), atoi(*(argv+2)));
} }
} }
else else
@ -118,6 +119,6 @@ int scm_start_servie(DWORD num, LPWSTR* args) {
ReportSvcStatus(SERVICE_START_PENDING, NO_ERROR, 300); ReportSvcStatus(SERVICE_START_PENDING, NO_ERROR, 300);
ReportSvcStatus(SERVICE_RUNNING, NO_ERROR, 0); ReportSvcStatus(SERVICE_RUNNING, NO_ERROR, 0);
log_init("ssh-agent", config_log_level(), 1, 0); log_init("ssh-agent", config_log_level(), 1, 0);
return agent_start(FALSE, FALSE, 0); return agent_start(FALSE, FALSE, 0, 0);
} }

View File

@ -32,108 +32,58 @@
#define AGENT_PIPE_ID L"\\\\.\\pipe\\ssh-agent" #define AGENT_PIPE_ID L"\\\\.\\pipe\\ssh-agent"
#define BUFSIZE 5 * 1024 #define BUFSIZE 5 * 1024
static HANDLE ioc_port; static HANDLE ioc_port = NULL;
static volatile long action_queue; static BOOL debug_mode = FALSE;
static struct agent_connection* list;
#define ACTION_LISTEN 0x80000000 #define NUM_LISTENERS 1
#define ACTION_SHUTDOWN 0x40000000 #define KEY_AGENT_PIPE_ID L"\\\\.\\pipe\\ssh-agent"
void agent_sm_process_action_queue() { static wchar_t *pipe_ids[NUM_LISTENERS] = { KEY_AGENT_PIPE_ID };
static enum agent_type types[NUM_LISTENERS] = { KEY_AGENT };
HANDLE event_stop_agent;
do { struct listener {
if (action_queue & ACTION_SHUTDOWN) { OVERLAPPED ol;
/* go through the list and disconect each connection */ HANDLE pipe;
struct agent_connection* tmp = list; wchar_t *pipe_id;
while (tmp) { enum agent_type type;
agent_connection_disconnect(tmp); SECURITY_ATTRIBUTES sa;
tmp = tmp->next; } listeners[NUM_LISTENERS];
}
/* remove unwanted queued actions */ static int
InterlockedAnd(&action_queue, ~ACTION_LISTEN); init_listeners() {
if (InterlockedAnd(&action_queue, ~ACTION_SHUTDOWN) == ACTION_SHUTDOWN) int i;
break; memset(listeners, 0, sizeof(listeners));
for (i = 0; i < NUM_LISTENERS; i++) {
if ((listeners[i].ol.hEvent = CreateEvent(NULL, TRUE, FALSE, NULL)) == NULL) {
debug("cannot create event ERROR:%d", GetLastError());
return GetLastError();
} }
else if (action_queue & ACTION_LISTEN) { listeners[i].pipe_id = pipe_ids[i];
HANDLE h; listeners[i].type = types[i];
long prev_queue; listeners[i].pipe = INVALID_HANDLE_VALUE;
SECURITY_ATTRIBUTES sa; listeners[i].sa.bInheritHandle = TRUE;
struct agent_connection* con = }
(struct agent_connection*)malloc(sizeof(struct agent_connection));
memset(con, 0, sizeof(struct agent_connection));
memset(&sa, 0, sizeof(sa));
sa.bInheritHandle = FALSE;
sa.lpSecurityDescriptor = NULL;
h = CreateNamedPipeW(
AGENT_PIPE_ID, // pipe name
PIPE_ACCESS_DUPLEX | FILE_FLAG_OVERLAPPED, // read/write access
PIPE_TYPE_BYTE | // message type pipe
PIPE_READMODE_BYTE | // message-read mode
PIPE_WAIT, // blocking mode
PIPE_UNLIMITED_INSTANCES, // max. instances
BUFSIZE, // output buffer size
BUFSIZE, // input buffer size
0, // client time-out
&sa);
/* remove action from queue before assigning iocp port*/ return 0;
con->connection = h;
con->next = list;
list = con;
prev_queue = InterlockedAnd(&action_queue, ~ACTION_LISTEN);
CreateIoCompletionPort(h, ioc_port, (ULONG_PTR)con, 0);
ConnectNamedPipe(h, &con->ol);
if (prev_queue == ACTION_LISTEN)
break;
}
else {
/* cleanup up a done connection*/
struct agent_connection *prev = NULL, *tmp = list;
while (tmp) {
if (tmp->state == DONE) {
if (prev == NULL)
list = tmp->next;
else
prev->next = tmp->next;
CloseHandle(tmp->connection);
printf("deleting %p\n", tmp);
free(tmp);
break;
}
prev = tmp;
tmp = tmp->next;
}
if (InterlockedDecrement(&action_queue) == 0)
break;
}
} while (1);
} }
void static void
agent_cleanup_connection(struct agent_connection* con) { agent_cleanup() {
if (InterlockedIncrement(&action_queue) == 1) int i;
agent_sm_process_action_queue(); for (i = 0; i < NUM_LISTENERS; i++) {
if (listeners[i].ol.hEvent != NULL)
CloseHandle(listeners[i].ol.hEvent);
if (listeners[i].pipe != INVALID_HANDLE_VALUE)
CloseHandle(listeners[i].pipe);
}
if (ioc_port)
CloseHandle(ioc_port);
return;
} }
void static DWORD WINAPI
agent_listen() { iocp_work(LPVOID lpParam) {
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);
CloseHandle(ioc_port);
}
HANDLE iocp_workers[4];
DWORD WINAPI iocp_work(LPVOID lpParam) {
DWORD bytes; DWORD bytes;
struct agent_connection* con = NULL; struct agent_connection* con = NULL;
OVERLAPPED *p_ol; OVERLAPPED *p_ol;
@ -141,35 +91,167 @@ DWORD WINAPI iocp_work(LPVOID lpParam) {
con = NULL; con = NULL;
p_ol = NULL; p_ol = NULL;
if (GetQueuedCompletionStatus(ioc_port, &bytes, &(ULONG_PTR)con, &p_ol, INFINITE) == FALSE) { if (GetQueuedCompletionStatus(ioc_port, &bytes, &(ULONG_PTR)con, &p_ol, INFINITE) == FALSE) {
printf("error: %d on %p \n", GetLastError(), con); debug("iocp error: %d on %p \n", GetLastError(), con);
if (con) if (con)
agent_connection_on_error(con, GetLastError()); agent_connection_on_error(con, GetLastError());
else else
return 0; return 0;
} }
//printf("io on %p state %d bytes %d\n", con, con->state, bytes); else
agent_connection_on_io(con, bytes, p_ol); agent_connection_on_io(con, bytes, p_ol);
} }
} }
int agent_start(BOOL dbg, BOOL child, HANDLE pipe) {
int i; static int
process_connection(HANDLE pipe, int type) {
struct agent_connection* con;
if ((con = malloc(sizeof(struct agent_connection))) == NULL) {
debug("out of memory");
return ERROR_OUTOFMEMORY;
}
memset(con, 0, sizeof(struct agent_connection));
con->connection = pipe;
con->type = type;
CreateIoCompletionPort(pipe, ioc_port, (ULONG_PTR)con, 0);
agent_connection_on_io(con, 0, &con->ol);
iocp_work(NULL);
}
static void
agent_listen_loop() {
DWORD i, r;
HANDLE wait_events[NUM_LISTENERS + 1];
wait_events[0] = event_stop_agent;
for (i = 0; i < NUM_LISTENERS; i++)
wait_events[i + 1] = listeners[i].ol.hEvent;
while (1) {
for (i = 0; i < NUM_LISTENERS; i++) {
if (listeners[i].pipe == INVALID_HANDLE_VALUE) {
listeners[i].pipe = CreateNamedPipeW(
listeners[i].pipe_id, // pipe name
PIPE_ACCESS_DUPLEX | FILE_FLAG_OVERLAPPED, // read/write access
PIPE_TYPE_BYTE | // message type pipe
PIPE_READMODE_BYTE | // message-read mode
PIPE_WAIT, // blocking mode
PIPE_UNLIMITED_INSTANCES, // max. instances
BUFSIZE, // output buffer size
BUFSIZE, // input buffer size
0, // client time-out
&listeners[i].sa);
if (listeners[i].pipe == INVALID_HANDLE_VALUE) {
debug("cannot create listener pipe ERROR:%d", GetLastError());
SetEvent(event_stop_agent);
}
else if (ConnectNamedPipe(listeners[i].pipe, &listeners[i].ol) != FALSE) {
debug("ConnectNamedPipe returned unexpectedly");
SetEvent(event_stop_agent);
}
if (GetLastError() == ERROR_PIPE_CONNECTED)
SetEvent(listeners[i].ol.hEvent);
if (GetLastError() != ERROR_IO_PENDING) {
debug("ConnectNamedPipe failed ERROR: %d", GetLastError());
SetEvent(event_stop_agent);
}
}
}
r = WaitForMultipleObjects(NUM_LISTENERS + 1, wait_events, FALSE, INFINITE);
if (r == WAIT_OBJECT_0) {
//received signal to shutdown
debug("shutting down");
agent_cleanup();
return;
}
else if ((r > WAIT_OBJECT_0) && (r <= (WAIT_OBJECT_0 + NUM_LISTENERS))) {
/* process incoming connection */
HANDLE con = listeners[r - 1].pipe;
listeners[r - 1].pipe = INVALID_HANDLE_VALUE;
if (debug_mode) {
process_connection(con, listeners[r - 1].type);
agent_cleanup();
return;
}
else {
/* todo - spawn a child to take care of this*/
wchar_t path[MAX_PATH], module_path[MAX_PATH];
PROCESS_INFORMATION pi;
STARTUPINFO si;
si.cb = sizeof(STARTUPINFO);
memset(&si, 0, sizeof(STARTUPINFO));
GetModuleFileNameW(NULL, module_path, MAX_PATH);
swprintf_s(path, MAX_PATH, L"%s %d %d", module_path, con, listeners[r - 1].type);
if (CreateProcessW(NULL, path, NULL, NULL, TRUE,
DETACHED_PROCESS, NULL, NULL,
&si, &pi) == FALSE) {
debug("CreateProcess failure: %d", GetLastError());
CloseHandle(con);
agent_cleanup();
return;
}
CloseHandle(con);
}
}
else {
debug("wait on events ended with %d ERROR:%d", r, GetLastError());
agent_cleanup();
return;
}
}
}
void agent_cleanup_connection(struct agent_connection* con) {
debug("connection %p clean up", con);
CloseHandle(con->connection);
free(con);
CloseHandle(ioc_port);
ioc_port = NULL;
}
void agent_shutdown() {
SetEvent(event_stop_agent);
}
int agent_start(BOOL dbg_mode, BOOL child, HANDLE pipe, enum agent_type type) {
int i, r;
HKEY agent_root; HKEY agent_root;
DWORD process_id = GetCurrentProcessId(); DWORD process_id = GetCurrentProcessId();
debug("agent_start pid:%d, dbg:%d, child:%d, pipe:%d", process_id, dbg, child, pipe); debug("agent_start pid:%d, dbg:%d, child:%d, pipe:%d", process_id, dbg_mode, child, pipe);
action_queue = 0; debug_mode = dbg_mode;
list = NULL;
ioc_port = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, (ULONG_PTR)NULL, 0);
for (i = 0; i < 3; i++) if ((ioc_port = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, (ULONG_PTR)NULL, 0)) == NULL) {
QueueUserWorkItem(iocp_work, NULL, 0); debug("cannot create ioc port ERROR:%d", GetLastError());
return GetLastError();
}
if (child == FALSE) {
RegCreateKeyExW(HKEY_LOCAL_MACHINE, SSH_AGENT_ROOT, 0, 0, 0, KEY_WRITE, 0, &agent_root, 0);
RegSetValueExW(agent_root, L"ProcessID", 0, REG_DWORD, (BYTE*)&process_id, 4);
if ((event_stop_agent = CreateEvent(NULL, TRUE, FALSE, NULL)) == NULL)
return GetLastError();
if ((r = init_listeners()) != 0)
return r;
agent_listen_loop();
}
else {
return process_connection(pipe, type);
}
agent_listen(); return 0;
RegCreateKeyExW(HKEY_LOCAL_MACHINE, SSH_AGENT_ROOT, 0, 0, 0, KEY_WRITE, 0, &agent_root, 0);
RegSetValueExW(agent_root, L"ProcessID", 0, REG_DWORD, (BYTE*)&process_id, 4);
iocp_work(NULL);
return 1;
} }

View File

@ -7,6 +7,13 @@
#define SSHD_KEYS_ROOT SSH_ROOT L"\\Keys" #define SSHD_KEYS_ROOT SSH_ROOT L"\\Keys"
#define HEADER_SIZE 4 #define HEADER_SIZE 4
enum agent_type {
KEY_AGENT,
PUBKEY_AGENT,
PUBKEY_AUTH_AGENT
};
struct agent_connection { struct agent_connection {
OVERLAPPED ol; OVERLAPPED ol;
HANDLE connection; HANDLE connection;
@ -23,14 +30,13 @@ struct agent_connection {
WRITING, WRITING,
DONE DONE
} state; } state;
struct agent_connection* next; enum agent_type type;
}; };
void agent_connection_on_io(struct agent_connection*, DWORD, OVERLAPPED*); void agent_connection_on_io(struct agent_connection*, DWORD, OVERLAPPED*);
void agent_connection_on_error(struct agent_connection* , DWORD ); void agent_connection_on_error(struct agent_connection* , DWORD );
void agent_connection_disconnect(struct agent_connection*); void agent_connection_disconnect(struct agent_connection*);
int agent_start(BOOL debug, BOOL child, HANDLE pipe); int agent_start(BOOL, BOOL, HANDLE, enum agent_type);
void agent_shutdown(); void agent_shutdown();
void agent_listen(); void agent_cleanup_connection(struct agent_connection*);
void agent_cleanup_connection(struct agent_connection*);

View File

@ -46,6 +46,7 @@ void agent_connection_on_error(struct agent_connection* con, DWORD error) {
void agent_connection_on_io(struct agent_connection* con, DWORD bytes, OVERLAPPED* ol) { void agent_connection_on_io(struct agent_connection* con, DWORD bytes, OVERLAPPED* ol) {
/* process error */ /* process error */
debug("connection io %p #bytes:%d state:%d", con, bytes, con->state);
if ((bytes == 0) && (GetOverlappedResult(con->connection, ol, &bytes, FALSE) == FALSE)) if ((bytes == 0) && (GetOverlappedResult(con->connection, ol, &bytes, FALSE) == FALSE))
ABORT_CONNECTION_RETURN(con); ABORT_CONNECTION_RETURN(con);
@ -56,7 +57,6 @@ void agent_connection_on_io(struct agent_connection* con, DWORD bytes, OVERLAPPE
{ {
switch (con->state) { switch (con->state) {
case LISTENING: case LISTENING:
agent_listen();
case WRITING: case WRITING:
/* Writing is done, read next request */ /* Writing is done, read next request */
/* assert on assumption that write always completes on sending all bytes*/ /* assert on assumption that write always completes on sending all bytes*/