//
//	InterlockedExecute.c
//
//	Execute code with exclusive access to system, with 
//  	all CPUs locked down. 
//
//	** Experimental code not fully tested **
//
//	Code assumes that KeQueryPerformanceCounter returns a
//  	true tick-count per CPU rather than a system-wide value
//
//	Make sure your callback function is prototyped as:
//
//		ULONG UserCallback(PVOID Param);
//
//	Call InterlockedExecute as follows:
//
//		InterlockedExecute(HIGH_LEVEL, UserCallback, userParam, &retVal);
//
#include <ntddk.h>

typedef ULONG (* USER_ROUTINE)(PVOID Param);

#define TIMEOUT (3000)	// 3 second timeout!!

typedef struct
{
	LONG				counter;
	LONG				exitLock;
	KIRQL				irqlLevel;

	LONG				funcStarted;
	LONG				funcTimeout;

	USER_ROUTINE		userFunc;
	PVOID				userParam;
	ULONG				userReturn;

	KDPC				dpcList   [MAXIMUM_PROCESSORS];
	KEVENT				eventList [MAXIMUM_PROCESSORS];

} SYNC_LIST, *PSYNC_LIST;

//
//	Return if specified time (in milliseconds) has passed
//
static BOOLEAN TimeExceeded(LARGE_INTEGER *ticks1, ULONG timeout)
{
	LARGE_INTEGER ticks2, freq;

	ticks2 = KeQueryPerformanceCounter(&freq); 

	if((ticks2.QuadPart - ticks1->QuadPart) * 1000 / freq.QuadPart >= timeout)
	{
		return TRUE;
	}
	else
	{
		return FALSE;
	}
}

//
//	Deferred Procedure, executes at DISPATCH_LEVEL therefore
//  cannot be interrupted by the dispatcher
//
static void SyncDpcProc
	(
		PKDPC		Dpc, 
		PSYNC_LIST	SyncList, 
		PKEVENT		ExitEvent, 
		PVOID		Unused
	)
{
	KIRQL oldIrql;
	
	KeRaiseIrql(SyncList->irqlLevel, &oldIrql);

	// The last DPC to be scheduled gets to execute the user-code
	if(InterlockedDecrement(&SyncList->counter) == 0)
	{
		// Signal that we have started executing, to prevent DPCs 
		// from exiting too early due to time-outs
		InterlockedExchange(&SyncList->funcStarted, 1);

		// if a DPC has timed-out then we must assume it has exited, therefore
		// the system is no longer locked and we must also exit
		if(SyncList->funcTimeout == 0)
		{
			// All the DPCs must be executing - so call the user function!
			SyncList->userReturn = SyncList->userFunc(SyncList->userParam);
		}
		
		InterlockedExchange(&SyncList->exitLock, 0);
	}
	// we weren't the last DPC so just wait until we are told to exit
	else
	{
		LARGE_INTEGER startTime;

		// record current cpu-tick-count
		startTime = KeQueryPerformanceCounter(NULL);

		// Try to acquire the exit lock
		while(InterlockedExchange(&SyncList->exitLock, 1))
		{
			// if we exceed our timeout period and the 
			// user-function hasn't yet started, break out and exit
			if(TimeExceeded(&startTime, TIMEOUT) && !SyncList->funcStarted)
			{
				InterlockedExchange(&SyncList->funcTimeout, 1);
				break;
			}
		}

		// Release the lock for the next thread so it can exit as well
		InterlockedExchange(&SyncList->exitLock, 0);
	}

	KeLowerIrql(oldIrql);

	// signal the main thread that we are exiting!
	KeSetEvent(ExitEvent, 0, FALSE);
}

//
//	InterlockedExecute
//
//	Execute the specified user function at minimum of DISPATCH_LEVEL, with 
//  exclusive access to the current processor (all other CPUs are locked)
//
NTSTATUS InterlockedExecute
	(
		IN	KIRQL				IrqlLevel, 
		IN	USER_ROUTINE		UserFunction, 
		IN	PVOID				UserParam, 
		OUT	PULONG				UserReturn	 OPTIONAL
	)
{
	PVOID		waitList[MAXIMUM_PROCESSORS];
	KIRQL		oldIrql;
	ULONG		numProcessors = KeNumberProcessors;
	ULONG		i;
	
	SYNC_LIST	syncList = 
	{ 
		numProcessors, 
		1,					// acquire the exit-lock 
		IrqlLevel,
		0,
		0,
		UserFunction, 
		UserParam, 
		0 
	};

	//
	// Paranoid parameter validation
	//
	if(UserFunction == 0 || numProcessors >= MAXIMUM_PROCESSORS)
		return STATUS_INVALID_PARAMETER;

	if(KeGetCurrentIrql() != PASSIVE_LEVEL || IrqlLevel < DISPATCH_LEVEL)
		return STATUS_INVALID_PARAMETER;

	//
	// Setup at PASSIVE_LEVEL
	//
	for(i = 0; i < numProcessors; i++)
	{
		// Setup an Event for each DPC to signal when it exits
		KeInitializeEvent(&syncList.eventList[i], SynchronizationEvent, FALSE);
		waitList[i]     = &syncList.eventList[i];

		// Initialize a DPC for each processor
		KeInitializeDpc			(&syncList.dpcList[i], SyncDpcProc, &syncList);
		KeSetImportanceDpc		(&syncList.dpcList[i], HighImportance);
		KeSetTargetProcessorDpc	(&syncList.dpcList[i], (CHAR)i);
	}

	//
	// Scheduling at DISPATCH_LEVEL, so that no DPCs get executed until
	// we drop back to passive - this prevents the current CPU from executing
	// any DPCs and protects us against deadlocking ourselves until all DPCs
	// have been set up
	//
	oldIrql = KeRaiseIrqlToDpcLevel();

	for(i = 0; i < numProcessors; i++)
	{
		KeInsertQueueDpc(&syncList.dpcList[i], &syncList.eventList[i], 0);	
	}

	KeLowerIrql(oldIrql);

	//
	// wait for all DPC routines to exit. KeFlushQueuedDpcs would also work
	// but doesn't seem to exist on NT->XP
	//
	KeWaitForMultipleObjects(	numProcessors, 
								waitList, 
								WaitAll, 
								Executive, 
								KernelMode, 
								FALSE, 0, 0
							);

	//
	// store return value if requested and return to caller
	//
	if(UserReturn != NULL)
		*UserReturn = syncList.userReturn;

	return syncList.funcTimeout ? STATUS_TIMEOUT : STATUS_SUCCESS;
}