libuv #3285

gethostnamew_nt60_by_kampeador.c

/* Fixed compatibility issue with windows 7 #3285 (kampeador) https://github.com/libuv/libuv/pull/3285 */

#pragma comment(lib, "ws2_32.lib")

#include <assert.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <wchar.h>

#include <winsock2.h>
#include <windows.h>
#include <svcguid.h>

/* Local buffer size for WSAQUERYSETW data inside uv__gethostnamew_nt60
   sizeof(WSAQUERYSETW) + 512 = 632 bytes to match GetHostNameW behavior */
#define WSAQ_LOCAL_BUF_LEN (sizeof(WSAQUERYSETW) + 512)

/* Parameters for WSAQUERYSETW inside uv__gethostnamew_nt60 */
static GUID guid_host_name = SVCID_HOSTNAME;
static AFPROTOCOLS af_protocols[2] = { {AF_INET, IPPROTO_UDP},
                                       {AF_INET, IPPROTO_TCP} };

typedef void* (*uv_malloc_func)(size_t size);
typedef void* (*uv_realloc_func)(void* ptr, size_t size);
typedef void* (*uv_calloc_func)(size_t count, size_t size);
typedef void (*uv_free_func)(void* ptr);

typedef struct {
    uv_malloc_func local_malloc;
    uv_realloc_func local_realloc;
    uv_calloc_func local_calloc;
    uv_free_func local_free;
} uv__allocator_t;

static uv__allocator_t uv__allocator = {
    malloc,
    realloc,
    calloc,
    free,
};


void* uv__malloc(size_t size) {
    if (size > 0)
        return uv__allocator.local_malloc(size);
    return NULL;
}


void uv__free(void* ptr) {
    int saved_errno;

    /* Libuv expects that free() does not clobber errno.  The system allocator
     * honors that assumption but custom allocators may not be so careful.
     */
    saved_errno = errno;
    uv__allocator.local_free(ptr);
    errno = saved_errno;
}


static int WSAAPI uv__gethostnamew_nt60(PWSTR name, int name_len) {
    int result_len;
    int error_code = NO_ERROR;

    /* WSALookupService stuff
     * Avoid dynamic memory allocation if possible */
    CHAR local_buf[WSAQ_LOCAL_BUF_LEN];
    DWORD dwlen = WSAQ_LOCAL_BUF_LEN;
    WSAQUERYSETW* pwsaq;
    /* hostname returned from WSALookupService stage */
    WCHAR* result_name = NULL;
    /* WSALookupService handle */
    HANDLE hlookup;
    /* Fallback to heap allocation if stack buffer is too small */
    WSAQUERYSETW* heap_data = NULL;

    /* check input */
    if (name == NULL) {
        error_code = WSAEFAULT;
        goto cleanup;
    }

    /*
     * Stage 1: Check environment variable
     * _CLUSTER_NETWORK_NAME_ len == ComputeName(NETBIOS) len.
     * i.e 15 characters + null.
     * It overrides the actual hostname, so application can
     * work when network name and computer name are different
     */
    result_len = GetEnvironmentVariableW(L"_CLUSTER_NETWORK_NAME_",
        name,
        name_len);
    if (result_len != 0) {
        if (result_len > name_len) {
            error_code = WSAEFAULT;
        }
        goto cleanup;
    }

    /* Stage 2: Do normal lookup through WSALookupServiceLookup */
    pwsaq = (WSAQUERYSETW*)local_buf;
    memset(pwsaq, 0, sizeof(*pwsaq));
    pwsaq->dwSize = sizeof(*pwsaq);
    pwsaq->lpszServiceInstanceName = NULL;
    pwsaq->lpServiceClassId = &guid_host_name;
    pwsaq->dwNameSpace = NS_ALL;
    pwsaq->lpafpProtocols = &af_protocols[0];
    pwsaq->dwNumberOfProtocols = 2;

    error_code = WSALookupServiceBeginW(pwsaq, LUP_RETURN_NAME, &hlookup);
    if (error_code == NO_ERROR) {
        /* Try stack allocation first */
        error_code = WSALookupServiceNextW(hlookup, 0, &dwlen, pwsaq);
        if (error_code == NO_ERROR) {
            result_name = pwsaq->lpszServiceInstanceName;
        }
        else {
            error_code = WSAGetLastError();

            if (error_code == WSAEFAULT) {
                /* Should never happen */
                assert(sizeof(CHAR) * dwlen >= sizeof(WSAQUERYSETW));

                /* Fallback to the heap allocation */
                heap_data = uv__malloc(sizeof(CHAR) * (size_t)dwlen);
                if (heap_data != NULL) {
                    error_code = WSALookupServiceNextW(hlookup, 0, &dwlen, heap_data);
                    if (error_code == NO_ERROR) {
                        result_name = heap_data->lpszServiceInstanceName;
                    }
                    else {
                        error_code = WSAGetLastError();
                    }
                }
                else {
                    error_code = WSA_NOT_ENOUGH_MEMORY;
                }
            }
        }

        WSALookupServiceEnd(hlookup);

        if (error_code != NO_ERROR) {
            WSASetLastError(error_code);
        }
    }

    if (result_name != NULL) {
        size_t wlen = wcslen(result_name) + 1;

        if (wlen <= (size_t)name_len) {
            wmemcpy(name, result_name, wlen);
        }
        else {
            error_code = WSAEFAULT;
        }
        goto cleanup;
    }

    /* Stage 3: If WSALookupServiceLookup fails, fallback to GetComputerName */
    result_len = name_len;
    /* Reset error code */
    error_code = NO_ERROR;

    if (GetComputerNameW(name, (PDWORD)&result_len) == FALSE) {
        error_code = WSAENETDOWN;
        if (result_len >= name_len) {
            error_code = WSAEFAULT;
        }
    }

cleanup:
    uv__free(heap_data);

    if (error_code == NO_ERROR) {
        return NO_ERROR;
    }
    else {
        WSASetLastError(error_code);
        return SOCKET_ERROR;
    }
}


int main()
{
    WCHAR buf[256];
    int result;

    result = uv__gethostnamew_nt60(buf, sizeof(buf));
    wprintf(L"result: %d, hostname: '%ls'\n", result, buf);

    return 0;
}