驱动中获取域名的IP地址

/*
功能：获取域名的IP地址。

made by correy
made at 2017.06.01
http://correy.webs.com
*/


#include <ntddk.h>
#include <wsk.h>
#include <ws2def.h>


#pragma warning(disable:4201) // nameless struct/union
#pragma warning(disable:4214) // bit field types other than int
#pragma warning(disable:4100) // 未引用的形参
#pragma warning(disable:4101) // 未引用的局部变量
#pragma warning(disable:4189) // 局部变量已初始化但不引用


const WSK_CLIENT_DISPATCH WskAppDispatch = {
  MAKE_WSK_VERSION(1,0), // Use WSK version 1.0
  0,    // Reserved
  NULL  // WskClientEvent callback not required for WSK version 1.0
};


WSK_REGISTRATION WskRegistration;


NTSTATUS
SyncIrpCompletionRoutine(
    __in PDEVICE_OBJECT Reserved,
    __in PIRP Irp,
    __in PVOID Context
    )
{    
    PKEVENT compEvent = (PKEVENT)Context;
    UNREFERENCED_PARAMETER(Reserved);
    UNREFERENCED_PARAMETER(Irp);
    KeSetEvent(compEvent, 2, FALSE);    
    return STATUS_MORE_PROCESSING_REQUIRED;
}


NTSTATUS
KernelNameResolutionSample(
    __in PCWSTR NodeName,
    __in_opt PCWSTR ServiceName,
    __in_opt PADDRINFOEXW Hints,
    __in PWSK_PROVIDER_NPI WskProviderNpi
    )
    //https://docs.microsoft.com/en-us/windows-hardware/drivers/network/resolving-host-names-and-ip-addresses
{
    NTSTATUS status;
    PIRP irp;
    KEVENT completionEvent;
    UNICODE_STRING uniNodeName, uniServiceName, *uniServiceNamePtr;
    PADDRINFOEXW results;
    SOCKADDR_IN * psi = NULL;
    wchar_t buffer[64] = {0};
    PWSTR p = NULL;
    UNICODE_STRING ip = {0};
    const struct in_addr * temp = NULL;

    PAGED_CODE();
    
    RtlInitUnicodeString(&uniNodeName, NodeName);// Initialize UNICODE_STRING structures for NodeName and ServiceName 

    if(ServiceName == NULL) {
        uniServiceNamePtr = NULL;
    }
    else {
        RtlInitUnicodeString(&uniServiceName, ServiceName);
        uniServiceNamePtr = &uniServiceName;
    }
    
    KeInitializeEvent(&completionEvent, SynchronizationEvent, FALSE);// Use an event object to synchronously wait for the WskGetAddressInfo request to be completed. 

    // Allocate an IRP for the WskGetAddressInfo request, and set the IRP completion routine, which will signal the completionEvent when the request is completed.
    irp = IoAllocateIrp(1, FALSE);
    if(irp == NULL) {
        return STATUS_INSUFFICIENT_RESOURCES;
    }        

    IoSetCompletionRoutine(irp, SyncIrpCompletionRoutine, &completionEvent, TRUE, TRUE, TRUE);

    // Make the WskGetAddressInfo request.
    WskProviderNpi->Dispatch->WskGetAddressInfo (
        WskProviderNpi->Client,
        &uniNodeName,
        uniServiceNamePtr,
        NS_ALL,
        NULL, // Provider
        Hints,
        &results, 
        NULL, // OwningProcess
        NULL, // OwningThread
        irp);

    // Wait for completion.
    // Note that processing of name resolution results can also be handled directly within the IRP completion routine,
    // but for simplicity, this example shows how to wait synchronously for completion.
    KeWaitForSingleObject(&completionEvent, Executive, KernelMode, FALSE, NULL);
    status = irp->IoStatus.Status;
    IoFreeIrp(irp);
    if(!NT_SUCCESS(status)) {
        return status;
    }

    // Process the name resolution results by iterating through the addresses within the returned ADDRINFOEXW structure.
    //results; // your code here
    psi = (SOCKADDR_IN *)results->ai_addr;//注意：这一行很重要。有时要该为IPV6版本的。
    temp = (const struct in_addr *)&psi->sin_addr;
    p = RtlIpv4AddressToString(temp, buffer);
    ASSERT(p);
    RtlInitUnicodeString(&ip, buffer);
    KdPrint(("ipv4:%wZ.\r\n", &ip));
    
    WskProviderNpi->Dispatch->WskFreeAddressInfo(WskProviderNpi->Client, results);// Release the returned ADDRINFOEXW structure when no longer needed.

    return status;
} 


VOID Unload(_In_ PDRIVER_OBJECT DriverObject)
{  
    UNREFERENCED_PARAMETER(DriverObject);

    PAGED_CODE();

    WskDeregister(&WskRegistration);
}


NTSTATUS DriverEntry(_In_ PDRIVER_OBJECT DriverObject, _In_ PUNICODE_STRING RegistryPath)
{
    NTSTATUS Status = STATUS_SUCCESS;
    WSK_CLIENT_NPI wskClientNpi;
    WSK_PROVIDER_NPI wskProviderNpi;
    //UNICODE_STRING test  = RTL_CONSTANT_STRING(L"DESKTOP-SQRJ1QU");
    
    UNREFERENCED_PARAMETER(RegistryPath);

    PAGED_CODE();

    KdBreakPoint();

    DriverObject->DriverUnload = Unload;
    
    wskClientNpi.ClientContext = NULL;
    wskClientNpi.Dispatch = &WskAppDispatch;
    Status = WskRegister(&wskClientNpi, &WskRegistration);
    ASSERT(NT_SUCCESS(Status));

    Status = WskCaptureProviderNPI(&WskRegistration, WSK_INFINITE_WAIT, &wskProviderNpi);
    ASSERT(NT_SUCCESS(Status));

    /*
    做一些事情。
    */
    Status = KernelNameResolutionSample(L"www.baidu.com", NULL, NULL, &wskProviderNpi);
    ASSERT(NT_SUCCESS(Status));

    WskReleaseProviderNPI(&WskRegistration);

    return Status;
}