windows-nt/Source/XPSP1/NT/inetsrv/iis/svcs/loadbal/driver/iislbdr.cxx
2020-09-26 16:20:57 +08:00

1533 lines
38 KiB
C++

/*++
Copyright (c) 1998 Microsoft Corporation
Module Name:
iislbdr.cxx
Abstract:
This module implements NAT entry points
Author:
Philippe Choquier ( phillich )
--*/
extern "C" {
//#include <nt.h>
//#include <ntrtl.h>
//#include <nturtl.h>
//#include <ntsecapi.h>
//#include <ntddnetd.h>
//#include <ntverp.h>
#define NOGDICAPMASKS
#define NOVIRTUALKEYCODES
#define NOWINMESSAGES
#define NOWINSTYLES
#define NOSYSMETRICS
#define NOMENUS
#define NOICONS
#define NOKEYSTATES
#define NOSYSCOMMANDS
#define NORASTEROPS
#define NOSHOWWINDOW
#define OEMRESOURCE
#define NOATOM
#define NOCLIPBOARD
#define NOCOLOR
#define NOCTLMGR
#define NODRAWTEXT
#define NOGDI
#define NOKERNEL
#if defined(KERNEL_MODE)
#define NOUSER
#endif
#define NONLS
#define NOMB
#define NOMEMMGR
#define NOMETAFILE
#define NOMINMAX
#define NOMSG
#define NOOPENFILE
#define NOSCROLL
#define NOSERVICE
#define NOSOUND
#define NOTEXTMETRIC
#define NOWH
#define NOWINOFFSETS
#define NOCOMM
#define NOKANJI
#define NOHELP
#define NOPROFILER
#define NODEFERWINDOWPOS
#if defined(KERNEL_MODE)
#include <strmini.h>
#else
#include <windows.h>
#include <wdbgexts.h>
#endif
#include <stdlib.h>
#include <ipnat.h>
#include "iisnatio.h"
#include <bootexp.hxx>
}
#pragma intrinsic(memcmp)
#include <iislbh.hxx>
#if DBG
#define NOISY_DBG 1
#else
#define NOISY_DBG 0
#endif
//
// Private strutures
//
#pragma pack(1)
typedef struct _HASH_BUCKET
{
ULONG m_HashEntry;
ULONG m_CheckExpiry;
} HASH_BUCKET ;
typedef struct _HASH_ENTRY
{
ULONG m_RemoteAddress;
USHORT m_PublicRef;
USHORT m_PrivateRef;
ULONG m_NextHashEntry;
ULONG m_Expiry;
} HASH_ENTRY;
#pragma pack()
#define NB_BUCKETS 251
#define HASH_ALLOC_GRAIN 64
#define EOCHAIN 0xffff
#define CHECK_EXPIRY 30 // in seconds
//
// Define macros allowing common source code for user & kernel mode
// This allows the mapping to be tested in user mode
//
#if defined(KERNEL_MODE)
#define DECLARE_LOCK_OBJECT(a) KSPIN_LOCK a
#define INIT_LOCK_OBJECT(a) KeInitializeSpinLock(a)
#define TERMINATE_LOCK_OBJECT(a)
#define LockIpMap() KeAcquireSpinLock(&IpMapLock, &oldIrql )
#define UnlockIpMap() KeReleaseSpinLock(&IpMapLock, oldIrql )
#define LockIpMapAtDpcLevel() KeAcquireSpinLock(&IpMapLock, &oldIrql)
#define UnlockIpMapFromDpcLevel() KeReleaseSpinLock(&IpMapLock, oldIrql)
//#define LockIpMapAtDpcLevel() KeAcquireSpinLockAtDpcLevel(&IpMapLock)
//#define UnlockIpMapFromDpcLevel() KeReleaseSpinLockFromDpcLevel(&IpMapLock)
#define AllocBuffer(a) ExAllocatePool( NonPagedPool, a )
#define FreeBuffer(a) ExFreePool( a )
#define Copy_Memory(a,b,c) RtlCopyBytes(a,b,c)
#define Equal_Memory(a,b,c) RtlEqualMemory(a,b,c)
#define DeclareTimeObject(a) ULONGLONG a
#define GetTime(a,b) KeQuerySystemTime( (PLARGE_INTEGER)&a ); b = (ULONG)(a / (ULONGLONG)(10*1000*1000))
#else
#include <time.h>
#define DECLARE_LOCK_OBJECT(a) CRITICAL_SECTION a
#define INIT_LOCK_OBJECT(a)
#define TERMINATE_LOCK_OBJECT(a)
#define LockIpMap()
#define UnlockIpMap()
#define LockIpMapAtDpcLevel()
#define UnlockIpMapFromDpcLevel()
#define AllocBuffer(a) LocalAlloc( LMEM_FIXED, a)
#define FreeBuffer(a) LocalFree( a )
#define Copy_Memory(a,b,c) memcpy(a,b,c)
#define Equal_Memory(a,b,c) 0 // !memcmp(a,b,c)
#define DeclareTimeObject(a)
#define GetTime(a,b) b = time(NULL)
#undef IoCompleteRequest
#define IoCompleteRequest(a,b)
#undef ProbeForRead
#define ProbeForRead(a,b,c)
#define IoDeleteSymbolicLink(a)
#define IoDeleteDevice(a)
#define IoCreateSymbolicLink(a,b) STATUS_SUCCESS
#define IoCreateDevice(a,b,c,d,e,f,g) STATUS_SUCCESS
#define IoGetDeviceObjectPointer(a,b,c,d) (*c=NULL,*d=NULL,STATUS_SUCCESS)
#define IoBuildDeviceIoControlRequest(a,b,c,d,e,f,g,h,i) (PIRP)1
#define ObfDereferenceObject(a)
#include <stdio.h>
#include <stdarg.h>
ULONG
_cdecl
DbgPrint(
PCH Format,
...
)
{
va_list marker;
CHAR achD[128];
va_start( marker, Format );
wvsprintf( achD, Format, marker );
OutputDebugString( achD );
va_end( marker );
return 0;
}
#endif
//
// Private prototypes
//
NTSTATUS
NotifyPort(
CKernelIpEndpointEx* pEndp
);
NTSTATUS
UnnotifyPort(
CKernelIpEndpointEx* pEndp
);
//
// Globals
//
HASH_BUCKET HashBuckets[NB_BUCKETS];
HASH_ENTRY* HashEntries;
unsigned int cHashEntries;
unsigned int iFirstFreeHashEntry;
CKernelIpMapMinHelper IpMap;
ULONG iCurrentServer;
PNAT_HELPER_QUERY_INFO_SESSION pQueryInfoSession;
PNAT_HELPER_DEREGISTER_DIRECTOR pDeregisterNat;
PDEVICE_OBJECT NatDeviceObject;
PFILE_OBJECT NatFileObject;
DECLARE_LOCK_OBJECT(IpMapLock);
//
// NAT functions
//
extern "C"
NTSTATUS
NatRegisterSessionControl(
IN ULONG Version
)
/*++
NatRegisterSessionControl
Routine Description:
This routine is exported by the session-control module.
It is invoked by the NAT to determine whether the module will be
directing incoming sessions.
Arguments:
Version - the version of the NAT which is running.
Return Value:
STATUS_SUCCESS if the module wishes to direct incoming sessions,
error code otherwise.
--*/
{
DeclareTimeObject( RegSes );
ULONG Now;
ULONG i;
GetTime( RegSes, Now );
HashEntries = NULL;
IpMap.SetBuffer( NULL );
iFirstFreeHashEntry = EOCHAIN;
cHashEntries = 0;
iCurrentServer = 0;
for ( i = 0 ; i < NB_BUCKETS ; ++i )
{
HashBuckets[i].m_HashEntry = EOCHAIN;
HashBuckets[i].m_CheckExpiry = Now + CHECK_EXPIRY;
}
INIT_LOCK_OBJECT(&IpMapLock);
pQueryInfoSession = NULL;
pDeregisterNat = NULL;
#if NOISY_DBG
DbgPrint( "NatRegisterSessionControl:\n" );
#endif
return STATUS_SUCCESS;
}
extern "C"
NTSTATUS
NatIoctlSessionControl(
IN PIRP Irp
)
/*++
NatIoctlSessionControl
Routine Description:
It handles IOCTL_IISNATIO_SET_CONFIG IRP, dispatched from IISNATDispatchIoctl()
It is also exported to be called directly when called from user mode
The module should treat the IRP as if it were received in an actual
dispatch routine; the NAT does no pre-processing or post-processing,
so the module must update the IRP on its own.
Arguments:
Irp - the IRP received by the NAT from the I/O manager.
Return Value:
The status code to be returned from the NAT's dispatch routine.
--*/
{
ULONG cNewIpMapSize;
LPVOID NewIpMap;
HASH_ENTRY* pTopHash;
HASH_ENTRY* pHash;
CKernelIpMapMinHelper NewIpMapHelper;
UINT iOld;
UINT iNew;
UINT iHash;
UINT iF;
UINT iPrevF;
UINT iNext;
DWORD iNewPub;
DWORD iNewPrv;
BOOL fUpdate = FALSE;
PIO_STACK_LOCATION IrpSp;
NTSTATUS status = STATUS_SUCCESS;
KIRQL oldIrql;
CKernelIpEndpointEx* pOld;
CKernelIpEndpointEx* pNew;
UINT cPublicIp;
UINT cPrivateIp;
IPREF* pIp;
IPREF* pTopIp;
UINT cSize;
UINT iIpRef;
IrpSp = IoGetCurrentIrpStackLocation( Irp );
cNewIpMapSize = IrpSp->Parameters.DeviceIoControl.InputBufferLength;
NewIpMap = AllocBuffer( cNewIpMapSize );
if ( NewIpMap == NULL )
{
status = STATUS_NO_MEMORY;
goto complete;
}
NewIpMapHelper.SetBuffer( NewIpMap );
//
// Copy over the user's buffer into our information structure.
//
Copy_Memory( NewIpMap,
Irp->AssociatedIrp.SystemBuffer,
cNewIpMapSize );
//
// Configuration validation
//
if ( NewIpMapHelper.GetSize() != cNewIpMapSize )
{
#if NOISY_DBG
DbgPrint( "NatIoctlSessionControl: size mismatch: %u %u\n", NewIpMapHelper.GetSize(), cNewIpMapSize );
#endif
FreeBuffer( NewIpMap );
status = STATUS_INVALID_PARAMETER;
goto complete;
}
cSize = NewIpMapHelper.GetKernelServerDescriptionSize();
pIp = NewIpMapHelper.GetPrivateIpRef( NewIpMapHelper.GetServerPtr( 0 ), 0 );
pTopIp = (IPREF *)((LPBYTE)pIp + cSize * NewIpMapHelper.ServerCount());
cPublicIp = NewIpMapHelper.PublicIpCount();
cPrivateIp = NewIpMapHelper.PrivateIpCount();
//
// Check that private IP references are in range
//
for ( ;
pIp < pTopIp ;
pIp = (IPREF*)(((LPBYTE)pIp) + cSize) )
{
for ( iIpRef = 0 ; iIpRef < cPublicIp ; ++iIpRef )
{
if ( pIp[iIpRef] >= cPrivateIp &&
pIp[iIpRef] != -1 )
{
#if NOISY_DBG
DbgPrint( "NatIoctlSessionControl: invalid private IP ref: %u max is %u\n",
pIp[iIpRef],
cPrivateIp );
#endif
FreeBuffer( NewIpMap );
status = STATUS_INVALID_PARAMETER;
goto complete;
}
}
}
//
// check if existing configuration not empty. If empty then no cache update
// necessary : there is no mapping to update
//
if ( IpMap.GetBuffer() )
{
//
// Check if each unique port in old cnfg is present in new cnfg
// if present in both then nothing to do
// if present in old but not in new then unnotify NAT
// if present in new but not in old then notify NAT
//
for ( iOld = 0 ;
iOld < IpMap.PublicIpCount() ;
++iOld )
{
pOld = IpMap.GetPublicIpPtr( iOld );
//
// need to reset m_dwNotifyPort : might have been set to 0 during previous
// IOCTL
// Only consider if != 0 : 0 means non unique port ( i.e already present in a
// previous entry )
//
if ( pOld->m_dwNotifyPort = pOld->m_usUniquePort )
{
//
// Check if present in new. If yes mark as no notification
// in both old & new configuration
//
for ( iNew = 0 ;
iNew < NewIpMapHelper.PublicIpCount() ;
++iNew )
{
pNew = NewIpMapHelper.GetPublicIpPtr( iNew );
if ( pNew->m_usUniquePort == pOld->m_usUniquePort )
{
pNew->m_dwNotifyPort = 0;
pOld->m_dwNotifyPort = 0;
pNew->m_pvDirectorHandle = pOld->m_pvDirectorHandle;
}
}
//
// If still !=0 then was not found in new configuration,
// so we have to unnotify NAT
//
if ( pOld->m_dwNotifyPort )
{
if ( (status = UnnotifyPort( pOld )) != STATUS_SUCCESS )
{
FreeBuffer( NewIpMap );
goto complete;
}
}
}
}
//
// build old public IP -> new public IP, old prv IP -> new private IP map
// if no match then -1
//
for ( iOld = 0 ; iOld < IpMap.PublicIpCount() ; ++iOld )
{
IpMap.GetPublicIpPtr( iOld )->m_dwIndex = (DWORD)-1;
for ( iNew = 0 ; iNew < NewIpMapHelper.PublicIpCount() ; ++iNew )
{
if ( Equal_Memory( IpMap.GetPublicIpPtr( iOld ),
NewIpMapHelper.GetPublicIpPtr( iNew ),
sizeof(CKernelIpEndpoint) ) )
{
IpMap.GetPublicIpPtr( iOld )->m_dwIndex = iNew;
break;
}
}
//
// If new index not equal to old, then must update cache entries
//
if ( iNew != iOld ||
iNew == NewIpMapHelper.PublicIpCount() )
{
#if NOISY_DBG
DbgPrint( "NatIoctlSessionControl: public IP mismatch: %u/%u %08x:%04x -> %u/%u\n",
iOld, IpMap.PublicIpCount(),
IpMap.GetPublicIpPtr( iOld )->m_dwIpAddress,
IpMap.GetPublicIpPtr( iOld )->m_usPort,
iNew, NewIpMapHelper.PublicIpCount() );
#endif
fUpdate = TRUE;
}
}
for ( iOld = 0 ; iOld < IpMap.PrivateIpCount() ; ++iOld )
{
IpMap.GetPrivateIpEndpoint( iOld )->m_dwIndex = (DWORD)-1;
for ( iNew = 0 ; iNew < NewIpMapHelper.PrivateIpCount() ; ++iNew )
{
//
// Check for identity between old and new private IP address
// We also check reference count on new private IP > 0
// ( or was already 0 for old private IP )
//
if ( Equal_Memory( IpMap.GetPrivateIpEndpoint( iOld ),
NewIpMapHelper.GetPrivateIpEndpoint( iNew ),
sizeof(CKernelIpEndpoint) ) &&
( NewIpMapHelper.GetPrivateIpEndpoint( iNew )->m_dwRefCount ||
!IpMap.GetPrivateIpEndpoint( iOld )->m_dwRefCount) )
{
IpMap.GetPrivateIpEndpoint( iOld )->m_dwIndex = iNew;
break;
}
}
//
// If new index not equal to old, then must update cache entries
//
if ( iNew != iOld ||
iNew == NewIpMapHelper.PrivateIpCount() )
{
#if NOISY_DBG
DbgPrint( "NatIoctlSessionControl: private IP mismatch: %u/%u %08x:%04x -> %u/%u\n",
iOld, IpMap.PrivateIpCount(),
IpMap.GetPrivateIpEndpoint( iOld )->m_dwIpAddress,
IpMap.GetPrivateIpEndpoint( iOld )->m_usPort,
iNew, NewIpMapHelper.PrivateIpCount() );
#endif
fUpdate = TRUE;
}
}
LockIpMap();
if ( fUpdate )
{
#if NOISY_DBG
DbgPrint( "NatIoctlSessionControl: update cache\n" );
#endif
// update cache, if no map for either public or private ref then delete entry
// scan all cache entries for all hash buckets
for ( iHash = 0 ; iHash < NB_BUCKETS ; ++iHash )
{
for ( iF = HashBuckets[iHash].m_HashEntry, iPrevF = EOCHAIN ;
iF != EOCHAIN ;
iF = iNext )
{
iNext = HashEntries[iF].m_NextHashEntry;
iNewPrv = IpMap.GetPrivateIpEndpoint( HashEntries[iF].m_PrivateRef )->m_dwIndex;
iNewPub = IpMap.GetPublicIpPtr( HashEntries[iF].m_PublicRef )->m_dwIndex;
if ( iNewPrv == (DWORD)-1 ||
iNewPub == (DWORD)-1 )
{
// remove it from hash chain
if ( iPrevF != EOCHAIN )
{
HashEntries[iPrevF].m_NextHashEntry = iNext;
}
else
{
HashBuckets[iHash].m_HashEntry = iNext;
}
// insert it in free list
HashEntries[iF].m_NextHashEntry = iFirstFreeHashEntry;
iFirstFreeHashEntry = iF;
}
else
{
HashEntries[iF].m_PrivateRef = (USHORT)iNewPrv;
HashEntries[iF].m_PublicRef = (USHORT)iNewPub;
iPrevF = iF;
}
}
}
}
FreeBuffer( IpMap.GetBuffer() );
IpMap.SetBuffer( NewIpMap );
UnlockIpMap();
}
else
{
LockIpMap();
IpMap.SetBuffer( NewIpMap );
UnlockIpMap();
}
//
// scan new configuration for all m_dwNotifyPort != 0
// which means unique ports not found in old cnfg
// so we have to notify NAT
//
for ( iNew = 0 ;
iNew < NewIpMapHelper.PublicIpCount() ;
++iNew )
{
pNew = NewIpMapHelper.GetPublicIpPtr( iNew );
if ( pNew->m_dwNotifyPort )
{
if ( (status = NotifyPort( pNew )) != STATUS_SUCCESS )
{
goto complete;
}
}
}
iCurrentServer = 0;
complete:
Irp->IoStatus.Information = 0; // nothing to send back to user mode
Irp->IoStatus.Status = status;
IoCompleteRequest( Irp, 0 );
#if NOISY_DBG
DbgPrint( "NatIoctlSessionControl: status=%u\n", status );
#endif
return status;
}
extern "C"
NTSTATUS
NatQuerySessionControl(
IN PVOID DirectorContext,
IN UCHAR Protocol,
IN ULONG PublicAddress,
IN USHORT PublicPort,
IN ULONG RemoteAddress,
IN USHORT RemotePort,
OUT PULONG PrivateAddress,
OUT PUSHORT PrivatePort,
OUT PVOID* DirectorSessionContextp OPTIONAL
)
/*++
NatQuerySessionControl
Routine Description:
This routine is exported by the session-control module.
It is invoked by the NAT to obtain the IP address and port to which
a given incoming session should be directed.
Arguments:
Protocol - NAT_PROTOCOL_TCP or NAT_PROTOCOL_UDP (See IPNAT.H).
PublicAddress - the public address on which the session was received
PublicPort - the public port on which the session was received
RemoteAddress - the remote address from which the session originated
RemotePort - the remote port from which the session originated
PrivateAddress - receives the private address for the session
PrivatePort - receives the private port for the session
Return Value:
STATUS_SUCCESS if the session-control module has stored directions
in 'PrivateAddress' and 'PrivatePort'; failure code otherwise.
--*/
{
ULONG iHash;
ULONG iF;
ULONG iPrevF;
ULONG iNext;
DeclareTimeObject( SesCtl );
ULONG Now;
ULONG cGrain;
int fAtLeastOneServer;
int fAnyServerAvailable;
CKernelIpEndpointEx* pIp;
IPREF BestIpSoFar = (IPREF)-1;
IPREF* pTopIp;
IPREF* pEndIp;
IPREF* pCurrentIp;
CKernelServerDescription* pServer;
CKernelServerDescription* BestServerSoFar = NULL;
ULONG iP;
ULONG cPublicIp;
ULONG cServer;
ULONG iS;
ULONG cSize;
HASH_ENTRY* NewHashEntries;
HASH_ENTRY* pHash;
int BestLoadSoFar = 0;
KIRQL oldIrql;
DWORD dwSticky;
if ( IpMap.GetBuffer() == NULL )
{
return STATUS_ADDRESS_NOT_ASSOCIATED;
}
*DirectorSessionContextp = NULL;
GetTime( SesCtl, Now );
#if NOISY_DBG
DbgPrint( "NatQuerySessionControl: protocol %02x remote %08x:%04x to %08x:%04x time %u\n",
Protocol,
RemoteAddress,RemotePort,
PublicAddress,PublicPort,
Now );
#endif
// hash bucket determined by XORing all bytes of public & remote address
iHash = (PublicAddress ^ RemoteAddress) % NB_BUCKETS;
LockIpMapAtDpcLevel();
// check entries for expiration
if ( HashBuckets[iHash].m_CheckExpiry < Now )
{
for ( iF = HashBuckets[iHash].m_HashEntry, iPrevF = EOCHAIN ;
iF != EOCHAIN ;
iF = iNext )
{
iNext = HashEntries[iF].m_NextHashEntry;
if ( HashEntries[iF].m_Expiry < Now )
{
#if NOISY_DBG
DbgPrint( "NatQuerySessionControl: remove entry %u, expire %u prev link %u\n",
iF,
HashEntries[iF].m_Expiry, iPrevF );
#endif
// remove it from hash chain
if ( iPrevF != EOCHAIN )
{
HashEntries[iPrevF].m_NextHashEntry = iNext;
}
else
{
HashBuckets[iHash].m_HashEntry = iNext;
}
// insert it in free list
HashEntries[iF].m_NextHashEntry = iFirstFreeHashEntry;
iFirstFreeHashEntry = iF;
}
else
{
iPrevF = iF;
}
}
HashBuckets[iHash].m_CheckExpiry = Now + CHECK_EXPIRY;
}
cPublicIp = IpMap.PublicIpCount();
cServer = IpMap.ServerCount();
// get index of public IP/port
for ( iP = 0 ;
iP < cPublicIp ;
++iP )
{
pIp = IpMap.GetPublicIpPtr( iP );
if ( pIp->m_dwIpAddress == PublicAddress &&
pIp->m_usPort == PublicPort )
{
dwSticky = pIp->m_dwSticky;
break;
}
}
#if NOISY_DBG
DbgPrint( "NatQuerySessionControl: %d public, %d servers, index public %d\n",
cPublicIp, cServer, iP );
#endif
if ( iP == cPublicIp )
{
// public IP not found in configuration : fail request
UnlockIpMapFromDpcLevel();
#if NOISY_DBG
DbgPrint( "NatQuerySessionControl: public IP not found, fail request\n" );
#endif
return STATUS_ADDRESS_NOT_ASSOCIATED;
}
// check if RemoteAddress in sticky cache ( also check expiration for this cache line )
// if yes update expiration time, return addr
for ( iF = HashBuckets[iHash].m_HashEntry ;
iF != EOCHAIN ;
iF = pHash->m_NextHashEntry )
{
pHash = HashEntries + iF;
if ( pHash->m_RemoteAddress == RemoteAddress &&
pHash->m_PublicRef == iP )
{
// we have a match : return cached address
pIp = IpMap.GetPrivateIpEndpoint( pHash->m_PrivateRef );
*PrivateAddress = pIp->m_dwIpAddress;
*PrivatePort = pIp->m_usPort;
// update mapping expiration time
pHash->m_Expiry = Now + dwSticky;
#if NOISY_DBG
DbgPrint( "NatQuerySessionControl: found in cache %08x:%04x, sticky %d new expire %u\n",
*PrivateAddress,*PrivatePort,
dwSticky, pHash->m_Expiry );
#endif
UnlockIpMapFromDpcLevel();
return STATUS_SUCCESS;
}
}
// get private IP : find entry with non null private IP, if > 0 then decrease by
// LOADBAL_NORMALIZED_TOTAL / cServer / 10.
// if no valid entry return public IP/port if no load avail entry reset all avail
fAtLeastOneServer = FALSE;
fAnyServerAvailable = FALSE;
//
// If no server then fail request, as at this point we know that the requested
// public IP address is in our list, so access is denied by configuration
//
if ( cServer == 0 )
{
UnlockIpMapFromDpcLevel();
#if NOISY_DBG
DbgPrint( "NatQuerySessionControl: no server avail, fail\n" );
#endif
return STATUS_ADDRESS_NOT_ASSOCIATED;
}
//
// load balancing granularity : based on # servers.
// note that if cGrain > all base load availability for all servers
// then load availability adjustment will loop several times, decreasing
// efficiency.
//
cGrain = LOADBAL_NORMALIZED_TOTAL / cServer / 10;
again:
pServer = IpMap.GetServerPtr( iCurrentServer );
pEndIp = pCurrentIp = IpMap.GetPrivateIpRef( pServer, iP );
pTopIp = IpMap.GetPrivateIpRef( IpMap.GetServerPtr( cServer ), iP );
cSize = IpMap.GetKernelServerDescriptionSize();
// locate entry with highest load available
for ( ;
;
)
{
if ( *pCurrentIp != (IPREF)-1 )
{
// mapping exist from current server to public IP/port
if ( pServer->m_dwLoadAvail )
{
fAtLeastOneServer = TRUE;
// check if available
if ( pServer->m_LoadbalancedLoadAvail > BestLoadSoFar )
{
BestServerSoFar = pServer;
BestIpSoFar = *pCurrentIp;
BestLoadSoFar = pServer->m_LoadbalancedLoadAvail;
fAnyServerAvailable = TRUE;
}
}
}
// advance to next server
pServer = (CKernelServerDescription*)(((LPBYTE)pServer) + cSize);
if ( (pCurrentIp = (IPREF *)(((LPBYTE)pCurrentIp) + cSize )) == pTopIp )
{
pServer = IpMap.GetServerPtr( 0 );
pCurrentIp = IpMap.GetPrivateIpRef( pServer, iP );
}
// check if looked at all servers
if ( pCurrentIp == pEndIp )
{
if ( ++iCurrentServer == cServer )
{
iCurrentServer = 0;
}
break;
}
}
if ( !fAnyServerAvailable )
{
// check if at least one server can map the public IP/port
if ( fAtLeastOneServer )
{
// adjust load balanced availability
for ( iS = 0 ; iS < cServer ; ++iS )
{
pServer = IpMap.GetServerPtr( iS );
pServer->m_LoadbalancedLoadAvail += pServer->m_dwLoadAvail;
}
goto again;
}
// no server for this public IP/port : fail request
UnlockIpMapFromDpcLevel();
#if NOISY_DBG
DbgPrint( "NatQuerySessionControl: no server for this addr avail, fail\n" );
#endif
return STATUS_ADDRESS_NOT_ASSOCIATED;
}
else
{
BestServerSoFar->m_LoadbalancedLoadAvail -= cGrain;
pIp = IpMap.GetPrivateIpEndpoint( BestIpSoFar );
*PrivateAddress = pIp->m_dwIpAddress;
*PrivatePort = pIp->m_usPort;
#if NOISY_DBG
DbgPrint( "NatQuerySessionControl: found %08x:%04x, sticky %d new avail %d grain %d\n",
*PrivateAddress,*PrivatePort,
dwSticky, BestServerSoFar->m_LoadbalancedLoadAvail, cGrain );
#endif
}
//
// insert in sticky cache if sticky duration not null,
// don't check if already there ( dups are OK )
//
if ( dwSticky != 0 )
{
//
// if no free entry need to extend hash table, insert in iFirstFreeHashEntry
//
if ( iFirstFreeHashEntry == EOCHAIN )
{
NewHashEntries = (HASH_ENTRY*)AllocBuffer( sizeof(HASH_ENTRY) * (cHashEntries + HASH_ALLOC_GRAIN) );
if ( NewHashEntries == NULL )
{
UnlockIpMapFromDpcLevel();
#if NOISY_DBG
DbgPrint( "NatQuerySessionControl: out of memory for %u hash entries\n",
cHashEntries );
#endif
return STATUS_NO_MEMORY;
}
if ( HashEntries )
{
Copy_Memory( NewHashEntries, HashEntries, sizeof(HASH_ENTRY) * cHashEntries );
FreeBuffer( HashEntries );
}
HashEntries = NewHashEntries;
//
// insert new entries in free chain
//
iFirstFreeHashEntry = cHashEntries;
for ( iF = cHashEntries, cHashEntries += HASH_ALLOC_GRAIN - 1;
iF < cHashEntries;
++iF
)
{
HashEntries[iF].m_NextHashEntry = iF + 1;
}
HashEntries[cHashEntries++].m_NextHashEntry = EOCHAIN;
}
//
// add new entry to cache
//
iF = iFirstFreeHashEntry;
pHash = HashEntries + iF;
iFirstFreeHashEntry = pHash->m_NextHashEntry;
pHash->m_RemoteAddress = RemoteAddress;
pHash->m_PrivateRef = (USHORT)BestIpSoFar;
pHash->m_PublicRef = (USHORT)iP;
pHash->m_NextHashEntry = HashBuckets[iHash].m_HashEntry;
pHash->m_Expiry = Now + dwSticky;
HashBuckets[iHash].m_HashEntry = iF;
}
UnlockIpMapFromDpcLevel();
return STATUS_SUCCESS;
}
extern "C"
NTSTATUS
IISNATDispatchIoctl(
IN PDEVICE_OBJECT DeviceObject,
IN PIRP Irp
)
/*++
Routine Description:
This function handled IOCTL to the IIS LB device driver.
Arguments:
DeviceObject - Pointer to device object created by the system.
Irp - I/O request packet
Return Value:
NT status
--*/
{
PIO_STACK_LOCATION IrpSp;
NTSTATUS status = STATUS_SUCCESS;
IrpSp = IoGetCurrentIrpStackLocation( Irp );
switch ( IrpSp->Parameters.DeviceIoControl.IoControlCode )
{
case IOCTL_IISNATIO_SET_CONFIG:
status = NatIoctlSessionControl( Irp );
break;
default:
status = STATUS_INVALID_PARAMETER;
}
return status;
}
extern "C"
NTSTATUS
IISNATDispatchCreate(
IN PDEVICE_OBJECT DeviceObject,
IN PIRP Irp
)
/*++
Routine Description:
This function Create for IIS LB device driver.
Arguments:
DeviceObject - Pointer to device object created by the system.
Irp - I/O request packet
Return Value:
NT status
--*/
{
NTSTATUS status = STATUS_SUCCESS;
Irp->IoStatus.Information = 0; // nothing to send back to user mode
Irp->IoStatus.Status = status;
IoCompleteRequest( Irp, IO_NO_INCREMENT );
#if NOISY_DBG
DbgPrint( "IISNATDispatchCreate: status=%u\n", status );
#endif
return status;
}
extern "C"
NTSTATUS
IISNATDispatchClose(
IN PDEVICE_OBJECT DeviceObject,
IN PIRP Irp
)
/*++
Routine Description:
This function Close for IIS LB device driver.
Arguments:
DeviceObject - Pointer to device object created by the system.
Irp - I/O request packet
Return Value:
NT status
--*/
{
NTSTATUS status = STATUS_SUCCESS;
Irp->IoStatus.Information = 0; // nothing to send back to user mode
Irp->IoStatus.Status = status;
IoCompleteRequest( Irp, IO_NO_INCREMENT );
#if NOISY_DBG
DbgPrint( "IISNATDispatchClose: status=%u\n", status );
#endif
return status;
}
extern "C"
VOID
IISNATUnloadDriver(
IN PDRIVER_OBJECT DriverObject
)
/*++
Routine Description:
This function is called when driver is unloaded
Arguments:
DriverObject - Pointer to driver object created by the system.
Return Value:
None
--*/
{
UNICODE_STRING uszDeviceString;
UINT iOld;
CKernelIpEndpointEx* pOld;
IoDeleteDevice( DriverObject->DeviceObject );
RtlInitUnicodeString( &uszDeviceString, L"\\DosDevices\\" LSZDRIVERNAME );
IoDeleteSymbolicLink( &uszDeviceString );
if ( IpMap.GetBuffer() )
{
//
// Unnotify all ports
//
for ( iOld = 0 ;
iOld < IpMap.PublicIpCount() ;
++iOld )
{
pOld = IpMap.GetPublicIpPtr( iOld );
if ( pOld->m_usUniquePort )
{
//
// Ignore error : nothing we can do at this point
//
UnnotifyPort( pOld );
}
}
FreeBuffer( IpMap.GetBuffer() );
}
ObDereferenceObject((PVOID)NatFileObject);
ObDereferenceObject(NatDeviceObject);
if ( HashEntries )
{
FreeBuffer( HashEntries );
}
#if NOISY_DBG
DbgPrint( "IISNATUnloadDriver: unloaded\n" );
#endif
}
extern "C"
NTSTATUS
DriverEntry (
IN PDRIVER_OBJECT DriverObject,
IN PUNICODE_STRING RegistryPath
)
/*++
Routine Description:
This is the initialization routine for the IIS LB device driver.
Arguments:
DriverObject - Pointer to driver object created by the system.
Return Value:
The function value is the final status from the initialization operation.
--*/
{
NTSTATUS ntStatus = STATUS_SUCCESS;
UNICODE_STRING uszDriverString;
UNICODE_STRING uszDeviceString;
UNICODE_STRING NatDeviceString;
PDEVICE_OBJECT pDeviceObject;
//
// Get reference to NAT driver, necessary to allow safe calls to deregister port
//
RtlInitUnicodeString(&NatDeviceString, DD_IP_NAT_DEVICE_NAME);
ntStatus =
IoGetDeviceObjectPointer(
&NatDeviceString,
SYNCHRONIZE|GENERIC_READ|GENERIC_WRITE,
&NatFileObject,
&NatDeviceObject
);
if (!NT_SUCCESS(ntStatus))
{
NatFileObject = NULL;
NatDeviceObject = NULL;
#if NOISY_DBG
DbgPrint( "DriverEntry: IoGetDeviceObjectPointer for NAT failed %08x\n",
ntStatus );
#endif
return ntStatus;
}
ObReferenceObject( NatDeviceObject );
RtlInitUnicodeString( &uszDriverString, L"\\Device\\" LSZDRIVERNAME );
ntStatus = IoCreateDevice( DriverObject,
0,
&uszDriverString,
FILE_DEVICE_UNKNOWN,
0,
FALSE,
&pDeviceObject );
#if NOISY_DBG
DbgPrint( "DriverEntry: IoCreateDevice=%x\n", ntStatus );
#endif
if ( ntStatus != STATUS_SUCCESS )
{
goto cleanup;
}
RtlInitUnicodeString( &uszDeviceString, L"\\DosDevices\\" LSZDRIVERNAME);
ntStatus = IoCreateSymbolicLink( &uszDeviceString, &uszDriverString );
#if NOISY_DBG
DbgPrint( "DriverEntry: IoCreateSymbolicLink=%x\n", ntStatus );
#endif
if ( ntStatus != STATUS_SUCCESS )
{
IoDeleteDevice( pDeviceObject );
goto cleanup;
}
NatRegisterSessionControl( 0 );
DriverObject->DriverUnload = IISNATUnloadDriver;
DriverObject->MajorFunction[IRP_MJ_CREATE] = IISNATDispatchCreate;
DriverObject->MajorFunction[IRP_MJ_CLOSE] = IISNATDispatchClose;
DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = IISNATDispatchIoctl;
cleanup:
if ( ntStatus != STATUS_SUCCESS )
{
ObDereferenceObject((PVOID)NatFileObject);
ObDereferenceObject(NatDeviceObject);
}
return ntStatus;
}
extern "C"
NTSTATUS
NatDeleteSession(
IN PVOID SessionHandle,
IN PVOID DirectorContext,
IN PVOID DirectorSessionContext
)
/*++
Routine Description:
This function is called by NAT for session deletion notification
Arguments:
SessionHandle - NAT session handle
DirectorContext - IIS LB per port context
DirectorSessionContext - IIS LB per session context
Return Value:
NT status
--*/
{
#if NOISY_DBG
DbgPrint( "NatDeleteSession: session %08x\n",
SessionHandle );
#endif
return STATUS_SUCCESS;
}
NTSTATUS
NotifyPort(
CKernelIpEndpointEx* pEndp
)
/*++
Routine Description:
This function notify NAT that load balancing driver wants to handle
specified port
Arguments:
pEndp - endpoint descriptor
Return Value:
NT status
--*/
{
IO_STATUS_BLOCK IoStatusBlock;
PIRP Irp;
NTSTATUS status;
IP_NAT_REGISTER_DIRECTOR RegisterDirector;
RtlZeroMemory(&RegisterDirector, sizeof(RegisterDirector));
RegisterDirector.Version = IP_NAT_VERSION;
RegisterDirector.Flags = 0;
RegisterDirector.Protocol = NAT_PROTOCOL_TCP;
RegisterDirector.Port = pEndp->m_usUniquePort;
RegisterDirector.DirectorContext = NULL;
RegisterDirector.QueryHandler = NatQuerySessionControl;
RegisterDirector.DeleteHandler = NatDeleteSession;
#if defined(KERNEL_MODE)
Irp =
IoBuildDeviceIoControlRequest(
IOCTL_IP_NAT_REGISTER_DIRECTOR,
NatDeviceObject,
(PVOID)&RegisterDirector,
sizeof(RegisterDirector),
(PVOID)&RegisterDirector,
sizeof(RegisterDirector),
FALSE,
NULL,
&IoStatusBlock
);
if (!Irp)
{
status = STATUS_UNSUCCESSFUL;
}
else
{
status = IoCallDriver(NatDeviceObject, Irp);
#if NOISY_DBG
if ( !NT_SUCCESS(status) )
{
DbgPrint( "NotifyPort: IoCallDriver failed %08x\n",
status );
}
#endif
}
if ( NT_SUCCESS(status) )
{
pEndp->m_pvDirectorHandle = RegisterDirector.DirectorHandle;
pQueryInfoSession = RegisterDirector.QueryInfoSession;
pDeregisterNat = RegisterDirector.Deregister;
}
#else
pEndp->m_pvDirectorHandle = (LPVOID)0x12345678;
status = STATUS_SUCCESS;
#endif
#if NOISY_DBG
DbgPrint( "NotifyPort: %u, status %x handle %08x\n", pEndp->m_usUniquePort, status, pEndp->m_pvDirectorHandle );
#endif
return status;
}
NTSTATUS
UnnotifyPort(
CKernelIpEndpointEx* pEndp
)
/*++
Routine Description:
This function notify NAT that load balancing driver does not want to handle
specified port anymore
Arguments:
pEndp - endpoint descriptor
Return Value:
NT status
--*/
{
NTSTATUS status;
if ( pDeregisterNat )
{
status = pDeregisterNat( pEndp->m_pvDirectorHandle );
}
else
{
//
// nothing to unregister in this case, so report success
//
status = STATUS_SUCCESS; // STATUS_INVALID_PARAMETER;
}
#if NOISY_DBG
DbgPrint( "UnnotifyPort: %u, handle %08x Deregister %08x status %x\n",
pEndp->m_usUniquePort,
pEndp->m_pvDirectorHandle,
pDeregisterNat,
status );
#endif
return status;
}