Kernel Mode Sockets Library

来源:百度文库 编辑:神马文学网 时间:2024/04/19 15:56:15

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__