#include "UdpVpn.hpp"

#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <string.h>
#include <poll.h>
#include <errno.h>

UdpVpn::UdpVpn()
    : _bound(false), _stopped(false), _has_peer(false), _tun_dev("cvpn%d")
{
    memset(&_serv_addr, 0, sizeof(_serv_addr));
    memset(&_peer_addr, 0, sizeof(_peer_addr));

    _socket = socket(AF_INET6, SOCK_DGRAM, 0);
    if(_socket < 0)
        throw UdpVpn::InitializationError("Cannot create socket", errno, true);
}

UdpVpn::~UdpVpn() {
    close(_socket);
}

void UdpVpn::bind(in_port_t port) {
    bind(in6addr_any, port);
}
void UdpVpn::bind(const struct in6_addr& bind_addr6, in_port_t port) {
    int rc;

    _serv_addr.sin6_family = AF_INET6;
    _serv_addr.sin6_port = htons(port);
    _serv_addr.sin6_addr = bind_addr6;

    rc = ::bind(
            _socket, (const struct sockaddr*)&_serv_addr, sizeof(_serv_addr));
    if(rc < 0) {
        throw UdpVpn::InitializationError("Cannot bind socket", errno, true);
    }

    debugf("> Listening on port %d\n", port);
    _bound = true;
}

void UdpVpn::set_peer(const sockaddr_in6& peer_addr) {
    memcpy(&_peer_addr, &peer_addr, sizeof(_peer_addr));
    _has_peer = true;

    char peer_addr_str[INET6_ADDRSTRLEN];
    inet_ntop(
            AF_INET6, &(peer_addr.sin6_addr), peer_addr_str, INET6_ADDRSTRLEN);
    debugf("Set peer to %s:%d\n", peer_addr_str, ntohs(peer_addr.sin6_port));
}

void UdpVpn::run() {
    int rc;
    int start_at_fd = 0;  // read from polled fds in round-robin fashion
    int cur_fd;
    int nfds = 2;
    struct pollfd poll_fds[2];

    // poll_fds[0]: tun device
    poll_fds[0].fd = _tun_dev.get_fd();
    poll_fds[0].events = POLLIN;

    // poll_fds[1]: UDP socket device
    poll_fds[1].fd = _socket;
    poll_fds[1].events = POLLIN;

    while(!_stopped) {
        rc = poll(poll_fds, nfds, -1);

        if(rc < 0) {
            if(errno == EINTR)  // Interrupt.
                continue;
            throw UdpVpn::NetError(
                    "Error polling from interface", errno, true);
        }
        else if(rc == 0)  // Nothing to read
            continue;

        cur_fd = start_at_fd;
        do {
            if(poll_fds[cur_fd].revents & POLLIN) {
                if(cur_fd == 0)
                    receive_from_tun();
                else if(cur_fd == 1)
                    receive_from_udp();
                break;
            }

            cur_fd = (cur_fd + 1) % nfds;
        } while(cur_fd != start_at_fd);

        start_at_fd = (start_at_fd + 1) % nfds;
    }
}

void UdpVpn::receive_from_tun() {
    // We know that there is data available -- use `read()`
    char buffer[1500];
    size_t nread = _tun_dev.read(buffer, 1500);

    if(nread == 0)
        return;

    kdebugf("Transmitting packet of size %d to peer\n", nread);
    send_over_udp(buffer, nread);
}

void UdpVpn::receive_from_udp() {
    ssize_t nread;
    char buffer[1500];
    struct sockaddr_in6 peer_addr;
    socklen_t peer_addr_len = sizeof(peer_addr);

    nread = recvfrom(_socket, buffer, 1500, MSG_WAITALL,
            (struct sockaddr*) &peer_addr, &peer_addr_len);

    if(nread < 0)
        throw UdpVpn::NetError("Cannot receive datagram", errno, true);
    if(nread == 0)
        return;

    if(peer_addr.sin6_family != AF_INET6) {
        debugf("WARNING: Received non-ipv6 family datagram %d. Ignoring.\n",
                peer_addr.sin6_family);
        return;
    }
    if(peer_addr_len != sizeof(peer_addr)) {
        debugf("WARNING: received unexpected source address length %u."
                "Ignoring.\n",
                peer_addr_len);
        return;
    }

    set_peer(peer_addr);

    // Reinject into tun
    kdebugf("Receiving packet of size %d from peer\n", nread);
    _tun_dev.write(buffer, nread);
}

size_t UdpVpn::send_over_udp(const char* data, size_t len) {
    ssize_t nsent;

    if(!_has_peer) {
        debugf("Dropping packet to be transmitted: no peer.\n");
        return 0;
    }

    nsent = sendto(_socket, data, len, MSG_CONFIRM,
            (const struct sockaddr*) &_peer_addr, sizeof(_peer_addr));
    if(nsent < 0)
        throw NetError("Could not send UDP packet", errno, true);

    return (size_t) nsent;
}