

#include "internal/quic_rcidm.h"
#include "internal/priority_queue.h"
#include "internal/list.h"
#include "internal/common.h"


static void rcidm_update(QUIC_RCIDM *rcidm);
static void rcidm_set_preferred_rcid(QUIC_RCIDM *rcidm,
                                     const QUIC_CONN_ID *rcid);

#define PACKETS_PER_RCID        10000

#define INITIAL_SEQ_NUM         0
#define PREF_ADDR_SEQ_NUM       1


enum {
    RCID_STATE_PENDING,
    RCID_STATE_CUR,
    RCID_STATE_RETIRING
};

enum {
    RCID_TYPE_INITIAL,      
    RCID_TYPE_PREF_ADDR,    
    RCID_TYPE_NCID          
    
};

typedef struct rcid_st {
    OSSL_LIST_MEMBER(retiring, struct rcid_st); 

    QUIC_CONN_ID    cid;        
    uint64_t        seq_num;
    size_t          pq_idx;     
    unsigned int    state  : 2; 
    unsigned int    type   : 2; 
} RCID;

DEFINE_PRIORITY_QUEUE_OF(RCID);
DEFINE_LIST_OF(retiring, RCID);


struct quic_rcidm_st {
    
    QUIC_CONN_ID                preferred_rcid;

    
    QUIC_CONN_ID                initial_odcid, retry_odcid;

    
    uint64_t                    packets_sent;

    
    uint64_t                    num_changes;

    
    uint64_t                    retire_prior_to;

    
    PRIORITY_QUEUE_OF(RCID)     *rcids;

    
    RCID                        *cur_rcid;

    
    OSSL_LIST(retiring)         retiring_list;

    
    size_t                      num_retiring;

    
    unsigned int    preferred_rcid_changed          : 1;

    
    unsigned int    have_preferred_rcid             : 1;

    
    unsigned int    handshake_complete              : 1;

    
    unsigned int    added_initial_odcid             : 1;
    
    unsigned int    added_retry_odcid               : 1;
    
    unsigned int    added_initial_rcid              : 1;
    
    unsigned int    roll_requested                  : 1;
};


#define MAX_NUMBERED_RCIDS      (SIZE_MAX / 2)

static void rcidm_transition_rcid(QUIC_RCIDM *rcidm, RCID *rcid,
                                  unsigned int state);


static void rcidm_check_rcid(QUIC_RCIDM *rcidm, RCID *rcid)
{
    assert(rcid->state == RCID_STATE_PENDING
           || rcid->state == RCID_STATE_CUR
           || rcid->state == RCID_STATE_RETIRING);
    assert((rcid->state == RCID_STATE_PENDING)
           == (rcid->pq_idx != SIZE_MAX));
    assert((rcid->state == RCID_STATE_CUR)
           == (rcidm->cur_rcid == rcid));
    assert((ossl_list_retiring_next(rcid) != NULL
            || ossl_list_retiring_prev(rcid) != NULL
            || ossl_list_retiring_head(&rcidm->retiring_list) == rcid)
           == (rcid->state == RCID_STATE_RETIRING));
    assert(rcid->type != RCID_TYPE_INITIAL || rcid->seq_num == 0);
    assert(rcid->type != RCID_TYPE_PREF_ADDR || rcid->seq_num == 1);
    assert(rcid->seq_num <= OSSL_QUIC_VLINT_MAX);
    assert(rcid->cid.id_len > 0 && rcid->cid.id_len <= QUIC_MAX_CONN_ID_LEN);
    assert(rcid->seq_num >= rcidm->retire_prior_to
            || rcid->state == RCID_STATE_RETIRING);
    assert(rcidm->num_changes == 0 || rcidm->handshake_complete);
    assert(rcid->state != RCID_STATE_RETIRING || rcidm->num_retiring > 0);
}

static int rcid_cmp(const RCID *a, const RCID *b)
{
    if (a->seq_num < b->seq_num)
        return -1;
    if (a->seq_num > b->seq_num)
        return 1;
    return 0;
}

QUIC_RCIDM *ossl_quic_rcidm_new(const QUIC_CONN_ID *initial_odcid)
{
    QUIC_RCIDM *rcidm;

    if ((rcidm = OPENSSL_zalloc(sizeof(*rcidm))) == NULL)
        return NULL;

    if ((rcidm->rcids = ossl_pqueue_RCID_new(rcid_cmp)) == NULL) {
        OPENSSL_free(rcidm);
        return NULL;
    }

    if (initial_odcid != NULL) {
        rcidm->initial_odcid        = *initial_odcid;
        rcidm->added_initial_odcid  = 1;
    }

    rcidm_update(rcidm);
    return rcidm;
}

void ossl_quic_rcidm_free(QUIC_RCIDM *rcidm)
{
    RCID *rcid, *rnext;

    if (rcidm == NULL)
        return;

    OPENSSL_free(rcidm->cur_rcid);
    while ((rcid = ossl_pqueue_RCID_pop(rcidm->rcids)) != NULL)
        OPENSSL_free(rcid);

    LIST_FOREACH_DELSAFE(rcid, rnext, retiring, &rcidm->retiring_list)
        OPENSSL_free(rcid);

    ossl_pqueue_RCID_free(rcidm->rcids);
    OPENSSL_free(rcidm);
}

static void rcidm_set_preferred_rcid(QUIC_RCIDM *rcidm,
                                     const QUIC_CONN_ID *rcid)
{
    if (rcid == NULL) {
        rcidm->preferred_rcid_changed   = 1;
        rcidm->have_preferred_rcid      = 0;
        return;
    }

    if (ossl_quic_conn_id_eq(&rcidm->preferred_rcid, rcid))
        return;

    rcidm->preferred_rcid           = *rcid;
    rcidm->preferred_rcid_changed   = 1;
    rcidm->have_preferred_rcid      = 1;
}


static RCID *rcidm_create_rcid(QUIC_RCIDM *rcidm, uint64_t seq_num,
                               const QUIC_CONN_ID *cid,
                               unsigned int type)
{
    RCID *rcid;

    if (cid->id_len < 1 || cid->id_len > QUIC_MAX_CONN_ID_LEN
        || seq_num > OSSL_QUIC_VLINT_MAX
        || ossl_pqueue_RCID_num(rcidm->rcids) + rcidm->num_retiring
            > MAX_NUMBERED_RCIDS)
        return NULL;

    if ((rcid = OPENSSL_zalloc(sizeof(*rcid))) == NULL)
        return NULL;

    rcid->seq_num           = seq_num;
    rcid->cid               = *cid;
    rcid->type              = type;

    if (rcid->seq_num >= rcidm->retire_prior_to) {
        rcid->state = RCID_STATE_PENDING;

        if (!ossl_pqueue_RCID_push(rcidm->rcids, rcid, &rcid->pq_idx)) {
            OPENSSL_free(rcid);
            return NULL;
        }
    } else {
        
        rcid->state     = RCID_STATE_RETIRING;
        rcid->pq_idx    = SIZE_MAX;
        ossl_list_retiring_insert_tail(&rcidm->retiring_list, rcid);
        ++rcidm->num_retiring;
    }

    rcidm_check_rcid(rcidm, rcid);
    return rcid;
}

static void rcidm_transition_rcid(QUIC_RCIDM *rcidm, RCID *rcid,
                                  unsigned int state)
{
    unsigned int old_state = rcid->state;

    assert(state >= old_state && state <= RCID_STATE_RETIRING);
    rcidm_check_rcid(rcidm, rcid);
    if (state == old_state)
        return;

    if (rcidm->cur_rcid != NULL && state == RCID_STATE_CUR) {
        rcidm_transition_rcid(rcidm, rcidm->cur_rcid, RCID_STATE_RETIRING);
        assert(rcidm->cur_rcid == NULL);
    }

    if (old_state == RCID_STATE_PENDING) {
        ossl_pqueue_RCID_remove(rcidm->rcids, rcid->pq_idx);
        rcid->pq_idx = SIZE_MAX;
    }

    rcid->state = state;

    if (state == RCID_STATE_CUR) {
        rcidm->cur_rcid = rcid;
    } else if (state == RCID_STATE_RETIRING) {
        if (old_state == RCID_STATE_CUR)
            rcidm->cur_rcid = NULL;

        ossl_list_retiring_insert_tail(&rcidm->retiring_list, rcid);
        ++rcidm->num_retiring;
    }

    rcidm_check_rcid(rcidm, rcid);
}

static void rcidm_free_rcid(QUIC_RCIDM *rcidm, RCID *rcid)
{
    if (rcid == NULL)
        return;

    rcidm_check_rcid(rcidm, rcid);

    switch (rcid->state) {
    case RCID_STATE_PENDING:
        ossl_pqueue_RCID_remove(rcidm->rcids, rcid->pq_idx);
        break;
    case RCID_STATE_CUR:
        rcidm->cur_rcid = NULL;
        break;
    case RCID_STATE_RETIRING:
        ossl_list_retiring_remove(&rcidm->retiring_list, rcid);
        --rcidm->num_retiring;
        break;
    default:
        assert(0);
        break;
    }

    OPENSSL_free(rcid);
}

static void rcidm_handle_retire_prior_to(QUIC_RCIDM *rcidm,
                                         uint64_t retire_prior_to)
{
    RCID *rcid;

    if (retire_prior_to <= rcidm->retire_prior_to)
        return;

    
    if (rcidm->cur_rcid != NULL && rcidm->cur_rcid->seq_num < retire_prior_to)
        rcidm_transition_rcid(rcidm, rcidm->cur_rcid, RCID_STATE_RETIRING);

    
    while ((rcid = ossl_pqueue_RCID_peek(rcidm->rcids)) != NULL
           && rcid->seq_num < retire_prior_to)
        rcidm_transition_rcid(rcidm, rcid, RCID_STATE_RETIRING);

    rcidm->retire_prior_to = retire_prior_to;
}



static void rcidm_roll(QUIC_RCIDM *rcidm)
{
    RCID *rcid;

    if ((rcid = ossl_pqueue_RCID_peek(rcidm->rcids)) == NULL)
        return;

    rcidm_transition_rcid(rcidm, rcid, RCID_STATE_CUR);

    ++rcidm->num_changes;
    rcidm->roll_requested = 0;

    if (rcidm->packets_sent >= PACKETS_PER_RCID)
        rcidm->packets_sent %= PACKETS_PER_RCID;
    else
        rcidm->packets_sent = 0;
}

static void rcidm_update(QUIC_RCIDM *rcidm)
{
    RCID *rcid;

    
    if (rcidm->cur_rcid == NULL
        && (rcid = ossl_pqueue_RCID_peek(rcidm->rcids)) != NULL) {
        rcidm_transition_rcid(rcidm, rcid, RCID_STATE_CUR);
        assert(rcidm->cur_rcid != NULL);
    }

    
    if (rcidm->cur_rcid != NULL) {
        rcidm_check_rcid(rcidm, rcidm->cur_rcid);
        rcidm_set_preferred_rcid(rcidm, &rcidm->cur_rcid->cid);
        return;
    }

    
    if (rcidm->added_retry_odcid && !rcidm->handshake_complete) {
        rcidm_set_preferred_rcid(rcidm, &rcidm->retry_odcid);
        return;
    }

    if (rcidm->added_initial_odcid && !rcidm->handshake_complete) {
        rcidm_set_preferred_rcid(rcidm, &rcidm->initial_odcid);
        return;
    }

    
    rcidm_set_preferred_rcid(rcidm, NULL);
}

static int rcidm_should_roll(QUIC_RCIDM *rcidm)
{
    
    return rcidm->handshake_complete
        && (rcidm->num_changes == 0
            || rcidm->packets_sent >= PACKETS_PER_RCID
            || rcidm->roll_requested);
}

static void rcidm_tick(QUIC_RCIDM *rcidm)
{
    if (rcidm_should_roll(rcidm))
        rcidm_roll(rcidm);

    rcidm_update(rcidm);
}


void ossl_quic_rcidm_on_handshake_complete(QUIC_RCIDM *rcidm)
{
    if (rcidm->handshake_complete)
        return;

    rcidm->handshake_complete = 1;
    rcidm_tick(rcidm);
}

void ossl_quic_rcidm_on_packet_sent(QUIC_RCIDM *rcidm, uint64_t num_packets)
{
    if (num_packets == 0)
        return;

    rcidm->packets_sent += num_packets;
    rcidm_tick(rcidm);
}

void ossl_quic_rcidm_request_roll(QUIC_RCIDM *rcidm)
{
    rcidm->roll_requested = 1;
    rcidm_tick(rcidm);
}


int ossl_quic_rcidm_add_from_initial(QUIC_RCIDM *rcidm,
                                     const QUIC_CONN_ID *rcid)
{
    RCID *rcid_obj;

    if (rcidm->added_initial_rcid || rcidm->handshake_complete)
        return 0;

    rcid_obj = rcidm_create_rcid(rcidm, INITIAL_SEQ_NUM,
                                 rcid, RCID_TYPE_INITIAL);
    if (rcid_obj == NULL)
        return 0;

    rcidm->added_initial_rcid = 1;
    rcidm_tick(rcidm);
    return 1;
}

int ossl_quic_rcidm_add_from_server_retry(QUIC_RCIDM *rcidm,
                                          const QUIC_CONN_ID *retry_odcid)
{
    if (rcidm->added_retry_odcid || rcidm->handshake_complete)
        return 0;

    rcidm->retry_odcid          = *retry_odcid;
    rcidm->added_retry_odcid    = 1;
    rcidm_tick(rcidm);
    return 1;
}

int ossl_quic_rcidm_add_from_ncid(QUIC_RCIDM *rcidm,
                                  const OSSL_QUIC_FRAME_NEW_CONN_ID *ncid)
{
    RCID *rcid;

    rcid = rcidm_create_rcid(rcidm, ncid->seq_num, &ncid->conn_id, RCID_TYPE_NCID);
    if (rcid == NULL)
        return 0;

    rcidm_handle_retire_prior_to(rcidm, ncid->retire_prior_to);
    rcidm_tick(rcidm);
    return 1;
}



static int rcidm_get_retire(QUIC_RCIDM *rcidm, uint64_t *seq_num, int peek)
{
    RCID *rcid = ossl_list_retiring_head(&rcidm->retiring_list);

    if (rcid == NULL)
        return 0;

    if (seq_num != NULL)
        *seq_num = rcid->seq_num;

    if (!peek)
        rcidm_free_rcid(rcidm, rcid);

    return 1;
}

int ossl_quic_rcidm_pop_retire_seq_num(QUIC_RCIDM *rcidm,
                                       uint64_t *seq_num)
{
    return rcidm_get_retire(rcidm, seq_num, 0);
}

int ossl_quic_rcidm_peek_retire_seq_num(QUIC_RCIDM *rcidm,
                                        uint64_t *seq_num)
{
    return rcidm_get_retire(rcidm, seq_num, 1);
}

int ossl_quic_rcidm_get_preferred_tx_dcid(QUIC_RCIDM *rcidm,
                                          QUIC_CONN_ID *tx_dcid)
{
    if (!rcidm->have_preferred_rcid)
        return 0;

    *tx_dcid = rcidm->preferred_rcid;
    return 1;
}

int ossl_quic_rcidm_get_preferred_tx_dcid_changed(QUIC_RCIDM *rcidm,
                                                  int clear)
{
    int r = rcidm->preferred_rcid_changed;

    if (clear)
        rcidm->preferred_rcid_changed = 0;

    return r;
}

size_t ossl_quic_rcidm_get_num_active(const QUIC_RCIDM *rcidm)
{
    return ossl_pqueue_RCID_num(rcidm->rcids)
        + (rcidm->cur_rcid != NULL ? 1 : 0)
        + ossl_quic_rcidm_get_num_retiring(rcidm);
}

size_t ossl_quic_rcidm_get_num_retiring(const QUIC_RCIDM *rcidm)
{
    return rcidm->num_retiring;
}
