/*

 * Copyright 2022-2023 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 "internal/quic_stream.h"

#include "internal/uint_set.h"

#include "internal/common.h"

#include "internal/ring_buf.h"

/*

 * ==================================================================

 * QUIC Send Stream

 */

struct quic_sstream_st {

    struct ring_buf ring_buf;

    /*

     * Any logical byte in the stream is in one of these states:

     *

     *   - NEW: The byte has not yet been transmitted, or has been lost and is

     *     in need of retransmission.

     *

     *   - IN_FLIGHT: The byte has been transmitted but is awaiting

     *     acknowledgement. We continue to store the data in case we return

     *     to the NEW state.

     *

     *   - ACKED: The byte has been acknowledged and we can cease storing it.

     *     We do not necessarily cull it immediately, so there may be a delay

     *     between reaching the ACKED state and the buffer space actually being

     *     recycled.

     *

     * A logical byte in the stream is

     *

     *   - in the NEW state if it is in new_set;

     *   - is in the ACKED state if it is in acked_set

     *       (and may or may not have been culled);

     *   - is in the IN_FLIGHT state otherwise.

     *

     * Invariant: No logical byte is ever in both new_set and acked_set.

     */

    UINT_SET        new_set, acked_set;

    /*

     * The current size of the stream is ring_buf.head_offset. If

     * have_final_size is true, this is also the final size of the stream.

     */

    unsigned int    have_final_size     : 1;

    unsigned int    sent_final_size     : 1;

    unsigned int    acked_final_size    : 1;

    unsigned int    cleanse             : 1;

};

static void qss_cull(QUIC_SSTREAM *qss);

QUIC_SSTREAM *ossl_quic_sstream_new(size_t init_buf_size)

{

    QUIC_SSTREAM *qss;

    qss = OPENSSL_zalloc(sizeof(QUIC_SSTREAM));

    if (qss == NULL)

        return NULL;

    ring_buf_init(&qss->ring_buf);

    if (!ring_buf_resize(&qss->ring_buf, init_buf_size, 0)) {

        ring_buf_destroy(&qss->ring_buf, 0);

        OPENSSL_free(qss);

        return NULL;

    }

    ossl_uint_set_init(&qss->new_set);

    ossl_uint_set_init(&qss->acked_set);

    return qss;

}

void ossl_quic_sstream_free(QUIC_SSTREAM *qss)

{

    if (qss == NULL)

        return;

    ossl_uint_set_destroy(&qss->new_set);

    ossl_uint_set_destroy(&qss->acked_set);

    ring_buf_destroy(&qss->ring_buf, qss->cleanse);

    OPENSSL_free(qss);

}

int ossl_quic_sstream_get_stream_frame(QUIC_SSTREAM *qss,

                                       size_t skip,

                                       OSSL_QUIC_FRAME_STREAM *hdr,

                                       OSSL_QTX_IOVEC *iov,

                                       size_t *num_iov)

{

    size_t num_iov_ = 0, src_len = 0, total_len = 0, i;

    uint64_t max_len;

    const unsigned char *src = NULL;

    UINT_SET_ITEM *range = ossl_list_uint_set_head(&qss->new_set);

    if (*num_iov < 2)

        return 0;

    for (i = 0; i < skip && range != NULL; ++i)

        range = ossl_list_uint_set_next(range);

    if (range == NULL) {

        if (i < skip)

            /* Don't return FIN for infinitely increasing skip */

            return 0;

        /* No new bytes to send, but we might have a FIN */

        if (!qss->have_final_size || qss->sent_final_size)

            return 0;

        hdr->offset = qss->ring_buf.head_offset;

        hdr->len    = 0;

        hdr->is_fin = 1;

        *num_iov    = 0;

        return 1;

    }

    /*

     * We can only send a contiguous range of logical bytes in a single

     * stream frame, so limit ourselves to the range of the first set entry.

     *

     * Set entries never have 'adjacent' entries so we don't have to worry

     * about them here.

     */

    max_len = range->range.end - range->range.start + 1;

    for (i = 0;; ++i) {

        if (total_len >= max_len)

            break;

        if (!ring_buf_get_buf_at(&qss->ring_buf,

                                 range->range.start + total_len,

                                 &src, &src_len))

            return 0;

        if (src_len == 0)

            break;

        assert(i < 2);

        if (total_len + src_len > max_len)

            src_len = (size_t)(max_len - total_len);

        iov[num_iov_].buf       = src;

        iov[num_iov_].buf_len   = src_len;

        total_len += src_len;

        ++num_iov_;

    }

    hdr->offset = range->range.start;

    hdr->len    = total_len;

    hdr->is_fin = qss->have_final_size

        && hdr->offset + hdr->len == qss->ring_buf.head_offset;

    *num_iov    = num_iov_;

    return 1;

}

int ossl_quic_sstream_has_pending(QUIC_SSTREAM *qss)

{

    OSSL_QUIC_FRAME_STREAM shdr;

    OSSL_QTX_IOVEC iov[2];

    size_t num_iov = OSSL_NELEM(iov);

    return ossl_quic_sstream_get_stream_frame(qss, 0, &shdr, iov, &num_iov);

}

uint64_t ossl_quic_sstream_get_cur_size(QUIC_SSTREAM *qss)

{

    return qss->ring_buf.head_offset;

}

int ossl_quic_sstream_mark_transmitted(QUIC_SSTREAM *qss,

                                       uint64_t start,

                                       uint64_t end)

{

    UINT_RANGE r;

    r.start = start;

    r.end   = end;

    if (!ossl_uint_set_remove(&qss->new_set, &r))

        return 0;

    return 1;

}

int ossl_quic_sstream_mark_transmitted_fin(QUIC_SSTREAM *qss,

                                           uint64_t final_size)

{

    /*

     * We do not really need final_size since we already know the size of the

     * stream, but this serves as a sanity check.

     */

    if (!qss->have_final_size || final_size != qss->ring_buf.head_offset)

        return 0;

    qss->sent_final_size = 1;

    return 1;

}

int ossl_quic_sstream_mark_lost(QUIC_SSTREAM *qss,

                                uint64_t start,

                                uint64_t end)

{

    UINT_RANGE r;

    r.start = start;

    r.end   = end;

    /*

     * We lost a range of stream data bytes, so reinsert them into the new set,

     * so that they are returned once more by ossl_quic_sstream_get_stream_frame.

     */

    if (!ossl_uint_set_insert(&qss->new_set, &r))

        return 0;

    return 1;

}

int ossl_quic_sstream_mark_lost_fin(QUIC_SSTREAM *qss)

{

    if (qss->acked_final_size)

        /* Does not make sense to lose a FIN after it has been ACKed */

        return 0;

    /* FIN was lost, so we need to transmit it again. */

    qss->sent_final_size = 0;

    return 1;

}

int ossl_quic_sstream_mark_acked(QUIC_SSTREAM *qss,

                                 uint64_t start,

                                 uint64_t end)

{

    UINT_RANGE r;

    r.start = start;

    r.end   = end;

    if (!ossl_uint_set_insert(&qss->acked_set, &r))

        return 0;

    qss_cull(qss);

    return 1;

}

int ossl_quic_sstream_mark_acked_fin(QUIC_SSTREAM *qss)

{

    if (!qss->have_final_size)

        /* Cannot ack final size before we have a final size */

        return 0;

    qss->acked_final_size = 1;

    return 1;

}

void ossl_quic_sstream_fin(QUIC_SSTREAM *qss)

{

    if (qss->have_final_size)

        return;

    qss->have_final_size = 1;

}

int ossl_quic_sstream_get_final_size(QUIC_SSTREAM *qss, uint64_t *final_size)

{

    if (!qss->have_final_size)

        return 0;

    if (final_size != NULL)

        *final_size = qss->ring_buf.head_offset;

    return 1;

}

int ossl_quic_sstream_append(QUIC_SSTREAM *qss,

                             const unsigned char *buf,

                             size_t buf_len,

                             size_t *consumed)

{

    size_t l, consumed_ = 0;

    UINT_RANGE r;

    struct ring_buf old_ring_buf = qss->ring_buf;

    if (qss->have_final_size) {

        *consumed = 0;

        return 0;

    }

    /*

     * Note: It is assumed that ossl_quic_sstream_append will be called during a

     * call to e.g. SSL_write and this function is therefore designed to support

     * such semantics. In particular, the buffer pointed to by buf is only

     * assumed to be valid for the duration of this call, therefore we must copy

     * the data here. We will later copy-and-encrypt the data during packet

     * encryption, so this is a two-copy design. Supporting a one-copy design in

     * the future will require applications to use a different kind of API.

     * Supporting such changes in future will require corresponding enhancements

     * to this code.

     */

    while (buf_len > 0) {

        l = ring_buf_push(&qss->ring_buf, buf, buf_len);

        if (l == 0)

            break;

        buf         += l;

        buf_len     -= l;

        consumed_   += l;

    }

    if (consumed_ > 0) {

        r.start = old_ring_buf.head_offset;

        r.end   = r.start + consumed_ - 1;

        assert(r.end + 1 == qss->ring_buf.head_offset);

        if (!ossl_uint_set_insert(&qss->new_set, &r)) {

            qss->ring_buf = old_ring_buf;

            *consumed = 0;

            return 0;

        }

    }

    *consumed = consumed_;

    return 1;

}

static void qss_cull(QUIC_SSTREAM *qss)

{

    UINT_SET_ITEM *h = ossl_list_uint_set_head(&qss->acked_set);

    /*

     * Potentially cull data from our ring buffer. This can happen once data has

     * been ACKed and we know we are never going to have to transmit it again.

     *

     * Since we use a ring buffer design for simplicity, we cannot cull byte n +

     * k (for k > 0) from the ring buffer until byte n has also been culled.

     * This means if parts of the stream get acknowledged out of order we might

     * keep around some data we technically don't need to for a while. The

     * impact of this is likely to be small and limited to quite a short

     * duration, and doesn't justify the use of a more complex design.

     */

    /*

     * We only need to check the first range entry in the integer set because we

     * can only cull contiguous areas at the start of the ring buffer anyway.

     */

    if (h != NULL)

        ring_buf_cpop_range(&qss->ring_buf, h->range.start, h->range.end,

                            qss->cleanse);

}

int ossl_quic_sstream_set_buffer_size(QUIC_SSTREAM *qss, size_t num_bytes)

{

    return ring_buf_resize(&qss->ring_buf, num_bytes, qss->cleanse);

}

size_t ossl_quic_sstream_get_buffer_size(QUIC_SSTREAM *qss)

{

    return qss->ring_buf.alloc;

}

size_t ossl_quic_sstream_get_buffer_used(QUIC_SSTREAM *qss)

{

    return ring_buf_used(&qss->ring_buf);

}

size_t ossl_quic_sstream_get_buffer_avail(QUIC_SSTREAM *qss)

{

    return ring_buf_avail(&qss->ring_buf);

}

int ossl_quic_sstream_is_totally_acked(QUIC_SSTREAM *qss)

{

    UINT_RANGE r;

    uint64_t cur_size;

    if (qss->have_final_size && !qss->acked_final_size)

        return 0;

    if (ossl_quic_sstream_get_cur_size(qss) == 0)

        return 1;

    if (ossl_list_uint_set_num(&qss->acked_set) != 1)

        return 0;

    r = ossl_list_uint_set_head(&qss->acked_set)->range;

    cur_size = qss->ring_buf.head_offset;

    /*

     * The invariants of UINT_SET guarantee a single list element if we have a

     * single contiguous range, which is what we should have if everything has

     * been acked.

     */

    assert(r.end + 1 <= cur_size);

    return r.start == 0 && r.end + 1 == cur_size;

}

void ossl_quic_sstream_adjust_iov(size_t len,

                                  OSSL_QTX_IOVEC *iov,

                                  size_t num_iov)

{

    size_t running = 0, i, iovlen;

    for (i = 0, running = 0; i < num_iov; ++i) {

        iovlen = iov[i].buf_len;

        if (running >= len)

            iov[i].buf_len = 0;

        else if (running + iovlen > len)

            iov[i].buf_len = len - running;

        running += iovlen;

    }

}

void ossl_quic_sstream_set_cleanse(QUIC_SSTREAM *qss, int cleanse)

{

    qss->cleanse = cleanse;

}