// Kernel socket layer. This file owns process socket objects and provider-neutral socket dispatch.

#include <stdint.h>
#include <stddef.h>
#include <stdbool.h>
#include <string.h>

#include <kernel/net_socket.h>
#include <kernel/ke_event.h>
#include <kernel/ke_irp.h>
#include <kernel/ke_process.h>
#include <kernel/ke_sched.h>
#include <kernel/ke_semaphore.h>
#include <kernel/ke_spinlock.h>
#include <kernel/ke_thread.h>
#include <kernel/mm_virt.h>
#include <kernel/ob_object.h>
#include <ntlite/ddi/io_driver.h>
#include <ntlite/ddi/net_stack.h>
#include <ntlite/syscall.h>

#define NET_SOCKET_IO_MAX_BYTES  (1024u * 1024u)
#define NET_SOCKET_UDP_MAX_BYTES 65507u
#define NET_SOCKET_RAW_MAX_BYTES 65535u

static SPINLOCK gNetIrpStatsLock = SPINLOCK_INIT;
static NET_IRP_STATS gNetIrpStats;

// Records that a socket operation entered the IRP boundary.
static void NetIrpStatsSubmit(uint32_t operation) {
    uintptr_t flags = KeAcquireSpinlock(&gNetIrpStatsLock);
    gNetIrpStats.Submitted++;
    if (operation < NetSocketIrpOperationCount)
        gNetIrpStats.ByOperation[operation]++;
    KeReleaseSpinlock(&gNetIrpStatsLock, flags);
}

// Records that a submitted socket IRP returned pending to the socket facade.
static void NetIrpStatsPending(void) {
    uintptr_t flags = KeAcquireSpinlock(&gNetIrpStatsLock);
    gNetIrpStats.Pending++;
    KeReleaseSpinlock(&gNetIrpStatsLock, flags);
}

// Records the final status of a socket IRP after synchronous wait completion.
static void NetIrpStatsComplete(KSTATUS status, bool asynchronous) {
    uintptr_t flags = KeAcquireSpinlock(&gNetIrpStatsLock);
    gNetIrpStats.Completed++;
    if (asynchronous)
        gNetIrpStats.CompletedAsynchronously++;
    else
        gNetIrpStats.CompletedSynchronously++;

    if (status == NTLSTATUS_CANCELLED)
        gNetIrpStats.Cancelled++;
    else if (status == NTLSTATUS_WOULD_BLOCK)
        gNetIrpStats.WouldBlock++;
    else if (status == NTLSTATUS_TIMEOUT)
        gNetIrpStats.Timeout++;
    else if (KFAILED(status))
        gNetIrpStats.Failed++;

    KeReleaseSpinlock(&gNetIrpStatsLock, flags);
}

// Returns a stable snapshot of socket IRP diagnostics.
__declspec(dllexport) void NetQueryIrpStats(NET_IRP_STATS* out) {
    if (!out)
        return;

    uintptr_t flags = KeAcquireSpinlock(&gNetIrpStatsLock);
    *out = gNetIrpStats;
    KeReleaseSpinlock(&gNetIrpStatsLock, flags);
}

// Records that a public async socket request requested cancellation.
static void NetIrpStatsCancelRequested(void) {
    uintptr_t flags = KeAcquireSpinlock(&gNetIrpStatsLock);
    gNetIrpStats.CancelRequested++;
    KeReleaseSpinlock(&gNetIrpStatsLock, flags);
}

struct NET_ASYNC_SOCKET_REQUEST_OBJECT {
    IO_ASYNC_REQUEST_OBJECT Async;

    PROCESS* OwnerProcess;
    SOCKET_OBJECT* Socket;
    uint32_t Operation;
    uint32_t SocketFlags;
    uint32_t Backlog;
    uint32_t ShutdownHow;
    uint32_t Option;
    uint32_t PollEvents;

    void* KernelBuffer;
    void* UserBuffer;
    size_t Length;

    SYS_SOCKET_ADDRESS_K Address;
    size_t AddressSize;
    SYS_SOCKET_ADDRESS_K* AddressOut;
    size_t* AddressSizeInOut;

    void* OptionValue;
    void* UserOptionValue;
    size_t OptionValueSize;
    size_t* OptionActualSizeOut;

    SYS_HANDLE* HandleOut;
    uint32_t* PollReturnedEventsOut;

    MM_LOCKED_USER_BUFFER* LockedAddressOut;
    MM_LOCKED_USER_BUFFER* LockedAddressSize;
    MM_LOCKED_USER_BUFFER* LockedOptionValue;
    MM_LOCKED_USER_BUFFER* LockedOptionActualSize;
    MM_LOCKED_USER_BUFFER* LockedHandleOut;
    MM_LOCKED_USER_BUFFER* LockedPollReturnedEvents;
};


static SPINLOCK gAsyncSocketInitLock = SPINLOCK_INIT;
static int gAsyncSocketInitState;
static SPINLOCK gAsyncSocketQueueLock = SPINLOCK_INIT;
static LIST_ENTRY gAsyncSocketQueue;
static SEMAPHORE* gAsyncSocketSemaphore;

// Releases operation-specific user locks that the common async object does not own.
static void NetAsyncSocketReleaseExtraLocks(NET_ASYNC_SOCKET_REQUEST_OBJECT* req) {
    if (!req)
        return;

    if (req->LockedAddressOut) {
        MmUnlockUserBuffer(req->LockedAddressOut);
        req->LockedAddressOut = (MM_LOCKED_USER_BUFFER*)0;
    }
    if (req->LockedAddressSize) {
        MmUnlockUserBuffer(req->LockedAddressSize);
        req->LockedAddressSize = (MM_LOCKED_USER_BUFFER*)0;
    }
    if (req->LockedOptionValue) {
        MmUnlockUserBuffer(req->LockedOptionValue);
        req->LockedOptionValue = (MM_LOCKED_USER_BUFFER*)0;
    }
    if (req->LockedOptionActualSize) {
        MmUnlockUserBuffer(req->LockedOptionActualSize);
        req->LockedOptionActualSize = (MM_LOCKED_USER_BUFFER*)0;
    }
    if (req->LockedHandleOut) {
        MmUnlockUserBuffer(req->LockedHandleOut);
        req->LockedHandleOut = (MM_LOCKED_USER_BUFFER*)0;
    }
    if (req->LockedPollReturnedEvents) {
        MmUnlockUserBuffer(req->LockedPollReturnedEvents);
        req->LockedPollReturnedEvents = (MM_LOCKED_USER_BUFFER*)0;
    }
}

// Releases heap storage no longer needed after a socket operation has finished.
static void NetAsyncSocketReleaseKernelStorage(NET_ASYNC_SOCKET_REQUEST_OBJECT* req) {
    if (!req)
        return;

    if (req->KernelBuffer) {
        MmFreePool(req->KernelBuffer);
        req->KernelBuffer = (void*)0;
    }
    if (req->OptionValue) {
        MmFreePool(req->OptionValue);
        req->OptionValue = (void*)0;
    }
}

// Releases socket-specific references owned by an async socket request.
static void NetAsyncSocketRequestCleanup(IO_ASYNC_REQUEST_OBJECT* async, void* context) {
    (void)context;
    NET_ASYNC_SOCKET_REQUEST_OBJECT* req =
        (NET_ASYNC_SOCKET_REQUEST_OBJECT*)async;

    NetAsyncSocketReleaseExtraLocks(req);
    NetAsyncSocketReleaseKernelStorage(req);

    if (req->Socket) {
        ObDereferenceObject(&req->Socket->Header);
        req->Socket = (SOCKET_OBJECT*)0;
    }
    if (req->OwnerProcess) {
        ObDereferenceObject(&req->OwnerProcess->Header);
        req->OwnerProcess = (PROCESS*)0;
    }
}

// Cancels an async socket request from the shared async I/O handle.
static KSTATUS NetCancelAsyncSocketRequest(IO_ASYNC_REQUEST_OBJECT* async, void* context) {
    (void)context;
    if (!async)
        return NTLSTATUS_INVALID_PARAMETER;

    NetIrpStatsCancelRequested();

    uintptr_t flags = KeAcquireSpinlock(&gAsyncSocketQueueLock);
    IO_ASYNC_REQUEST_STATE state = IoAsyncRequestGetState(async);

    if (state == IoAsyncRequestQueued) {
        ListRemove(&async->Link);
        IoAsyncRequestRequestCancel(async);
        IoAsyncRequestSetState(async, IoAsyncRequestCancelled);
        KeReleaseSpinlock(&gAsyncSocketQueueLock, flags);

        IoCompleteAsyncRequest(async, NTLSTATUS_CANCELLED, 0);
        ObDereferenceObject(&async->Header);
        return NTLSTATUS_SUCCESS;
    }

    if (state == IoAsyncRequestRunning) {
        IoAsyncRequestRequestCancel(async);
        KeReleaseSpinlock(&gAsyncSocketQueueLock, flags);
        return NTLSTATUS_PENDING;
    }

    KeReleaseSpinlock(&gAsyncSocketQueueLock, flags);
    return NTLSTATUS_OBJECT_DEAD;
}

// Removes one queued async socket request for the worker.
static NET_ASYNC_SOCKET_REQUEST_OBJECT* NetAsyncSocketDequeue(void) {
    uintptr_t flags = KeAcquireSpinlock(&gAsyncSocketQueueLock);
    LIST_ENTRY* node = ListPopHead(&gAsyncSocketQueue);
    if (!node) {
        KeReleaseSpinlock(&gAsyncSocketQueueLock, flags);
        return (NET_ASYNC_SOCKET_REQUEST_OBJECT*)0;
    }

    NET_ASYNC_SOCKET_REQUEST_OBJECT* req =
        (NET_ASYNC_SOCKET_REQUEST_OBJECT*)CONTAINING_RECORD(node, IO_ASYNC_REQUEST_OBJECT, Link);
    if (IoAsyncRequestIsCancelRequested(&req->Async))
        IoAsyncRequestSetState(&req->Async, IoAsyncRequestCancelled);
    else
        IoAsyncRequestSetState(&req->Async, IoAsyncRequestRunning);
    KeReleaseSpinlock(&gAsyncSocketQueueLock, flags);
    return req;
}

// Publishes a remote address output for accept and recvfrom.
static size_t NetAsyncSocketPublishAddress(NET_ASYNC_SOCKET_REQUEST_OBJECT* req,
                                           const SYS_SOCKET_ADDRESS_K* address) {
    if (!req || !req->AddressOut || !address)
        return 0;

    memcpy(req->AddressOut, address, sizeof(*address));
    if (req->AddressSizeInOut)
        *req->AddressSizeInOut = sizeof(*address);
    return sizeof(*address);
}

// Executes one queued async socket operation on the worker thread.
static void NetAsyncSocketExecute(NET_ASYNC_SOCKET_REQUEST_OBJECT* req,
                                  KSTATUS* statusOut,
                                  size_t* informationOut) {
    KSTATUS status = NTLSTATUS_INVALID_PARAMETER;
    size_t information = 0;

    if (!req || !statusOut || !informationOut)
        return;

    switch (req->Operation) {
    case SYS_ASYNC_SOCKET_OP_BIND:
        status = NetSocketBind(req->Socket, &req->Address, req->AddressSize);
        break;

    case SYS_ASYNC_SOCKET_OP_LISTEN:
        status = NetSocketListen(req->Socket, req->Backlog);
        break;

    case SYS_ASYNC_SOCKET_OP_ACCEPT: {
        SOCKET_OBJECT* accepted = (SOCKET_OBJECT*)0;
        SYS_SOCKET_ADDRESS_K remote;
        memset(&remote, 0, sizeof(remote));
        status = NetSocketAccept(req->Socket, &accepted, req->AddressOut ? &remote : (SYS_SOCKET_ADDRESS_K*)0, &req->Async);
        if (KSUCCESS(status)) {
            SYS_HANDLE handle = SYS_HANDLE_INVALID;
            status = KeCreateHandle(req->OwnerProcess,
                                    &accepted->Header,
                                    KeGetValidHandleRights(ObjectTypeSocket),
                                    &handle);
            ObDereferenceObject(&accepted->Header);
            if (KSUCCESS(status)) {
                if (req->HandleOut)
                    *req->HandleOut = handle;
                information = handle;
                (void)NetAsyncSocketPublishAddress(req, &remote);
            }
        }
        break;
    }

    case SYS_ASYNC_SOCKET_OP_CONNECT:
        status = NetSocketConnect(req->Socket, &req->Address, req->AddressSize, &req->Async);
        break;

    case SYS_ASYNC_SOCKET_OP_SEND: {
        size_t sent = 0;
        status = NetSocketSend(req->Socket, req->KernelBuffer, req->Length, req->SocketFlags, &sent);
        information = sent;
        break;
    }

    case SYS_ASYNC_SOCKET_OP_RECV: {
        size_t received = 0;
        status = NetSocketReceive(req->Socket, req->KernelBuffer, req->Length, req->SocketFlags, &received, &req->Async);
        if (KSUCCESS(status) && received)
            memcpy(req->UserBuffer, req->KernelBuffer, received);
        information = received;
        break;
    }

    case SYS_ASYNC_SOCKET_OP_SEND_TO: {
        size_t sent = 0;
        status = NetSocketSendTo(req->Socket, req->KernelBuffer, req->Length,
                                 req->SocketFlags, &req->Address, req->AddressSize, &sent);
        information = sent;
        break;
    }

    case SYS_ASYNC_SOCKET_OP_RECV_FROM: {
        size_t received = 0;
        SYS_SOCKET_ADDRESS_K remote;
        memset(&remote, 0, sizeof(remote));
        status = NetSocketReceiveFrom(req->Socket, req->KernelBuffer, req->Length,
                                      req->SocketFlags,
                                      req->AddressOut ? &remote : (SYS_SOCKET_ADDRESS_K*)0,
                                      &received,
                                      &req->Async);
        if (KSUCCESS(status) && received)
            memcpy(req->UserBuffer, req->KernelBuffer, received);
        if (KSUCCESS(status))
            (void)NetAsyncSocketPublishAddress(req, &remote);
        information = received;
        break;
    }

    case SYS_ASYNC_SOCKET_OP_SHUTDOWN:
        status = NetSocketShutdown(req->Socket, req->ShutdownHow);
        break;

    case SYS_ASYNC_SOCKET_OP_SET_OPT:
        status = NetSocketSetOption(req->Socket, req->Option, req->OptionValue, req->OptionValueSize);
        break;

    case SYS_ASYNC_SOCKET_OP_GET_OPT: {
        size_t actual = 0;
        status = NetSocketGetOption(req->Socket, req->Option, req->OptionValue, req->OptionValueSize, &actual);
        if (KSUCCESS(status) && actual && req->UserOptionValue)
            memcpy(req->UserOptionValue, req->OptionValue, actual);
        if ((KSUCCESS(status) || status == NTLSTATUS_BUFFER_TOO_SMALL) && req->OptionActualSizeOut)
            *req->OptionActualSizeOut = actual;
        information = actual;
        break;
    }

    case SYS_ASYNC_SOCKET_OP_POLL: {
        uint32_t returnedEvents = 0;
        status = NetSocketPoll(req->Socket, req->PollEvents, &returnedEvents);
        if (req->PollReturnedEventsOut)
            *req->PollReturnedEventsOut = returnedEvents;
        information = returnedEvents;
        break;
    }

    default:
        status = NTLSTATUS_INVALID_PARAMETER;
        break;
    }

    *statusOut = status;
    *informationOut = information;
}

// Worker for public async socket requests.
static void NetAsyncSocketWorker(void* arg) {
    (void)arg;

    while (true) {
        KSTATUS waitStatus = KeWaitSemaphore(gAsyncSocketSemaphore, KE_WAIT_INFINITE);
        if (!KSUCCESS(waitStatus)) {
            KeYield();
            continue;
        }

        NET_ASYNC_SOCKET_REQUEST_OBJECT* req = NetAsyncSocketDequeue();
        if (!req)
            continue;

        KSTATUS status = NTLSTATUS_CANCELLED;
        size_t information = 0;
        if (!IoAsyncRequestIsCancelRequested(&req->Async) &&
            IoAsyncRequestGetState(&req->Async) != IoAsyncRequestCancelled) {
            NetAsyncSocketExecute(req, &status, &information);
        }

        NetAsyncSocketReleaseExtraLocks(req);
        NetAsyncSocketReleaseKernelStorage(req);
        IoCompleteAsyncRequest(&req->Async, status, information);
        ObDereferenceObject(&req->Async.Header);
    }
}

// Starts the async socket worker on first use.
static KSTATUS NetEnsureAsyncSocketWorkerStarted(void) {
    uintptr_t flags = KeAcquireSpinlock(&gAsyncSocketInitLock);
    if (gAsyncSocketInitState == 2) {
        KeReleaseSpinlock(&gAsyncSocketInitLock, flags);
        return NTLSTATUS_SUCCESS;
    }
    if (gAsyncSocketInitState == 1) {
        KeReleaseSpinlock(&gAsyncSocketInitLock, flags);
        while (gAsyncSocketInitState == 1)
            KeYield();
        return (gAsyncSocketInitState == 2) ? NTLSTATUS_SUCCESS : NTLSTATUS_FAILURE;
    }

    gAsyncSocketInitState = 1;
    KeReleaseSpinlock(&gAsyncSocketInitLock, flags);

    LIST_INIT(&gAsyncSocketQueue);

    KSTATUS status = KeSemaphoreCreate(0, 0x7fffffff, &gAsyncSocketSemaphore);
    if (KSUCCESS(status)) {
        THREAD* worker = (THREAD*)0;
        const KE_THREAD_CREATE_PARAMETERS parameters =
            KE_THREAD_CREATE_NAMED_NORMAL(PriorityNormal, "net.socket.async");
        status = KeCreateKernelThread((uintptr_t)NetAsyncSocketWorker,
                                      0,
                                      &parameters,
                                      &worker);
        (void)worker;
    }

    flags = KeAcquireSpinlock(&gAsyncSocketInitLock);
    gAsyncSocketInitState = KSUCCESS(status) ? 2 : 0;
    KeReleaseSpinlock(&gAsyncSocketInitLock, flags);
    return status;
}

// Submits a public async socket request to the socket worker.
KSTATUS NetSocketSubmitAsyncKernel(const NET_ASYNC_SOCKET_SUBMIT* submit,
                                                         NET_ASYNC_SOCKET_REQUEST_OBJECT** requestOut) {
    if (requestOut)
        *requestOut = (NET_ASYNC_SOCKET_REQUEST_OBJECT*)0;
    if (!submit || !submit->OwnerProcess || !submit->Socket || !submit->Iosb)
        return NTLSTATUS_INVALID_PARAMETER;

    KSTATUS status = NetEnsureAsyncSocketWorkerStarted();
    if (!KSUCCESS(status))
        return status;

    IO_ASYNC_REQUEST_OBJECT* async = (IO_ASYNC_REQUEST_OBJECT*)0;
    status = IoCreateAsyncRequest(sizeof(NET_ASYNC_SOCKET_REQUEST_OBJECT),
                                  NetAsyncSocketRequestCleanup,
                                  (void*)0,
                                  NetCancelAsyncSocketRequest,
                                  (void*)0,
                                  &async);
    if (!KSUCCESS(status))
        return status;

    NET_ASYNC_SOCKET_REQUEST_OBJECT* req = (NET_ASYNC_SOCKET_REQUEST_OBJECT*)async;
    req->OwnerProcess = submit->OwnerProcess;
    req->Socket = submit->Socket;
    req->Operation = submit->Operation;
    req->SocketFlags = submit->SocketFlags;
    req->Backlog = submit->Backlog;
    req->ShutdownHow = submit->ShutdownHow;
    req->Option = submit->Option;
    req->PollEvents = submit->PollEvents;
    req->KernelBuffer = submit->KernelBuffer;
    req->UserBuffer = submit->UserBuffer;
    req->Length = submit->Length;
    req->Address = submit->Address;
    req->AddressSize = submit->AddressSize;
    req->AddressOut = submit->AddressOut;
    req->AddressSizeInOut = submit->AddressSizeInOut;
    req->OptionValue = submit->OptionValue;
    req->UserOptionValue = submit->UserOptionValue;
    req->OptionValueSize = submit->OptionValueSize;
    req->OptionActualSizeOut = submit->OptionActualSizeOut;
    req->HandleOut = submit->HandleOut;
    req->PollReturnedEventsOut = submit->PollReturnedEventsOut;
    req->LockedAddressOut = submit->LockedAddressOut;
    req->LockedAddressSize = submit->LockedAddressSize;
    req->LockedOptionValue = submit->LockedOptionValue;
    req->LockedOptionActualSize = submit->LockedOptionActualSize;
    req->LockedHandleOut = submit->LockedHandleOut;
    req->LockedPollReturnedEvents = submit->LockedPollReturnedEvents;

    ObReferenceObject(&req->OwnerProcess->Header);
    ObReferenceObject(&req->Socket->Header);
    IoInitializeAsyncRequestCompletion(&req->Async,
                                       submit->Iosb,
                                       submit->CompletionEvent,
                                       submit->CompletionApc,
                                       submit->LockedBuffer,
                                       submit->LockedIosb);

    ObReferenceObject(&req->Async.Header);
    uintptr_t qflags = KeAcquireSpinlock(&gAsyncSocketQueueLock);
    ListInsertTail(&gAsyncSocketQueue, &req->Async.Link);
    KeReleaseSpinlock(&gAsyncSocketQueueLock, qflags);
    KeSignalSemaphore(gAsyncSocketSemaphore);

    if (requestOut)
        *requestOut = req;
    else
        ObDereferenceObject(&req->Async.Header);
    return NTLSTATUS_PENDING;
}

// Returns the object header embedded in a socket async request.
OBJECT* NetAsyncSocketRequestObject(NET_ASYNC_SOCKET_REQUEST_OBJECT* request) {
    return request ? &request->Async.Header : (OBJECT*)0;
}

// Submits a socket provider request through the active stack DEVICE_OBJECT.
static KSTATUS NetSubmitSocketIrp(const NET_STACK_PROVIDER* provider,
                                  NET_SOCKET_IRP_REQUEST* request,
                                  IO_ASYNC_REQUEST_OBJECT* async) {
    if (!provider || !provider->StackDevice || !request)
        return NTLSTATUS_INVALID_PARAMETER;
    if (request->Size != sizeof(*request))
        return NTLSTATUS_INVALID_PARAMETER;
    if (request->Operation == NetSocketIrpInvalid ||
        request->Operation >= NetSocketIrpOperationCount)
        return NTLSTATUS_INVALID_PARAMETER;

    request->CancelRequested = async ? &async->CancelRequested : (volatile uint8_t*)0;

    EVENT* event = NULL;
    KSTATUS st = KeEventCreate(EventAutoReset, 0, &event);
    if (!KSUCCESS(st))
        return st;

    IRP irp;
    IoInitializeIrp(&irp, 1);
    IoBuildDeviceIoControlIrp(&irp,
                              IRP_MJ_INTERNAL_DEVICE_CONTROL,
                              NET_IOCTL_SOCKET,
                              request,
                              sizeof(*request),
                              NULL,
                              0);
    irp.CompletionEvent = event;

    NetIrpStatsSubmit(request->Operation);
    st = IoSubmitIrp(provider->StackDevice, &irp);

    bool asynchronous = false;
    if (st == NTLSTATUS_PENDING) {
        asynchronous = true;
        NetIrpStatsPending();
        (void)KeWaitForEvent(event, 0);
        st = irp.Status;
    }

    NetIrpStatsComplete(st, asynchronous);
    ObDereferenceObject(&event->Header);
    return st;
}


// Closes a provider socket through the IRP path while object teardown is still safe.
static void NetSocketCloseProvider(const NET_STACK_PROVIDER* provider, void* providerSocket) {
    if (!provider || !providerSocket)
        return;

    NET_SOCKET_IRP_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Operation = NetSocketIrpClose;
    request.u.Close.Socket = providerSocket;
    (void)NetSubmitSocketIrp(provider, &request, (IO_ASYNC_REQUEST_OBJECT*)0);
}

// Releases provider state owned by a process socket object.
static void NetSocketDestroy(OBJECT* object) {
    SOCKET_OBJECT* socket = (SOCKET_OBJECT*)object;
    if (socket->Provider && socket->ProviderSocket) {
        NetSocketCloseProvider(socket->Provider, socket->ProviderSocket);
        socket->ProviderSocket = NULL;
    }
}

// Allocates a process-visible socket object around provider-owned state.
static KSTATUS NetSocketAllocate(const NET_STACK_PROVIDER* provider,
                                 void* providerSocket,
                                 uint32_t domain,
                                 uint32_t type,
                                 uint32_t protocol,
                                 SOCKET_OBJECT** out) {
    if (!provider || !providerSocket || !out)
        return NTLSTATUS_INVALID_PARAMETER;

    SOCKET_OBJECT* socket = (SOCKET_OBJECT*)ObAllocateObject(ObjectTypeSocket, sizeof(SOCKET_OBJECT));
    if (!socket)
        return NTLSTATUS_NO_MEMORY;

    socket->Header.Destroy = NetSocketDestroy;
    socket->Provider = provider;
    socket->ProviderSocket = providerSocket;
    socket->Domain = domain;
    socket->Type = type;
    socket->Protocol = protocol;
    *out = socket;
    return NTLSTATUS_SUCCESS;
}

// Returns the provider and provider-private state behind a live socket object.
static KSTATUS NetSocketGetProviderSocket(SOCKET_OBJECT* socket,
                                          const NET_STACK_PROVIDER** providerOut,
                                          void** providerSocketOut) {
    if (!socket || !providerOut || !providerSocketOut)
        return NTLSTATUS_INVALID_PARAMETER;
    if (!socket->Provider || !socket->ProviderSocket)
        return NTLSTATUS_OBJECT_DEAD;

    *providerOut = socket->Provider;
    *providerSocketOut = socket->ProviderSocket;
    return NTLSTATUS_SUCCESS;
}

// Creates a socket by submitting a create request to the active stack device.
KSTATUS NetSocketCreate(uint32_t domain, uint32_t type, uint32_t protocol, SOCKET_OBJECT** out) {
    if (!out)
        return NTLSTATUS_INVALID_PARAMETER;
    *out = NULL;

    const NET_STACK_PROVIDER* provider = NetGetStackProvider();
    if (!provider)
        return NTLSTATUS_NOT_SUPPORTED;

    void* providerSocket = NULL;
    NET_SOCKET_IRP_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Operation = NetSocketIrpCreate;
    request.u.Create.Domain = domain;
    request.u.Create.Type = type;
    request.u.Create.Protocol = protocol;
    request.u.Create.ProviderSocketOut = &providerSocket;

    KSTATUS st = NetSubmitSocketIrp(provider, &request, (IO_ASYNC_REQUEST_OBJECT*)0);
    if (!KSUCCESS(st))
        return st;
    if (!providerSocket)
        return NTLSTATUS_FAILURE;

    st = NetSocketAllocate(provider, providerSocket, domain, type, protocol, out);
    if (!KSUCCESS(st))
        NetSocketCloseProvider(provider, providerSocket);
    return st;
}

// Binds a socket address through the stack device IRP path.
KSTATUS NetSocketBind(SOCKET_OBJECT* socket, const SYS_SOCKET_ADDRESS_K* address, size_t addressSize) {
    const NET_STACK_PROVIDER* provider = NULL;
    void* providerSocket = NULL;
    KSTATUS st = NetSocketGetProviderSocket(socket, &provider, &providerSocket);
    if (!KSUCCESS(st))
        return st;

    NET_SOCKET_IRP_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Operation = NetSocketIrpBind;
    request.u.Bind.Socket = providerSocket;
    request.u.Bind.Address = address;
    request.u.Bind.AddressSize = addressSize;
    return NetSubmitSocketIrp(provider, &request, (IO_ASYNC_REQUEST_OBJECT*)0);
}

// Marks a stream socket as listening through the stack device IRP path.
KSTATUS NetSocketListen(SOCKET_OBJECT* socket, uint32_t backlog) {
    const NET_STACK_PROVIDER* provider = NULL;
    void* providerSocket = NULL;
    KSTATUS st = NetSocketGetProviderSocket(socket, &provider, &providerSocket);
    if (!KSUCCESS(st))
        return st;

    NET_SOCKET_IRP_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Operation = NetSocketIrpListen;
    request.u.Listen.Socket = providerSocket;
    request.u.Listen.Backlog = backlog;
    return NetSubmitSocketIrp(provider, &request, (IO_ASYNC_REQUEST_OBJECT*)0);
}

// Accepts a provider socket through IRP and wraps the result in a process socket object.
KSTATUS NetSocketAccept(SOCKET_OBJECT* socket, SOCKET_OBJECT** acceptedOut,
                        SYS_SOCKET_ADDRESS_K* remoteOut,
                        IO_ASYNC_REQUEST_OBJECT* async) {
    if (!acceptedOut)
        return NTLSTATUS_INVALID_PARAMETER;
    *acceptedOut = NULL;

    const NET_STACK_PROVIDER* provider = NULL;
    void* providerSocket = NULL;
    KSTATUS st = NetSocketGetProviderSocket(socket, &provider, &providerSocket);
    if (!KSUCCESS(st))
        return st;

    void* acceptedProviderSocket = NULL;
    NET_SOCKET_IRP_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Operation = NetSocketIrpAccept;
    request.u.Accept.Socket = providerSocket;
    request.u.Accept.AcceptedOut = &acceptedProviderSocket;
    request.u.Accept.RemoteOut = remoteOut;

    st = NetSubmitSocketIrp(provider, &request, async);
    if (!KSUCCESS(st))
        return st;
    if (!acceptedProviderSocket)
        return NTLSTATUS_FAILURE;

    st = NetSocketAllocate(provider,
                           acceptedProviderSocket,
                           socket->Domain,
                           socket->Type,
                           socket->Protocol,
                           acceptedOut);
    if (!KSUCCESS(st))
        NetSocketCloseProvider(provider, acceptedProviderSocket);
    return st;
}


// Connects a socket through the stack device IRP path.
KSTATUS NetSocketConnect(SOCKET_OBJECT* socket, const SYS_SOCKET_ADDRESS_K* address,
                         size_t addressSize, IO_ASYNC_REQUEST_OBJECT* async) {
    const NET_STACK_PROVIDER* provider = NULL;
    void* providerSocket = NULL;
    KSTATUS st = NetSocketGetProviderSocket(socket, &provider, &providerSocket);
    if (!KSUCCESS(st))
        return st;

    NET_SOCKET_IRP_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Operation = NetSocketIrpConnect;
    request.u.Connect.Socket = providerSocket;
    request.u.Connect.Address = address;
    request.u.Connect.AddressSize = addressSize;
    return NetSubmitSocketIrp(provider, &request, async);
}


// Sends stream bytes through the stack device IRP path.
KSTATUS NetSocketSend(SOCKET_OBJECT* socket, const void* buffer, size_t length, uint32_t flags, size_t* sentOut) {
    if (sentOut)
        *sentOut = 0;
    if (!buffer && length)
        return NTLSTATUS_INVALID_PARAMETER;
    if (length > NET_SOCKET_IO_MAX_BYTES)
        return NTLSTATUS_INVALID_PARAMETER;

    const NET_STACK_PROVIDER* provider = NULL;
    void* providerSocket = NULL;
    KSTATUS st = NetSocketGetProviderSocket(socket, &provider, &providerSocket);
    if (!KSUCCESS(st))
        return st;

    NET_SOCKET_IRP_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Operation = NetSocketIrpSend;
    request.u.Send.Socket = providerSocket;
    request.u.Send.Buffer = buffer;
    request.u.Send.Length = length;
    request.u.Send.Flags = flags;
    request.u.Send.SentOut = sentOut;
    return NetSubmitSocketIrp(provider, &request, (IO_ASYNC_REQUEST_OBJECT*)0);
}

// Receives stream bytes through the stack device IRP path.
KSTATUS NetSocketReceive(SOCKET_OBJECT* socket, void* buffer, size_t length,
                         uint32_t flags, size_t* receivedOut,
                         IO_ASYNC_REQUEST_OBJECT* async) {
    if (receivedOut)
        *receivedOut = 0;
    if (!buffer && length)
        return NTLSTATUS_INVALID_PARAMETER;
    if (length > NET_SOCKET_IO_MAX_BYTES)
        return NTLSTATUS_INVALID_PARAMETER;

    const NET_STACK_PROVIDER* provider = NULL;
    void* providerSocket = NULL;
    KSTATUS st = NetSocketGetProviderSocket(socket, &provider, &providerSocket);
    if (!KSUCCESS(st))
        return st;

    NET_SOCKET_IRP_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Operation = NetSocketIrpReceive;
    request.u.Receive.Socket = providerSocket;
    request.u.Receive.Buffer = buffer;
    request.u.Receive.Length = length;
    request.u.Receive.Flags = flags;
    request.u.Receive.ReceivedOut = receivedOut;
    return NetSubmitSocketIrp(provider, &request, async);
}


// Sends datagram bytes through the stack device IRP path.
KSTATUS NetSocketSendTo(SOCKET_OBJECT* socket, const void* buffer, size_t length, uint32_t flags,
                        const SYS_SOCKET_ADDRESS_K* remote, size_t remoteSize, size_t* sentOut) {
    if (sentOut)
        *sentOut = 0;
    if (!buffer && length)
        return NTLSTATUS_INVALID_PARAMETER;
    if (socket && socket->Type == SYS_SOCK_DGRAM && length > NET_SOCKET_UDP_MAX_BYTES)
        return NTLSTATUS_INVALID_PARAMETER;
    if (socket && socket->Type == SYS_SOCK_RAW && length > NET_SOCKET_RAW_MAX_BYTES)
        return NTLSTATUS_INVALID_PARAMETER;
    if (socket && socket->Type != SYS_SOCK_DGRAM && socket->Type != SYS_SOCK_RAW)
        return NTLSTATUS_INVALID_PARAMETER;

    const NET_STACK_PROVIDER* provider = NULL;
    void* providerSocket = NULL;
    KSTATUS st = NetSocketGetProviderSocket(socket, &provider, &providerSocket);
    if (!KSUCCESS(st))
        return st;

    NET_SOCKET_IRP_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Operation = NetSocketIrpSendTo;
    request.u.SendTo.Socket = providerSocket;
    request.u.SendTo.Buffer = buffer;
    request.u.SendTo.Length = length;
    request.u.SendTo.Flags = flags;
    request.u.SendTo.Remote = remote;
    request.u.SendTo.RemoteSize = remoteSize;
    request.u.SendTo.SentOut = sentOut;
    return NetSubmitSocketIrp(provider, &request, (IO_ASYNC_REQUEST_OBJECT*)0);
}

// Receives datagram bytes through the stack device IRP path.
KSTATUS NetSocketReceiveFrom(SOCKET_OBJECT* socket, void* buffer, size_t length,
                             uint32_t flags, SYS_SOCKET_ADDRESS_K* remoteOut,
                             size_t* receivedOut, IO_ASYNC_REQUEST_OBJECT* async) {
    if (receivedOut)
        *receivedOut = 0;
    if (!buffer && length)
        return NTLSTATUS_INVALID_PARAMETER;
    if (socket && socket->Type == SYS_SOCK_DGRAM && length > NET_SOCKET_UDP_MAX_BYTES)
        return NTLSTATUS_INVALID_PARAMETER;
    if (socket && socket->Type == SYS_SOCK_RAW && length > NET_SOCKET_RAW_MAX_BYTES)
        return NTLSTATUS_INVALID_PARAMETER;
    if (socket && socket->Type != SYS_SOCK_DGRAM && socket->Type != SYS_SOCK_RAW)
        return NTLSTATUS_INVALID_PARAMETER;

    const NET_STACK_PROVIDER* provider = NULL;
    void* providerSocket = NULL;
    KSTATUS st = NetSocketGetProviderSocket(socket, &provider, &providerSocket);
    if (!KSUCCESS(st))
        return st;

    NET_SOCKET_IRP_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Operation = NetSocketIrpReceiveFrom;
    request.u.ReceiveFrom.Socket = providerSocket;
    request.u.ReceiveFrom.Buffer = buffer;
    request.u.ReceiveFrom.Length = length;
    request.u.ReceiveFrom.Flags = flags;
    request.u.ReceiveFrom.RemoteOut = remoteOut;
    request.u.ReceiveFrom.ReceivedOut = receivedOut;
    return NetSubmitSocketIrp(provider, &request, async);
}


// Shuts down one or both socket directions through the stack device IRP path.
KSTATUS NetSocketShutdown(SOCKET_OBJECT* socket, uint32_t how) {
    const NET_STACK_PROVIDER* provider = NULL;
    void* providerSocket = NULL;
    KSTATUS st = NetSocketGetProviderSocket(socket, &provider, &providerSocket);
    if (!KSUCCESS(st))
        return st;

    NET_SOCKET_IRP_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Operation = NetSocketIrpShutdown;
    request.u.Shutdown.Socket = providerSocket;
    request.u.Shutdown.How = how;
    return NetSubmitSocketIrp(provider, &request, (IO_ASYNC_REQUEST_OBJECT*)0);
}

// Sets a socket option through the stack device IRP path.
KSTATUS NetSocketSetOption(SOCKET_OBJECT* socket, uint32_t option, const void* value, size_t valueSize) {
    const NET_STACK_PROVIDER* provider = NULL;
    void* providerSocket = NULL;
    KSTATUS st = NetSocketGetProviderSocket(socket, &provider, &providerSocket);
    if (!KSUCCESS(st))
        return st;

    NET_SOCKET_IRP_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Operation = NetSocketIrpSetOption;
    request.u.SetOption.Socket = providerSocket;
    request.u.SetOption.Option = option;
    request.u.SetOption.Value = value;
    request.u.SetOption.ValueSize = valueSize;
    return NetSubmitSocketIrp(provider, &request, (IO_ASYNC_REQUEST_OBJECT*)0);
}

// Reads a socket option through the stack device IRP path.
KSTATUS NetSocketGetOption(SOCKET_OBJECT* socket, uint32_t option, void* value, size_t valueSize, size_t* actualSizeOut) {
    const NET_STACK_PROVIDER* provider = NULL;
    void* providerSocket = NULL;
    KSTATUS st = NetSocketGetProviderSocket(socket, &provider, &providerSocket);
    if (!KSUCCESS(st))
        return st;

    NET_SOCKET_IRP_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Operation = NetSocketIrpGetOption;
    request.u.GetOption.Socket = providerSocket;
    request.u.GetOption.Option = option;
    request.u.GetOption.Value = value;
    request.u.GetOption.ValueSize = valueSize;
    request.u.GetOption.ActualSizeOut = actualSizeOut;
    return NetSubmitSocketIrp(provider, &request, (IO_ASYNC_REQUEST_OBJECT*)0);
}

// Polls socket readiness through the stack device IRP path.
KSTATUS NetSocketPoll(SOCKET_OBJECT* socket, uint32_t events, uint32_t* returnedEventsOut) {
    if (!returnedEventsOut)
        return NTLSTATUS_INVALID_PARAMETER;
    *returnedEventsOut = 0;

    const NET_STACK_PROVIDER* provider = NULL;
    void* providerSocket = NULL;
    KSTATUS st = NetSocketGetProviderSocket(socket, &provider, &providerSocket);
    if (!KSUCCESS(st))
        return st;

    NET_SOCKET_IRP_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Operation = NetSocketIrpPoll;
    request.u.Poll.Socket = providerSocket;
    request.u.Poll.Events = events;
    request.u.Poll.ReturnedEventsOut = returnedEventsOut;
    return NetSubmitSocketIrp(provider, &request, (IO_ASYNC_REQUEST_OBJECT*)0);
}

// Resolves a host name through the stack device IRP path.
KSTATUS NetResolveHost(const char* name, uint32_t family, SYS_SOCKET_ADDRESS_K* addressOut) {
    if (!name || !addressOut)
        return NTLSTATUS_INVALID_PARAMETER;

    const NET_STACK_PROVIDER* provider = NetGetStackProvider();
    if (!provider)
        return NTLSTATUS_NOT_SUPPORTED;

    NET_SOCKET_IRP_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Operation = NetSocketIrpResolveHost;
    request.u.ResolveHost.Name = name;
    request.u.ResolveHost.Family = family;
    request.u.ResolveHost.AddressOut = addressOut;
    return NetSubmitSocketIrp(provider, &request, (IO_ASYNC_REQUEST_OBJECT*)0);
}