ExRegisterExtension.c

NTSTATUS ExRegisterExtension(_Outptr_ PEX_EXTENSION *Extension, _In_ ULONG RegistrationVersion, _In_ PREGISTRATION_INFO RegistrationInfo)
{
	// Validate that version is ok and that FunctionTable is not sent without FunctionCount or vise-versa.
	if ( (RegistrationVersion & 0xFFFF0000 != 0x10000) || (RegistrationInfo->FunctionTable == nullptr && RegistrationInfo->FunctionCount != 0) ) 
	{ 
		return STATUS_INVALID_PARAMETER; 
	}

	// Skipping over some lock-related stuff,
	
	// Find the host with the matching version and id.
	PHOST_LIST_ENTRY pHostListEntry;
	pHostListEntry = ExpFindHost(RegistrationInfo->ExtensionId, RegistrationInfo->ExtensionVersion);

	// More lock-related stuff.	

	if (!pHostListEntry)
	{
		return STATUS_NOT_FOUND;
	}

	// Verify that the FunctionCount in the host doesn't exceed the FunctionCount supplied by the caller.
	if (RegistrationInfo->FunctionCount < pHostListEntry->FunctionCount)
	{
		ExpDereferenceHost(pHostListEntry);
		return STATUS_INVALID_PARAMETER;
	}

	// Check that the number of functions in FunctionTable matches the amount in FunctionCount.
	PVOID FunctionTable = RegistrationInfo->FunctionTable;
	for (int i = 0; i < RegistrationInfo->FunctionCount; i++)
	{	
		if ( RegistrationInfo->FunctionTable[i] == nullptr ) 
		{ 
			ExpDereferenceHost(pHostListEntry);
			return STATUS_ACCESS_DENIED; 
		}
	}

	// skipping over some more lock-related stuff

	// Check if there is already an extension registered for this host.
	if (pHostListEntry->FunctionTable != nullptr || FlagOn(pHostListEntry->Flags, 1) )
	{
		// There is something related to locks here
		ExpDereferenceHost(pHostListEntry);
		return STATUS_OBJECT_NAME_COLLISION;
	}

	// If there is a callback function for this host, call it before registering the extension, with 0 as the first parameter.
	if (pHostListEntry->FunctionAddress) 
	{ 
		pHostListEntry->FunctionAddress(0, pHostListEntry->ArgForFunction); 
	}

	// Set the FunctionTable in the host to the table supplied by the caller, or to MmBadPointer if a table wasn't supplied.
	if (RegistrationInfo->FunctionTable == nullptr)
	{
		pHostListEntry->FunctionTable = nt!MmBadPointer;
	}
	else
	{
		pHostListEntry->FunctionTable = RegistrationInfo->FunctionTable;
	}

	pHostListEntry->RundownRef = 0;

	// If there is a callback function for this host, call it after registering the extension, with 1 as the first parameter.
	if (pHostListEntry->FunctionAddress)
	{
		pHostListEntry->FunctionAddress(1, pHostListEntry->ArgForFunction);
	}

	// Here there is some more lock-related stuff

	// Set the HostTable of the calling driver to the table of functions listed in the host.
	if (RegistrationInfo->HostTable != nullptr)
	{
		*(PVOID)RegistrationInfo->HostTable = pHostListEntry->hostInterface;
	}

	// Return the initialized host to the caller in the output Extension parameter.
	*Extension = pHostListEntry;
	return STATUS_SUCCESS;
}