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

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

#include <kernel/net_socket.h>
#include <ntlite/ddi/net_stack.h>
#include <kernel/ob_object.h>

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

// Implements the net Socket Destroy operation.
static void NetSocketDestroy(OBJECT* object) {
    SOCKET_OBJECT* socket = (SOCKET_OBJECT*)object;
    if (socket->Provider && socket->ProviderSocket) {
        socket->Provider->SocketClose(socket->ProviderSocket);
        socket->ProviderSocket = NULL;
    }
}

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 STATUS_INVALID_PARAMETER;

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

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

static KSTATUS NetSocketGetProviderSocket(SOCKET_OBJECT* socket,
                                          const NET_STACK_PROVIDER** providerOut,
                                          void** providerSocketOut) {
    if (!socket || !providerOut || !providerSocketOut)
        return STATUS_INVALID_PARAMETER;
    if (!socket->Provider || !socket->ProviderSocket)
        return STATUS_OBJECT_DEAD;

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

// Implements the net Socket Create operation.
KSTATUS NetSocketCreate(uint32_t domain, uint32_t type, uint32_t protocol, SOCKET_OBJECT** out) {
    if (!out)
        return STATUS_INVALID_PARAMETER;
    *out = NULL;

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

    void* providerSocket = NULL;
    KSTATUS st = provider->SocketCreate(domain, type, protocol, &providerSocket);
    if (!KSUCCESS(st))
        return st;
    if (!providerSocket)
        return STATUS_FAILURE;

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

// Implements the net Socket Bind operation.
KSTATUS NetSocketBind(SOCKET_OBJECT* socket, const SYS_SOCKET_ADDRESS_IPV4_K* address, size_t addressSize) {
    const NET_STACK_PROVIDER* provider = NULL;
    void* providerSocket = NULL;
    KSTATUS st = NetSocketGetProviderSocket(socket, &provider, &providerSocket);
    return KSUCCESS(st) ? provider->Bind(providerSocket, address, addressSize) : st;
}

// Implements the net Socket Listen operation.
KSTATUS NetSocketListen(SOCKET_OBJECT* socket, uint32_t backlog) {
    const NET_STACK_PROVIDER* provider = NULL;
    void* providerSocket = NULL;
    KSTATUS st = NetSocketGetProviderSocket(socket, &provider, &providerSocket);
    return KSUCCESS(st) ? provider->Listen(providerSocket, backlog) : st;
}

// Implements the net Socket Accept operation.
KSTATUS NetSocketAccept(SOCKET_OBJECT* socket, SOCKET_OBJECT** acceptedOut, SYS_SOCKET_ADDRESS_IPV4_K* remoteOut) {
    if (!acceptedOut)
        return STATUS_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;
    st = provider->Accept(providerSocket, &acceptedProviderSocket, remoteOut);
    if (!KSUCCESS(st))
        return st;
    if (!acceptedProviderSocket)
        return STATUS_FAILURE;

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

// Implements the net Socket Connect operation.
KSTATUS NetSocketConnect(SOCKET_OBJECT* socket, const SYS_SOCKET_ADDRESS_IPV4_K* address, size_t addressSize) {
    const NET_STACK_PROVIDER* provider = NULL;
    void* providerSocket = NULL;
    KSTATUS st = NetSocketGetProviderSocket(socket, &provider, &providerSocket);
    return KSUCCESS(st) ? provider->Connect(providerSocket, address, addressSize) : st;
}

// Implements the net Socket Send operation.
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 STATUS_INVALID_PARAMETER;
    if (length > NET_SOCKET_IO_MAX_BYTES)
        return STATUS_INVALID_PARAMETER;

    const NET_STACK_PROVIDER* provider = NULL;
    void* providerSocket = NULL;
    KSTATUS st = NetSocketGetProviderSocket(socket, &provider, &providerSocket);
    return KSUCCESS(st) ? provider->Send(providerSocket, buffer, length, flags, sentOut) : st;
}

// Implements the net Socket Receive operation.
KSTATUS NetSocketReceive(SOCKET_OBJECT* socket, void* buffer, size_t length, uint32_t flags, size_t* receivedOut) {
    if (receivedOut)
        *receivedOut = 0;
    if (!buffer && length)
        return STATUS_INVALID_PARAMETER;
    if (length > NET_SOCKET_IO_MAX_BYTES)
        return STATUS_INVALID_PARAMETER;

    const NET_STACK_PROVIDER* provider = NULL;
    void* providerSocket = NULL;
    KSTATUS st = NetSocketGetProviderSocket(socket, &provider, &providerSocket);
    return KSUCCESS(st) ? provider->Receive(providerSocket, buffer, length, flags, receivedOut) : st;
}

KSTATUS NetSocketSendTo(SOCKET_OBJECT* socket, const void* buffer, size_t length, uint32_t flags,
                        const SYS_SOCKET_ADDRESS_IPV4_K* remote, size_t remoteSize, size_t* sentOut) {
    if (sentOut)
        *sentOut = 0;
    if (!buffer && length)
        return STATUS_INVALID_PARAMETER;
    if (socket && socket->Type == SYS_SOCK_DGRAM && length > NET_SOCKET_UDP_MAX_BYTES)
        return STATUS_INVALID_PARAMETER;
    if (socket && socket->Type == SYS_SOCK_RAW && length > NET_SOCKET_RAW_MAX_BYTES)
        return STATUS_INVALID_PARAMETER;
    if (socket && socket->Type != SYS_SOCK_DGRAM && socket->Type != SYS_SOCK_RAW)
        return STATUS_INVALID_PARAMETER;

    const NET_STACK_PROVIDER* provider = NULL;
    void* providerSocket = NULL;
    KSTATUS st = NetSocketGetProviderSocket(socket, &provider, &providerSocket);
    return KSUCCESS(st) ? provider->SendTo(providerSocket, buffer, length, flags, remote, remoteSize, sentOut) : st;
}

KSTATUS NetSocketReceiveFrom(SOCKET_OBJECT* socket, void* buffer, size_t length, uint32_t flags,
                             SYS_SOCKET_ADDRESS_IPV4_K* remoteOut, size_t* receivedOut) {
    if (receivedOut)
        *receivedOut = 0;
    if (!buffer && length)
        return STATUS_INVALID_PARAMETER;
    if (socket && socket->Type == SYS_SOCK_DGRAM && length > NET_SOCKET_UDP_MAX_BYTES)
        return STATUS_INVALID_PARAMETER;
    if (socket && socket->Type == SYS_SOCK_RAW && length > NET_SOCKET_RAW_MAX_BYTES)
        return STATUS_INVALID_PARAMETER;
    if (socket && socket->Type != SYS_SOCK_DGRAM && socket->Type != SYS_SOCK_RAW)
        return STATUS_INVALID_PARAMETER;

    const NET_STACK_PROVIDER* provider = NULL;
    void* providerSocket = NULL;
    KSTATUS st = NetSocketGetProviderSocket(socket, &provider, &providerSocket);
    return KSUCCESS(st) ? provider->ReceiveFrom(providerSocket, buffer, length, flags, remoteOut, receivedOut) : st;
}

// Implements the net Socket Shutdown operation.
KSTATUS NetSocketShutdown(SOCKET_OBJECT* socket, uint32_t how) {
    const NET_STACK_PROVIDER* provider = NULL;
    void* providerSocket = NULL;
    KSTATUS st = NetSocketGetProviderSocket(socket, &provider, &providerSocket);
    return KSUCCESS(st) ? provider->Shutdown(providerSocket, how) : st;
}

// Implements the net Socket Set Option operation.
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);
    return KSUCCESS(st) ? provider->SetOption(providerSocket, option, value, valueSize) : st;
}

// Implements the net Socket Get Option operation.
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);
    return KSUCCESS(st) ? provider->GetOption(providerSocket, option, value, valueSize, actualSizeOut) : st;
}

// Implements the net Socket Poll operation.
KSTATUS NetSocketPoll(SOCKET_OBJECT* socket, uint32_t events, uint32_t* returnedEventsOut) {
    if (!returnedEventsOut)
        return STATUS_INVALID_PARAMETER;
    *returnedEventsOut = 0;

    const NET_STACK_PROVIDER* provider = NULL;
    void* providerSocket = NULL;
    KSTATUS st = NetSocketGetProviderSocket(socket, &provider, &providerSocket);
    return KSUCCESS(st) ? provider->Poll(providerSocket, events, returnedEventsOut) : st;
}

// Implements the net Resolve Host Ipv4 operation.
KSTATUS NetResolveHostIpv4(const char* name, SYS_SOCKET_ADDRESS_IPV4_K* addressOut) {
    if (!name || !addressOut)
        return STATUS_INVALID_PARAMETER;

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

    return provider->ResolveHostIpv4(name, addressOut);
}