#include "transfd_shm.h"

#include <stdio.h>
#include <stdlib.h>
#include <stddef.h>
#include <string.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <errno.h>

int sendfd(int conn, int fd[], int fds){
    if (fds <= 0) return -1;

    struct msghdr msg = {};
    static char dummy = 'x';
    struct iovec iov = {&dummy, sizeof(dummy)};
    msg.msg_iov = &iov;
    msg.msg_iovlen = 1;

    struct cmsghdr* cmsg;

    const unsigned control_len = CMSG_SPACE(sizeof(int) * fds);
    char* control_buffer = (char*)malloc(control_len);

    msg.msg_control = control_buffer;
    msg.msg_controllen = control_len;

    cmsg = CMSG_FIRSTHDR(&msg);
    cmsg->cmsg_level = SOL_SOCKET;
    cmsg->cmsg_type = SCM_RIGHTS;
    cmsg->cmsg_len = CMSG_LEN(sizeof(int) * fds);
    memcpy(CMSG_DATA(cmsg), fd, sizeof(int) * fds);

    int ret = sendmsg(conn, &msg, MSG_NOSIGNAL);
    free(control_buffer);

    if (ret == -1){
        printf("sendmsg error %d -> %s\n", errno, strerror(errno));
        if (errno == EPIPE || errno == ECONNRESET || errno == ENOTCONN)
            return 0;
        return -1;
    }
    return ret;
}

int recvfd(int conn, int fd[], int fds){
    char data;
    struct msghdr msg = {};
    struct iovec iov = {&data, 1};
    msg.msg_iov = &iov;
    msg.msg_iovlen = 1;

    struct cmsghdr *cmsgp;

    const size_t control_len = CMSG_SPACE(sizeof(int) * 16);
    char control_buffer[control_len];

    msg.msg_control = control_buffer;
    msg.msg_controllen = sizeof(control_buffer);

    int ret = recvmsg(conn, &msg, MSG_CMSG_CLOEXEC);
    if (ret == -1){
        printf("recvmsg failed %d -> %s\n", errno, strerror(errno));
        if (errno == ECONNRESET)
            return 0;
        return -1;
    } else if (ret == 0){
        return 0;
    }

    int* wire_fds = NULL;
    unsigned wire_fds_len = 0;
    if (msg.msg_controllen > 0){
        struct cmsghdr* cmsg;
        for (cmsg = CMSG_FIRSTHDR(&msg); cmsg; cmsg = CMSG_NXTHDR(&msg, cmsg)) {
            const unsigned payload_len = cmsg->cmsg_len - CMSG_LEN(0);
            if (cmsg->cmsg_level == SOL_SOCKET && cmsg->cmsg_type == SCM_RIGHTS) {
                wire_fds = (int*)CMSG_DATA(cmsg);
                wire_fds_len = payload_len / sizeof(int);
            }
        }
    }

    const unsigned count = wire_fds_len > fds ? fds : wire_fds_len;
    for (unsigned i = 0; i < count; i++){
        fd[i] = wire_fds[i];
    }

    return wire_fds_len;
}