/* SPDX-License-Identifier: GPL-3.0-only */
#include "web_session_store.h"

#include <limits.h>
#include <string.h>
#include "esp_timer.h"
#include "freertos/FreeRTOS.h"
#include "mbedtls/sha256.h"
#include "secure_random.h"

typedef struct {
    web_session_id_t id;
    int64_t expires_at_us;
    user_principal_t principal;
    uint8_t token_digest[WEB_SESSION_STORE_SECRET_BYTES];
    uint8_t origin_digest[WEB_SESSION_STORE_SECRET_BYTES];
    uint8_t csrf[WEB_SESSION_STORE_SECRET_BYTES];
} session_entry_t;

static portMUX_TYPE s_lock = portMUX_INITIALIZER_UNLOCKED;
static struct {
    session_entry_t entries[WEB_SESSION_STORE_CAPACITY];
    uint64_t next_id;
    uint64_t epoch;
    bool ready;
    bool initializing;
    uint32_t issued;
    uint32_t capacity_rejections;
    uint32_t expired;
    uint32_t invalidated;
    uint32_t lookup_rejections;
    uint32_t init_failures;
} s_state;

static bool equal_bytes(const uint8_t *a, const uint8_t *b, size_t length)
{
    uint8_t difference = 0U;
    for (size_t i = 0U; i < length; ++i) {
        difference |= a[i] ^ b[i];
    }
    return difference == 0U;
}

static void encode_hex(const uint8_t *bytes, char *text)
{
    static const char hex[] = "0123456789abcdef";
    for (size_t i = 0U; i < WEB_SESSION_STORE_SECRET_BYTES; ++i) {
        text[2U * i] = hex[bytes[i] >> 4U];
        text[2U * i + 1U] = hex[bytes[i] & 15U];
    }
    text[WEB_SESSION_STORE_TOKEN_LENGTH] = '\0';
}

static esp_err_t digest(const void *input, size_t length, uint8_t *output)
{
    return mbedtls_sha256(input, length, output, 0) == 0 ? ESP_OK : ESP_FAIL;
}

static esp_err_t origin_digest(const char *origin, size_t length, uint8_t *output)
{
    if (origin == NULL || length <= 8U ||
        length > WEB_SESSION_STORE_ORIGIN_MAX_LENGTH ||
        memcmp(origin, "https://", 8U) != 0 || memchr(origin, '\0', length) != NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    return digest(origin, length, output);
}

static session_entry_t *find_locked(web_session_id_t id)
{
    if (id != 0U) {
        for (size_t i = 0U; i < WEB_SESSION_STORE_CAPACITY; ++i) {
            if (s_state.entries[i].id == id) {
                return &s_state.entries[i];
            }
        }
    }
    return NULL;
}

static void retire_locked(session_entry_t *entry, bool expired)
{
    if (entry->id != 0U) {
        if (expired) {
            ++s_state.expired;
        } else {
            ++s_state.invalidated;
        }
        secure_wipe(entry, sizeof(*entry));
    }
}

static void expire_locked(int64_t now)
{
    for (size_t i = 0U; i < WEB_SESSION_STORE_CAPACITY; ++i) {
        session_entry_t *entry = &s_state.entries[i];
        if (entry->id != 0U && entry->expires_at_us <= now) {
            retire_locked(entry, true);
        }
    }
}

/* Also cancels issuance already outside the lock, even if no record matched.
 * Exhaustion is fail-closed rather than allowing an epoch/identity ABA. */
static void advance_epoch_locked(void)
{
    if (s_state.epoch != UINT64_MAX) {
        ++s_state.epoch;
    } else {
        s_state.ready = false;
    }
}

static void export_view(const session_entry_t *entry, web_session_view_t *view)
{
    view->id = entry->id;
    view->expires_at_us = entry->expires_at_us;
    view->principal = entry->principal;
    encode_hex(entry->csrf, view->csrf);
}

esp_err_t web_session_store_init(void)
{
    taskENTER_CRITICAL(&s_lock);
    if (s_state.ready) {
        taskEXIT_CRITICAL(&s_lock);
        return ESP_OK;
    }
    if (s_state.initializing || s_state.epoch == UINT64_MAX) {
        taskEXIT_CRITICAL(&s_lock);
        return ESP_ERR_INVALID_STATE;
    }
    s_state.initializing = true;
    uint64_t epoch = s_state.epoch;
    taskEXIT_CRITICAL(&s_lock);

    uint8_t probe[WEB_SESSION_STORE_SECRET_BYTES] = {0};
    esp_err_t error = secure_random_fill(probe, sizeof(probe));
    secure_wipe(probe, sizeof(probe));

    taskENTER_CRITICAL(&s_lock);
    if (epoch != s_state.epoch) {
        error = ESP_ERR_INVALID_STATE;
    }
    s_state.initializing = false;
    s_state.ready = error == ESP_OK;
    if (error != ESP_OK) {
        ++s_state.init_failures;
    }
    taskEXIT_CRITICAL(&s_lock);
    return error;
}

void web_session_store_stop(void)
{
    taskENTER_CRITICAL(&s_lock);
    advance_epoch_locked();
    s_state.ready = false;
    for (size_t i = 0U; i < WEB_SESSION_STORE_CAPACITY; ++i) {
        retire_locked(&s_state.entries[i], false);
    }
    taskEXIT_CRITICAL(&s_lock);
}

/* Never enter the database while holding our lock. On return, re-find the
 * non-reused ID and deadline: logout/stop/slot reuse may have raced the call. */
static esp_err_t resolve(web_session_id_t id, web_session_view_t *view,
                         const user_principal_t *expected)
{
    session_entry_t candidate = {0};
    int64_t now = esp_timer_get_time();
    taskENTER_CRITICAL(&s_lock);
    expire_locked(now);
    session_entry_t *entry = find_locked(id);
    esp_err_t error = s_state.ready ? ESP_ERR_NOT_FOUND : ESP_ERR_INVALID_STATE;
    if (s_state.ready && entry != NULL) {
        candidate = *entry;
        error = ESP_OK;
    }
    taskEXIT_CRITICAL(&s_lock);

    if (error == ESP_OK && expected != NULL &&
        (candidate.principal.user_id != expected->user_id ||
         candidate.principal.auth_generation != expected->auth_generation ||
         candidate.principal.role != expected->role ||
         candidate.principal.method != expected->method ||
         candidate.principal.username_length != expected->username_length ||
         memcmp(candidate.principal.username, expected->username,
                candidate.principal.username_length) != 0)) {
        error = ESP_ERR_NOT_FOUND;
    }
    if (error == ESP_OK) {
        bool current = false;
        error = user_database_principal_is_current(&candidate.principal, &current);
        now = esp_timer_get_time();
        taskENTER_CRITICAL(&s_lock);
        expire_locked(now);
        entry = find_locked(id);
        if (!s_state.ready || entry == NULL) {
            error = ESP_ERR_NOT_FOUND;
        } else if (error != ESP_OK || !current) {
            retire_locked(entry, false);
            if (error == ESP_OK) {
                error = ESP_ERR_NOT_FOUND;
            }
        }
        taskEXIT_CRITICAL(&s_lock);
    }
    if (error == ESP_OK && view != NULL) {
        export_view(&candidate, view);
    }
    if (error != ESP_OK) {
        taskENTER_CRITICAL(&s_lock);
        ++s_state.lookup_rejections;
        taskEXIT_CRITICAL(&s_lock);
    }
    secure_wipe(&candidate, sizeof(candidate));
    return error;
}

void web_session_store_prune(void)
{
    for (size_t i = 0U; i < WEB_SESSION_STORE_CAPACITY; ++i) {
        taskENTER_CRITICAL(&s_lock);
        web_session_id_t id = s_state.entries[i].id;
        taskEXIT_CRITICAL(&s_lock);
        if (id != 0U) {
            (void)resolve(id, NULL, NULL);
        }
    }
}

esp_err_t web_session_store_issue(
    const user_principal_t *principal, const char *origin, size_t origin_length,
    char token[WEB_SESSION_STORE_TOKEN_LENGTH + 1U], web_session_view_t *view)
{
    if (token != NULL) {
        secure_wipe(token, WEB_SESSION_STORE_TOKEN_LENGTH + 1U);
    }
    if (view != NULL) {
        secure_wipe(view, sizeof(*view));
    }
    if (principal == NULL || token == NULL || view == NULL ||
        principal->user_id == 0U || principal->auth_generation == 0U ||
        principal->method != USER_AUTH_METHOD_PASSWORD ||
        (principal->role != USER_ROLE_USER && principal->role != USER_ROLE_ADMIN) ||
        principal->username_length == 0U ||
        principal->username_length > USER_DATABASE_USERNAME_CAPACITY ||
        principal->username[principal->username_length] != '\0') {
        return ESP_ERR_INVALID_ARG;
    }

    taskENTER_CRITICAL(&s_lock);
    bool ready = s_state.ready;
    uint64_t epoch = s_state.epoch;
    taskEXIT_CRITICAL(&s_lock);
    if (!ready) {
        return ESP_ERR_INVALID_STATE;
    }
    web_session_store_prune();

    session_entry_t candidate = {0};
    uint8_t random[2U * WEB_SESSION_STORE_SECRET_BYTES] = {0};
    esp_err_t error = origin_digest(origin, origin_length, candidate.origin_digest);
    if (error == ESP_OK) {
        error = secure_random_fill(random, sizeof(random));
    }
    if (error == ESP_OK) {
        encode_hex(random, token);
        memcpy(candidate.csrf, random + WEB_SESSION_STORE_SECRET_BYTES,
               sizeof(candidate.csrf));
        error = digest(token, WEB_SESSION_STORE_TOKEN_LENGTH, candidate.token_digest);
    }
    bool current = false;
    if (error == ESP_OK) {
        candidate.principal = *principal;
        error = user_database_principal_is_current(&candidate.principal, &current);
        if (error == ESP_OK && !current) {
            error = ESP_ERR_NOT_FOUND;
        }
    }
    if (error == ESP_OK) {
        int64_t now = esp_timer_get_time();
        taskENTER_CRITICAL(&s_lock);
        expire_locked(now);
        session_entry_t *available = NULL;
        bool duplicate = false;
        for (size_t i = 0U; i < WEB_SESSION_STORE_CAPACITY; ++i) {
            session_entry_t *entry = &s_state.entries[i];
            if (entry->id == 0U) {
                if (available == NULL) {
                    available = entry;
                }
            } else if (equal_bytes(entry->token_digest, candidate.token_digest,
                                   sizeof(entry->token_digest))) {
                duplicate = true;
            }
        }
        if (!s_state.ready || epoch != s_state.epoch ||
            s_state.next_id == UINT64_MAX || now < 0 ||
            now > INT64_MAX - WEB_SESSION_STORE_LIFETIME_US) {
            error = ESP_ERR_INVALID_STATE;
        } else if (duplicate) {
            error = ESP_FAIL;
        } else if (available == NULL) {
            ++s_state.capacity_rejections;
            error = ESP_ERR_NO_MEM;
        } else {
            candidate.id = ++s_state.next_id;
            candidate.expires_at_us = now + WEB_SESSION_STORE_LIFETIME_US;
            *available = candidate;
            ++s_state.issued;
        }
        taskEXIT_CRITICAL(&s_lock);
    }
    if (error == ESP_OK) {
        export_view(&candidate, view);
    } else {
        secure_wipe(token, WEB_SESSION_STORE_TOKEN_LENGTH + 1U);
    }
    secure_wipe(random, sizeof(random));
    secure_wipe(&candidate, sizeof(candidate));
    return error;
}

esp_err_t web_session_store_lookup(
    const char *token, size_t token_length, const char *origin,
    size_t origin_length, web_session_view_t *view)
{
    if (view == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    secure_wipe(view, sizeof(*view));
    if (token == NULL || token_length != WEB_SESSION_STORE_TOKEN_LENGTH) {
        return ESP_ERR_INVALID_ARG;
    }
    for (size_t i = 0U; i < token_length; ++i) {
        if (!((token[i] >= '0' && token[i] <= '9') ||
              (token[i] >= 'a' && token[i] <= 'f'))) {
            return ESP_ERR_INVALID_ARG;
        }
    }
    uint8_t token_hash[WEB_SESSION_STORE_SECRET_BYTES] = {0};
    uint8_t origin_hash[WEB_SESSION_STORE_SECRET_BYTES] = {0};
    esp_err_t error = origin_digest(origin, origin_length, origin_hash);
    if (error == ESP_OK) {
        error = digest(token, token_length, token_hash);
    }
    if (error == ESP_OK) {
        web_session_id_t id = 0U;
        taskENTER_CRITICAL(&s_lock);
        for (size_t i = 0U; i < WEB_SESSION_STORE_CAPACITY; ++i) {
            const session_entry_t *entry = &s_state.entries[i];
            bool matches = equal_bytes(entry->token_digest, token_hash, sizeof(token_hash));
            matches &= equal_bytes(entry->origin_digest, origin_hash, sizeof(origin_hash));
            if (entry->id != 0U && matches) {
                id = entry->id;
            }
        }
        taskEXIT_CRITICAL(&s_lock);
        error = resolve(id, view, NULL);
    }
    secure_wipe(token_hash, sizeof(token_hash));
    secure_wipe(origin_hash, sizeof(origin_hash));
    return error;
}

esp_err_t web_session_store_is_current(web_session_id_t id, bool *current)
{
    if (current == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    *current = false;
    esp_err_t error = resolve(id, NULL, NULL);
    if (error == ESP_OK) {
        *current = true;
    }
    return error;
}

esp_err_t web_session_store_check_principal(web_session_id_t id,
                                           const user_principal_t *principal,
                                           bool *current)
{
    if (current == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    *current = false;
    if (principal == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    esp_err_t error = resolve(id, NULL, principal);
    *current = error == ESP_OK;
    return error;
}

void web_session_store_invalidate_username(const uint8_t *username, size_t length)
{
    if ((username == NULL && length != 0U) ||
        (username != NULL && !user_database_username_valid(username, length))) {
        return;
    }
    taskENTER_CRITICAL(&s_lock);
    advance_epoch_locked();
    for (size_t i = 0U; i < WEB_SESSION_STORE_CAPACITY; ++i) {
        session_entry_t *entry = &s_state.entries[i];
        if (username == NULL ||
            (entry->principal.username_length == length &&
             memcmp(entry->principal.username, username, length) == 0)) {
            retire_locked(entry, false);
        }
    }
    taskEXIT_CRITICAL(&s_lock);
}

void web_session_store_invalidate(web_session_id_t id)
{
    if (id == 0U) {
        return;
    }
    taskENTER_CRITICAL(&s_lock);
    advance_epoch_locked();
    session_entry_t *entry = find_locked(id);
    if (entry != NULL) {
        retire_locked(entry, false);
    }
    taskEXIT_CRITICAL(&s_lock);
}

void web_session_store_invalidate_user(uint32_t user_id)
{
    if (user_id == 0U) {
        return;
    }
    taskENTER_CRITICAL(&s_lock);
    advance_epoch_locked();
    for (size_t i = 0U; i < WEB_SESSION_STORE_CAPACITY; ++i) {
        if (s_state.entries[i].principal.user_id == user_id) {
            retire_locked(&s_state.entries[i], false);
        }
    }
    taskEXIT_CRITICAL(&s_lock);
}

esp_err_t web_session_store_get_snapshot(web_session_store_snapshot_t *snapshot)
{
    if (snapshot == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    memset(snapshot, 0, sizeof(*snapshot));
    int64_t now = esp_timer_get_time();
    taskENTER_CRITICAL(&s_lock);
    expire_locked(now);
    snapshot->initialized = s_state.ready;
    for (size_t i = 0U; i < WEB_SESSION_STORE_CAPACITY; ++i) {
        snapshot->active += s_state.entries[i].id != 0U;
    }
    snapshot->issued = s_state.issued;
    snapshot->capacity_rejections = s_state.capacity_rejections;
    snapshot->expired = s_state.expired;
    snapshot->invalidated = s_state.invalidated;
    snapshot->lookup_rejections = s_state.lookup_rejections;
    snapshot->init_failures = s_state.init_failures;
    snapshot->storage_bytes = sizeof(s_state) + sizeof(s_lock);
    snapshot->slot_bytes = sizeof(session_entry_t);
    taskEXIT_CRITICAL(&s_lock);
    return ESP_OK;
}