udp.hpp

#pragma once


#include <type_traits>
#include <algorithm>
#include <iterator>
#include <utility>
#include <cstddef>
#include <memory>
#include <string>
#include <uv.h>
#include "request.hpp"
#include "handle.hpp"
#include "util.hpp"


namespace uvw {


struct SendEvent {};


struct UDPDataEvent {
    explicit UDPDataEvent(Addr sndr, std::unique_ptr<const char[]> buf, std::size_t len, bool part) noexcept
        : data{std::move(buf)}, length{len}, sender{std::move(sndr)}, partial{part}
    {}

    std::unique_ptr<const char[]> data; 
    std::size_t length;  
    Addr sender; 
    bool partial; 
};


namespace details {


enum class UVUdpFlags: std::underlying_type_t<uv_udp_flags> {
    IPV6ONLY = UV_UDP_IPV6ONLY,
    REUSEADDR = UV_UDP_REUSEADDR
};


enum class UVMembership: std::underlying_type_t<uv_membership> {
    LEAVE_GROUP = UV_LEAVE_GROUP,
    JOIN_GROUP = UV_JOIN_GROUP
};


class SendReq final: public Request<SendReq, uv_udp_send_t> {
public:
    using Deleter = void(*)(char *);

    SendReq(ConstructorAccess ca, std::shared_ptr<Loop> loop, std::unique_ptr<char[], Deleter> dt, unsigned int len)
        : Request<SendReq, uv_udp_send_t>{ca, std::move(loop)},
          data{std::move(dt)},
          buf{uv_buf_init(data.get(), len)}
    {}

    void send(uv_udp_t *handle, const struct sockaddr* addr) {
        invoke(&uv_udp_send, get(), handle, &buf, 1, addr, &defaultCallback<SendEvent>);
    }

private:
    std::unique_ptr<char[], Deleter> data;
    uv_buf_t buf;
};


}


class UDPHandle final: public Handle<UDPHandle, uv_udp_t> {
    template<typename I>
    static void recvCallback(uv_udp_t *handle, ssize_t nread, const uv_buf_t *buf, const sockaddr *addr, unsigned flags) {
        const typename details::IpTraits<I>::Type *aptr = reinterpret_cast<const typename details::IpTraits<I>::Type *>(addr);

        UDPHandle &udp = *(static_cast<UDPHandle*>(handle->data));
        // data will be destroyed no matter of what the value of nread is
        std::unique_ptr<const char[]> data{buf->base};

        if(nread > 0) {
            // data available (can be truncated)
            udp.publish(UDPDataEvent{details::address<I>(aptr), std::move(data), static_cast<std::size_t>(nread), !(0 == (flags & UV_UDP_PARTIAL))});
        } else if(nread == 0 && addr == nullptr) {
            // no more data to be read, doing nothing is fine
        } else if(nread == 0 && addr != nullptr) {
            // empty udp packet
            udp.publish(UDPDataEvent{details::address<I>(aptr), std::move(data), static_cast<std::size_t>(nread), false});
        } else {
            // transmission error
            udp.publish(ErrorEvent(nread));
        }
    }

public:
    using Membership = details::UVMembership;
    using Bind = details::UVUdpFlags;
    using IPv4 = uvw::IPv4;
    using IPv6 = uvw::IPv6;

    using Handle::Handle;

    explicit UDPHandle(ConstructorAccess ca, std::shared_ptr<Loop> ref, unsigned int f)
        : Handle{ca, std::move(ref)}, tag{FLAGS}, flags{f}
    {}

    bool init() {
        return (tag == FLAGS)
                ? initialize(&uv_udp_init_ex, flags)
                : initialize(&uv_udp_init);
    }

    void open(OSSocketHandle socket) {
        invoke(&uv_udp_open, get(), socket);
    }

    template<typename I = IPv4>
    void bind(std::string ip, unsigned int port, Flags<Bind> opts = Flags<Bind>{}) {
        typename details::IpTraits<I>::Type addr;
        details::IpTraits<I>::addrFunc(ip.data(), port, &addr);
        invoke(&uv_udp_bind, get(), reinterpret_cast<const sockaddr *>(&addr), opts);
    }

    template<typename I = IPv4>
    void bind(Addr addr, Flags<Bind> opts = Flags<Bind>{}) {
        bind<I>(addr.ip, addr.port, opts);
    }

    template<typename I = IPv4>
    Addr sock() const noexcept {
        return details::address<I>(&uv_udp_getsockname, get());
    }

    template<typename I = IPv4>
    bool multicastMembership(std::string multicast, std::string iface, Membership membership) {
        return (0 == uv_udp_set_membership(get(), multicast.data(), iface.data(), static_cast<uv_membership>(membership)));
    }

    bool multicastLoop(bool enable = true) {
        return (0 == uv_udp_set_multicast_loop(get(), enable));
    }

    bool multicastTtl(int val) {
        return (0 == uv_udp_set_multicast_ttl(get(), val > 255 ? 255 : val));
    }

    template<typename I = IPv4>
    bool multicastInterface(std::string iface) {
        return (0 == uv_udp_set_multicast_interface(get(), iface.data()));
    }

    bool broadcast(bool enable = false) {
        return (0 == uv_udp_set_broadcast(get(), enable));
    }

    bool ttl(int val) {
        return (0 == uv_udp_set_ttl(get(), val > 255 ? 255 : val));
    }

    template<typename I = IPv4>
    void send(std::string ip, unsigned int port, std::unique_ptr<char[]> data, unsigned int len) {
        typename details::IpTraits<I>::Type addr;
        details::IpTraits<I>::addrFunc(ip.data(), port, &addr);

        auto req = loop().resource<details::SendReq>(
                    std::unique_ptr<char[], details::SendReq::Deleter>{
                        data.release(), [](char *ptr) { delete[] ptr; }
                    }, len);

        auto listener = [ptr = shared_from_this()](const auto &event, const auto &) {
            ptr->publish(event);
        };

        req->once<ErrorEvent>(listener);
        req->once<SendEvent>(listener);
        req->send(get(), reinterpret_cast<const sockaddr *>(&addr));
    }

    template<typename I = IPv4>
    void send(Addr addr, std::unique_ptr<char[]> data, unsigned int len) {
        send<I>(addr.ip, addr.port, std::move(data), len);
    }

    template<typename I = IPv4>
    void send(std::string ip, unsigned int port, char *data, unsigned int len) {
        typename details::IpTraits<I>::Type addr;
        details::IpTraits<I>::addrFunc(ip.data(), port, &addr);

        auto req = loop().resource<details::SendReq>(
                    std::unique_ptr<char[], details::SendReq::Deleter>{
                        data, [](char *) {}
                    }, len);

        auto listener = [ptr = shared_from_this()](const auto &event, const auto &) {
            ptr->publish(event);
        };

        req->once<ErrorEvent>(listener);
        req->once<SendEvent>(listener);
        req->send(get(), reinterpret_cast<const sockaddr *>(&addr));
    }

    template<typename I = IPv4>
    void send(Addr addr, char *data, unsigned int len) {
        send<I>(addr.ip, addr.port, data, len);
    }

    template<typename I = IPv4>
    int trySend(std::string ip, unsigned int port, std::unique_ptr<char[]> data, unsigned int len) {
        typename details::IpTraits<I>::Type addr;
        details::IpTraits<I>::addrFunc(ip.data(), port, &addr);

        uv_buf_t bufs[] = { uv_buf_init(data.get(), len) };
        auto bw = uv_udp_try_send(get(), bufs, 1, reinterpret_cast<const sockaddr *>(&addr));

        if(bw < 0) {
            publish(ErrorEvent{bw});
            bw = 0;
        }

        return bw;
    }

    template<typename I = IPv4>
    int trySend(Addr addr, std::unique_ptr<char[]> data, unsigned int len) {
        return trySend<I>(addr.ip, addr.port, std::move(data), len);
    }

    template<typename I = IPv4>
    int trySend(std::string ip, unsigned int port, char *data, unsigned int len) {
        typename details::IpTraits<I>::Type addr;
        details::IpTraits<I>::addrFunc(ip.data(), port, &addr);

        uv_buf_t bufs[] = { uv_buf_init(data, len) };
        auto bw = uv_udp_try_send(get(), bufs, 1, reinterpret_cast<const sockaddr *>(&addr));

        if(bw < 0) {
            publish(ErrorEvent{bw});
            bw = 0;
        }

        return bw;
    }

    template<typename I = IPv4>
    int trySend(Addr addr, char *data, unsigned int len) {
        return trySend<I>(addr.ip, addr.port, data, len);
    }

    template<typename I = IPv4>
    void recv() {
        invoke(&uv_udp_recv_start, get(), &allocCallback, &recvCallback<I>);
    }

    void stop() {
        invoke(&uv_udp_recv_stop, get());
    }

private:
    enum { DEFAULT, FLAGS } tag{DEFAULT};
    unsigned int flags{};
};


}