/*

 * Copyright 2024-2025 The OpenSSL Project Authors. All Rights Reserved.

 *

 * Licensed under the Apache License 2.0 (the "License").  You may not use

 * this file except in compliance with the License.  You can obtain a copy

 * in the file LICENSE in the source distribution or at

 * https://www.openssl.org/source/license.html

 */

#include <assert.h>

#include <errno.h>

#include "internal/safe_math.h"

#include "poll_builder.h"

OSSL_SAFE_MATH_UNSIGNED(size_t, size_t)

int ossl_rio_poll_builder_init(RIO_POLL_BUILDER *rpb)

{

#if RIO_POLL_METHOD == RIO_POLL_METHOD_NONE

    return 0;

#elif RIO_POLL_METHOD == RIO_POLL_METHOD_SELECT

    FD_ZERO(&rpb->rfd);

    FD_ZERO(&rpb->wfd);

    FD_ZERO(&rpb->efd);

    rpb->hwm_fd = -1;

#elif RIO_POLL_METHOD == RIO_POLL_METHOD_POLL

    rpb->pfd_heap = NULL;

    rpb->pfd_num = 0;

    rpb->pfd_alloc = OSSL_NELEM(rpb->pfds);

#endif

    return 1;

}

void ossl_rio_poll_builder_cleanup(RIO_POLL_BUILDER *rpb)

{

    if (rpb == NULL)

        return;

#if RIO_POLL_METHOD == RIO_POLL_METHOD_POLL

    OPENSSL_free(rpb->pfd_heap);

#endif

}

#if RIO_POLL_METHOD == RIO_POLL_METHOD_POLL

static int rpb_ensure_alloc(RIO_POLL_BUILDER *rpb, size_t alloc)

{

    struct pollfd *pfd_heap_new;

    size_t total_size;

    int error = 0;

    if (alloc <= rpb->pfd_alloc)

        return 1;

    total_size = safe_mul_size_t(alloc, sizeof(struct pollfd), &error);

    if (error)

        return 0;

    pfd_heap_new = OPENSSL_realloc(rpb->pfd_heap, total_size);

    if (pfd_heap_new == NULL)

        return 0;

    if (rpb->pfd_heap == NULL) {

        /* Copy the contents of the stacked array. */

        memcpy(pfd_heap_new, rpb->pfds, sizeof(rpb->pfds));

    }

    rpb->pfd_heap = pfd_heap_new;

    rpb->pfd_alloc = alloc;

    return 1;

}

#endif

int ossl_rio_poll_builder_add_fd(RIO_POLL_BUILDER *rpb, int fd,

    int want_read, int want_write)

{

#if RIO_POLL_METHOD == RIO_POLL_METHOD_POLL

    size_t i;

    struct pollfd *pfds = (rpb->pfd_heap != NULL ? rpb->pfd_heap : rpb->pfds);

    struct pollfd *pfd;

#endif

    if (fd < 0)

        return 0;

#if RIO_POLL_METHOD == RIO_POLL_METHOD_SELECT

#ifndef OPENSSL_SYS_WINDOWS

    /*

     * On Windows there is no relevant limit to the magnitude of a fd value (see

     * above). On *NIX the fd_set uses a bitmap and we must check the limit.

     */

    if (fd >= FD_SETSIZE)

        return 0;

#endif

    if (want_read)

        openssl_fdset(fd, &rpb->rfd);

    if (want_write)

        openssl_fdset(fd, &rpb->wfd);

    openssl_fdset(fd, &rpb->efd);

    if (fd > rpb->hwm_fd)

        rpb->hwm_fd = fd;

    return 1;

#elif RIO_POLL_METHOD == RIO_POLL_METHOD_POLL

    for (i = 0; i < rpb->pfd_num; ++i) {

        pfd = &pfds[i];

        if (pfd->fd == -1 || pfd->fd == fd)

            break;

    }

    if (i >= rpb->pfd_alloc) {

        if (!rpb_ensure_alloc(rpb, rpb->pfd_alloc * 2))

            return 0;

        pfds = rpb->pfd_heap;

    }

    assert((rpb->pfd_heap != NULL && rpb->pfd_heap == pfds) || (rpb->pfd_heap == NULL && rpb->pfds == pfds));

    assert(i <= rpb->pfd_num && rpb->pfd_num <= rpb->pfd_alloc);

    /* Check the index first because an appended entry is uninitialised. */

    if (i == rpb->pfd_num || pfds[i].fd == -1)

        pfds[i].events = 0;

    pfds[i].fd = fd;

    if (want_read)

        pfds[i].events |= POLLIN;

    if (want_write)

        pfds[i].events |= POLLOUT;

    if (i == rpb->pfd_num)

        ++rpb->pfd_num;

    return 1;

#endif

}

/* Returns 1 if no file descriptors have been added to the poll builder. */

static int rio_poll_builder_is_empty(const RIO_POLL_BUILDER *rpb)

{

#if RIO_POLL_METHOD == RIO_POLL_METHOD_SELECT

    return rpb->hwm_fd < 0;

#elif RIO_POLL_METHOD == RIO_POLL_METHOD_POLL

    return rpb->pfd_num == 0;

#else

    return 1;

#endif

}

int ossl_rio_poll_builder_poll(RIO_POLL_BUILDER *rpb, OSSL_TIME deadline)

{

    int rc;

    /*

     * Waiting with no file descriptors is legitimate: a DTLS connection whose

     * BIO cannot provide a poll descriptor has no readiness to wait for, but

     * still has a retransmission deadline to wake for. Handle that here rather

     * than leaving it to the OS, because poll() treats an empty descriptor set

     * as a plain sleep whereas Windows' select() rejects it outright.

     *

     * With no descriptors and no deadline nothing could ever wake us, so that

     * is a caller error rather than an indefinite sleep.

     */

    if (rio_poll_builder_is_empty(rpb)) {

        if (ossl_time_is_infinite(deadline))

            return 0;

        OSSL_sleep(ossl_time2ms(ossl_time_subtract(deadline, ossl_time_now())));

        return 1;

    }

#if RIO_POLL_METHOD == RIO_POLL_METHOD_SELECT

    do {

        struct timeval timeout, *p_timeout = &timeout;

        /*

         * select expects a timeout, not a deadline, so do the conversion.

         * Update for each call to ensure the correct value is used if we repeat

         * due to EINTR.

         */

        if (ossl_time_is_infinite(deadline))

            p_timeout = NULL;

        else

            /*

             * ossl_time_subtract saturates to zero so we don't need to check if

             * now > deadline.

             */

            timeout = ossl_time_to_timeval(ossl_time_subtract(deadline,

                ossl_time_now()));

        rc = select(rpb->hwm_fd + 1, &rpb->rfd, &rpb->wfd, &rpb->efd, p_timeout);

    } while (rc == -1 && get_last_socket_error_is_eintr());

#elif RIO_POLL_METHOD == RIO_POLL_METHOD_POLL

    do {

        int timeout_ms;

        if (ossl_time_is_infinite(deadline))

            timeout_ms = -1;

        else

            timeout_ms = ossl_time2ms(ossl_time_subtract(deadline,

                ossl_time_now()));

        rc = poll(rpb->pfd_heap != NULL ? rpb->pfd_heap : rpb->pfds,

            rpb->pfd_num, timeout_ms);

    } while (rc == -1 && get_last_socket_error_is_eintr());

#endif

    return rc < 0 ? 0 : 1;

}