#include <windows.h>
#include <stdio.h>
#include <stdlib.h>
#include "thread.h"
#ifndef _WIN32
#include <pthread.h>

int __futex_wait(volatile void *ftx, int val, const struct timespec *timeout);
int __futex_wake(volatile void *ftx, int count);

#include <time.h>
#include <signal.h>

#endif

// thread.c contains wrappers for the primitives of locks, events and threads for use in
// the multithreaded meterpreter. This is the win32/win64 implementation.

/*****************************************************************************************/

/*
 * Create a new lock. We choose Mutex's over CriticalSections as their appears to be an issue
 * when using CriticalSections with OpenSSL on some Windows systems. Mutex's are not as optimal
 * as CriticalSections but they appear to resolve the OpenSSL deadlock issue.
 */
LOCK * lock_create( VOID )
{
	LOCK * lock = (LOCK *)malloc( sizeof( LOCK ) );
	if( lock != NULL )
	{
		memset( lock, 0, sizeof( LOCK ) );

#ifdef _WIN32
		lock->handle = CreateMutex( NULL, FALSE, NULL );
#else
		pthread_mutex_init(lock->handle, NULL);
#endif
	}
	return lock;
}

/*
 * Destroy a lock that is no longer required.
 */
VOID lock_destroy( LOCK * lock )
{
	if( lock != NULL  )
	{
		lock_release( lock );

#ifdef _WIN32
		CloseHandle( lock->handle );
#else
		pthread_mutex_destroy(lock->handle);
#endif

		free( lock );
	}
}

/*
 * Acquire a lock and block untill it is acquired.
 */
VOID lock_acquire( LOCK * lock )
{
	if( lock != NULL  ) {
#ifdef _WIN32
		WaitForSingleObject( lock->handle, INFINITE );
#else
		pthread_mutex_lock(lock->handle);
#endif
	}
}

/*
 * Release a lock previously held.
 */
VOID lock_release( LOCK * lock )
{
	if( lock != NULL  ) {
#ifdef _WIN32
		ReleaseMutex( lock->handle );
#else
		pthread_mutex_unlock(lock->handle);
#endif
	}
}

/*****************************************************************************************/

/*
 * Create a new event which can be signaled/polled/and blocked on.
 */
EVENT * event_create( VOID )
{
	EVENT * event = NULL;

	event = (EVENT *)malloc( sizeof( EVENT ) );
	if( event == NULL )
		return NULL;

	memset( event, 0, sizeof( EVENT ) );

#ifdef _WIN32
	event->handle = CreateEvent( NULL, FALSE, FALSE, NULL );
	if( event->handle == NULL )
	{
		free( event );
		return NULL;
	}
#endif

	return event;
}

/*
 * Destroy an event.
 */
BOOL event_destroy( EVENT * event )
{
	if( event == NULL )
		return FALSE;

#ifdef _WIN32
	CloseHandle( event->handle );
#endif

	free( event );

	return TRUE;
}

/*
 * Signal an event.
 */
BOOL event_signal( EVENT * event )
{
	if( event == NULL )
		return FALSE;

#ifdef _WIN32
	//dprintf( "Signalling 0x%x", event->handle );
	if( SetEvent( event->handle ) == 0 ) {
		//dprintf( "Signalling 0x%x failed %u", event->handle, GetLastError() );
		return FALSE;
	}
#else
	event->handle = (HANDLE)1;
	__futex_wake(&(event->handle), 1);
#endif

	return TRUE;
}

/*
 * Poll an event to see if it has been signaled. Set timeout to -1 to block indefinatly.
 * If timeout is 0 this function does not block but returns immediately.
 */
BOOL event_poll( EVENT * event, DWORD timeout )
{
#ifdef _WIN32
	if( event == NULL )
		return FALSE;

	if( WaitForSingleObject( event->handle, timeout ) == WAIT_OBJECT_0 )
		return TRUE;

	return FALSE;
#else
	BOOL result = FALSE;

	// DWORD WINAPI WaitForSingleObject(
	// __in  HANDLE hHandle,
	// __in  DWORD dwMilliseconds
	// );
	// http://msdn.microsoft.com/en-us/library/ms687032(VS.85).aspx

	if( event == NULL )
		return FALSE;

	if(timeout) {
		struct timespec ts;

		// XXX, need to verify for -1. below modified from bionic/pthread.c
		// and maybe loop if needed ;\

		ts.tv_sec = timeout / 1000;
		ts.tv_nsec = (timeout%1000)*1000000;
		if (ts.tv_nsec >= 1000000000) {
			ts.tv_sec++;
			ts.tv_nsec -= 1000000000;
		}

		// atomically checks if event->handle is 0, if so,
		// it sleeps for timeout. if event->handle is 1, it
		// returns straight away.

		__futex_wait(&(event->handle), 0, &ts);
	}

	// We should behave like an auto-reset event
	result = event->handle ? TRUE : FALSE;
	if( result )
		event->handle = (HANDLE)0;

	return result;
#endif
}

/*****************************************************************************************/

/*
 * Opens and create a THREAD item for the current/calling thread.
 */
THREAD * thread_open( VOID )
{
	THREAD * thread        = NULL;
#ifdef _WIN32
	OPENTHREAD pOpenThread = NULL;
	HMODULE hKernel32      = NULL;


	thread = (THREAD *)malloc( sizeof( THREAD ) );
	if( thread != NULL )
	{
		memset( thread, 0, sizeof(THREAD) );

		thread->id      = GetCurrentThreadId();
		thread->sigterm = event_create();

		// Windows specific process of opening a handle to the current thread which
		// works on NT4 up. We only want THREAD_TERMINATE|THREAD_SUSPEND_RESUME access
		// for now.

		// First we try to use the normal OpenThread function, available on Windows 2000 and up...
		hKernel32 = LoadLibrary( "kernel32.dll" );
		pOpenThread = (OPENTHREAD)GetProcAddress( hKernel32, "OpenThread" );
		if( pOpenThread )
		{
			thread->handle = pOpenThread( THREAD_TERMINATE|THREAD_SUSPEND_RESUME, FALSE, thread->id );
		}
		else
		{
			NTOPENTHREAD pNtOpenThread = NULL;
			// If we can't use OpenThread, we try the older NtOpenThread function as found on NT4 machines.
			HMODULE hNtDll = LoadLibrary( "ntdll.dll" );
			pNtOpenThread = (NTOPENTHREAD)GetProcAddress( hNtDll, "NtOpenThread" );
			if( pNtOpenThread )
			{
				_OBJECT_ATTRIBUTES oa = {0};
				_CLIENT_ID cid        = {0};

				cid.UniqueThread = (PVOID)thread->id;

				pNtOpenThread( &thread->handle, THREAD_TERMINATE|THREAD_SUSPEND_RESUME, &oa, &cid );
			}

			FreeLibrary( hNtDll );
		}

		FreeLibrary( hKernel32 );
	}

	return thread;
#else
	thread = (THREAD *)malloc( sizeof( THREAD ) );

	if( thread != NULL )
	{
		memset( thread, 0, sizeof(THREAD) );

		thread->id      = gettid();
		thread->sigterm = event_create();
		thread->pid	= pthread_self();
	}
	return thread;
#endif
}

#ifndef _WIN32

struct thread_conditional {
	pthread_mutex_t suspend_mutex;
	pthread_cond_t suspend_cond;
	int engine_running;
	LPVOID (*funk)(void *arg);
	THREAD *thread;
};

void __thread_cancelled(int signo)
{
	signal(SIGTERM, SIG_DFL);
	pthread_exit(NULL);
}

/*
 * This is the entry point for threads created with thread_create.
 *
 * To implement suspended threads, we need to do some messing around with
 * mutexes and conditional broadcasts ;\
 */

void *__paused_thread(void *req)
{
	LPVOID (*funk)(void *arg);
	THREAD *thread;

	struct thread_conditional *tc = (struct thread_conditional *)(req);
	tc->thread->id = gettid();

	signal(SIGTERM, __thread_cancelled);

	pthread_mutex_lock(&tc->suspend_mutex);

	while(tc->engine_running == FALSE) {
		pthread_cond_wait(&tc->suspend_cond, &tc->suspend_mutex);
	}

	pthread_mutex_unlock(&tc->suspend_mutex);

	funk = tc->funk;
	thread = tc->thread;
	free(tc);

	if(event_poll(thread->sigterm, 0) == TRUE) {
		/*
		 * In some cases, we might want to stop a thread before it does anything :/
		 */
		return NULL;
	}

	return funk(thread);
}
#endif

/*
 * Create a new thread in a suspended state.
 */
THREAD * thread_create( THREADFUNK funk, LPVOID param1, LPVOID param2, LPVOID param3 )
{
	THREAD * thread = NULL;

	if( funk == NULL )
		return NULL;

	thread = (THREAD *)malloc( sizeof( THREAD ) );
	if( thread == NULL )
		return NULL;

	memset( thread, 0, sizeof( THREAD ) );

	thread->sigterm = event_create();
	if( thread->sigterm == NULL )
	{
		free( thread );
		return NULL;
	}


	thread->parameter1 = param1;
	thread->parameter2 = param2;
	thread->parameter3 = param3;

#ifdef _WIN32
	thread->handle = CreateThread( NULL, 0, funk, thread, CREATE_SUSPENDED, &thread->id );

	if( thread->handle == NULL )
	{
		event_destroy( thread->sigterm );
		free( thread );
		return NULL;
	}

#else
	// PKS, this is fucky.
	// we need to use conditionals to implement this.

	thread->thread_started = FALSE;

	do {
		pthread_t pid;

		struct thread_conditional *tc;
		tc = (struct thread_conditional *) malloc(sizeof(struct thread_conditional));

		if( tc == NULL ) {
			event_destroy(thread->sigterm);
			free(thread);
			return NULL;
		}

		memset( tc, 0, sizeof(struct thread_conditional));

		pthread_mutex_init(&tc->suspend_mutex, NULL);
		pthread_cond_init(&tc->suspend_cond, NULL);

		tc->funk = funk;
		tc->thread = thread;

		thread->suspend_thread_data = (void *)(tc);

		if(pthread_create(&(thread->pid), NULL, __paused_thread, tc) == -1) {
			free(tc);
			event_destroy(thread->sigterm);
			free(thread);
			return NULL;
		}
		// __paused_thread free's the allocated memory.

	} while(0);
#endif

	return thread;
}

/*
 * Run a thread.
 */
BOOL thread_run( THREAD * thread )
{
	if( thread == NULL )
		return FALSE;

#ifdef _WIN32
	if( ResumeThread( thread->handle ) < 0 )
		return FALSE;

#else
	struct thread_conditional *tc;
	tc = (struct thread_conditional *)thread->suspend_thread_data;
	pthread_mutex_lock(&tc->suspend_mutex);
	tc->engine_running = TRUE;
	pthread_mutex_unlock(&tc->suspend_mutex);
	pthread_cond_signal(&tc->suspend_cond);

	thread->thread_started = TRUE;
#endif
	return TRUE;
}

/*
 * Signals the thread to terminate. It is the responsibility of the thread to wait for and process this signal.
 * Should be used to signal the thread to terminate.
 */
BOOL thread_sigterm( THREAD * thread )
{
	BOOL ret;

	if( thread == NULL )
		return FALSE;

	ret = event_signal( thread->sigterm );

#ifndef _WIN32
	/*
	 * If we sig term a thread before it's started execution, we will leak memory / not be
	 * able to join on the thread, etc.
	 *
	 * Therefore, we need to start the thread executing before calling thread_join
	 */
	if(thread->thread_started != TRUE) {
		thread_run(thread);
	}
#endif

	return ret;
}

/*
 * Terminate a thread. Use with caution! better to signal your thread to terminate and wait for it to do so.
 */
BOOL thread_kill( THREAD * thread )
{
	if( thread == NULL )
		return FALSE;

#ifdef _WIN32
	if( TerminateThread( thread->handle, -1 ) == 0 )
		return FALSE;

	return TRUE;
#else
	// bionic/libc/bionic/CAVEATS
	// - pthread cancellation is *not* supported. this seemingly simple "feature" is the source
	// of much bloat and complexity in a C library. Besides, you'd better write correct
	// multi-threaded code instead of relying on this stuff.

	// pthread_kill says: Note  that  pthread_kill()  only  causes  the
	// signal to be handled in the context of the given thread; the signal
	// action (termination or stopping) affects the process as a whole.

	// We send our thread a SIGTERM, and a signal handler calls pthread_exit().

	pthread_kill(thread->id, SIGTERM);
	return FALSE;
#endif
}


/*
 * Blocks untill the thread has terminated.
 */
BOOL thread_join( THREAD * thread )
{
	if( thread == NULL )
		return FALSE;

#ifdef _WIN32
	if( WaitForSingleObject( thread->handle, INFINITE ) == WAIT_OBJECT_0 )
		return TRUE;

	return FALSE;
#else
	if(pthread_join(thread->pid, NULL) == 0)
		return TRUE;

	return FALSE;
#endif
}

/*
 * Destroys a previously created thread. Note, this does not terminate the thread. You must signal your
 * thread to terminate and wait for it to do so (via thread_signal/thread_join).
 */
BOOL thread_destroy( THREAD * thread )
{
	if( thread == NULL )
		return FALSE;

	event_destroy( thread->sigterm );

#ifdef _WIN32
	CloseHandle( thread->handle );
#else
	pthread_detach(thread->pid);
#endif

	free( thread );

	return TRUE;
}