/* SPDX-License-Identifier: GPL-3.0-or-later */

#include "p2p.h"

#include <errno.h>
#include <netinet/in.h>
#include <stdint.h>
#include <sys/poll.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <unistd.h>

// Create a UDP socket with DF set.
static int create_socket() {
    int sockfd = socket(AF_INET, SOCK_DGRAM, 0);
    if (sockfd < 0) {
        return -1;
    }

    int value = IP_PMTUDISC_DO;
    if (setsockopt(sockfd, IPPROTO_IP, IP_MTU_DISCOVER, &value, sizeof(value)) < 0) {
        close(sockfd);
        return -1;
    }

    return sockfd;
}

int p2p_bind(uint16_t port, in_addr_t ip) {
    int sockfd = create_socket();
    if (sockfd < 0) {
        goto err;
    }

    struct sockaddr_in local_addr = {
        .sin_family = AF_INET,
        .sin_port = htons(port),
        .sin_addr = { .s_addr = INADDR_ANY },
    };
    if (bind(sockfd, (struct sockaddr*) &local_addr, sizeof(local_addr)) < 0) {
        goto err_close;
    }

    while (true) {
        // Wait for handshake from peer
        struct sockaddr_in remote_addr;
        socklen_t remote_addr_len = sizeof(remote_addr);

        uint8_t handshake[1];
        ssize_t handshake_len = recvfrom(sockfd, handshake, sizeof(handshake), 0,
            (struct sockaddr*) &remote_addr, &remote_addr_len);
        if (handshake_len < 0) {
            goto err_close;
        }

        if (handshake_len != 1 || *handshake != '?') {
            // Malformed handshake, ignore
            continue;
        }

        if (remote_addr.sin_addr.s_addr != ip) {
            // Wrong peer, reject handshake
            sendto(sockfd, "N", 1, 0, (struct sockaddr*) &remote_addr, remote_addr_len);
            continue;
        }

        // Accept handshake and lock into peer
        if (connect(sockfd, (struct sockaddr*) &remote_addr, remote_addr_len) < 0) {
            goto err_close;
        }

        if (send(sockfd, "Y", 1, 0) < 0) {
            goto err_close;
        }

        break;
    }

    return sockfd;
err_close:
    close(sockfd);
err:
    return -1;
}

int p2p_connect(uint16_t port, in_addr_t ip, int timeout) {
    int sockfd = create_socket();
    if (sockfd < 0) {
        goto err;
    }

    struct sockaddr_in remote_addr = {
        .sin_family = AF_INET,
        .sin_port = htons(port),
        .sin_addr = { .s_addr = ip },
    };
    if (connect(sockfd, (struct sockaddr*) &remote_addr, sizeof(remote_addr)) < 0) {
        goto err_close;
    }

    // Perform handshake
    if (send(sockfd, "?", 1, 0) < 0) {
        goto err_close;
    }

    struct pollfd pfd = {
        .fd = sockfd,
        .events = POLLIN,
    };
    int ret = poll(&pfd, 1, timeout);
    if (ret < 0) {
        goto err_close;
    } else if (ret == 0) {
        errno = ETIMEDOUT;
        goto err_close;
    }

    uint8_t handshake[1];
    ssize_t handshake_len = recv(sockfd, handshake, sizeof(handshake), 0);
    if (handshake_len < 0) {
        goto err_close;
    } else if (handshake_len != 1 || *handshake != 'Y') {
        errno = ECONNREFUSED;
        goto err_close;
    }

    return sockfd;
err_close:
    close(sockfd);
err:
    return -1;
}