/*
 * libwebsockets - small server side websockets and web server implementation
 *
 * Copyright (C) 2010 - 2026 Andy Green <andy@warmcat.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to
 * deal in the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
 * sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 */

#include "private-lib-core.h"

#define LWS_AUTH_MAX_COOKIE_LEN 4096

struct lws_jwt_auth {
	struct lws_context *cx;
	struct lws *wsi;
	struct lws_jwk *jwk;
	lws_sorted_usec_list_t sul;
	lws_jwt_auth_cb_t cb;
	void *user;

	struct lws_dll2_owner grants;
	uint64_t iat;
	uint64_t exp;
	char cookie_name[64];
	char sub[128];
	char did[128];
	uint32_t uid;
	uint32_t session_epoch;
};

struct lws_jwt_auth_grant {
	lws_dll2_t list;
	char service_name[64];
	int grant_level;
};

static void
lws_jwt_auth_sul_cb(lws_sorted_usec_list_t *sul)
{
	struct lws_jwt_auth *ja = lws_container_of(sul, struct lws_jwt_auth, sul);
	uint64_t now = (uint64_t)lws_now_secs();

	if (now >= ja->exp) {
		if (ja->cb)
			ja->cb(ja, LWS_JWT_AUTH_STATE_EXPIRED, ja->user);
	} else {
		if (ja->cb)
			ja->cb(ja, LWS_JWT_AUTH_STATE_REAUTH, ja->user);

		/* Reschedule for the actual expiration */
		lws_usec_t us = (lws_usec_t)(ja->exp - now) * LWS_US_PER_SEC;
		lws_sul_schedule(ja->cx, 0, &ja->sul, lws_jwt_auth_sul_cb, us);
	}
}

static void
lws_jwt_auth_schedule(struct lws_jwt_auth *ja)
{
	uint64_t now = (uint64_t)lws_now_secs();
	lws_usec_t us;

	lws_sul_cancel(&ja->sul);

	if (now >= ja->exp) {
		lws_sul_schedule(ja->cx, 0, &ja->sul, lws_jwt_auth_sul_cb, 1);
		return;
	}

	uint64_t total_val = 0;
	if (ja->iat && ja->exp > ja->iat)
		total_val = ja->exp - ja->iat;
	else
		total_val = ja->exp - now;

	uint64_t reauth_point = ja->exp - (total_val * 15) / 100; /* 85% */

	if (now >= reauth_point) {
		us = (lws_usec_t)(ja->exp - now) * LWS_US_PER_SEC;
		lws_sul_schedule(ja->cx, 0, &ja->sul, lws_jwt_auth_sul_cb, us);
	} else {
		us = (lws_usec_t)(reauth_point - now) * LWS_US_PER_SEC;
		lws_sul_schedule(ja->cx, 0, &ja->sul, lws_jwt_auth_sul_cb, us);
	}
}

struct jwt_auth_parse_ctx {
	struct lws_jwt_auth *ja;
	int parsing_grants;
};

static const char * const auth_paths[] = {
	"exp",
	"iat",
	"grants",
	"grants.*",
	"sub",
	"email",
	"uid",
	"did",
	"sec",
};

enum {
	JAP_EXP,
	JAP_IAT,
	JAP_GRANTS,
	JAP_GRANTS_ANY,
	JAP_SUB,
	JAP_EMAIL,
	JAP_UID,
	JAP_DID,
	JAP_SEC,
};

static signed char
jwt_auth_lejp_cb(struct lejp_ctx *ctx, char reason)
{
	struct jwt_auth_parse_ctx *pctx = (struct jwt_auth_parse_ctx *)ctx->user;

	if (reason == LEJPCB_OBJECT_START && ctx->path_match == JAP_GRANTS + 1) {
		pctx->parsing_grants = 1;
		return 0;
	}
	if (reason == LEJPCB_OBJECT_END && pctx->parsing_grants) {
		pctx->parsing_grants = 0;
		return 0;
	}

	if (reason == LEJPCB_VAL_NUM_INT) {
		if (ctx->path_match == JAP_EXP + 1) {
			pctx->ja->exp = (uint64_t)atoll(ctx->buf);
		} else if (ctx->path_match == JAP_IAT + 1) {
			pctx->ja->iat = (uint64_t)atoll(ctx->buf);
		} else if (ctx->path_match == JAP_UID + 1) {
			pctx->ja->uid = (uint32_t)atoi(ctx->buf);
		} else if (ctx->path_match == JAP_SEC + 1) {
			pctx->ja->session_epoch = (uint32_t)atoi(ctx->buf);
		} else if (ctx->path_match == JAP_GRANTS_ANY + 1 && pctx->parsing_grants) {
			struct lws_jwt_auth_grant *g = malloc(sizeof(*g));
			if (g) {
				memset(g, 0, sizeof(*g));
				lws_strncpy(g->service_name, ctx->path + 7, sizeof(g->service_name));
				g->grant_level = atoi(ctx->buf);
				lws_dll2_add_tail(&g->list, &pctx->ja->grants);
			}
		}
	} else if (reason == LEJPCB_VAL_STR_CHUNK || reason == LEJPCB_VAL_STR_END) {
		if (ctx->path_match == JAP_SUB + 1 || ctx->path_match == JAP_EMAIL + 1) {
			lws_strncpy(pctx->ja->sub, ctx->buf, sizeof(pctx->ja->sub));
		} else if (ctx->path_match == JAP_DID + 1) {
			lws_strncpy(pctx->ja->did, ctx->buf, sizeof(pctx->ja->did));
		}
	}

	return 0;
}

int
lws_jwt_auth_update(struct lws_jwt_auth *ja, const char *jwt, const char **reason)
{
	char temp[2048], out[2048];
	size_t out_len = sizeof(out);
	struct lejp_ctx ctx;
	struct jwt_auth_parse_ctx pctx;
	int m;

	lws_start_foreach_dll_safe(struct lws_dll2 *, d, d1, ja->grants.head) {
		struct lws_jwt_auth_grant *g = lws_container_of(d, struct lws_jwt_auth_grant, list);
		lws_dll2_remove(&g->list);
		free(g);
	} lws_end_foreach_dll_safe(d, d1);

	if (lws_jwt_signed_validate(ja->cx, ja->jwk, "ES256,ES384,ES512,RS256,RS384,RS512,HS256",
				    jwt, strlen(jwt), temp, sizeof(temp), out, &out_len)) {
		lwsl_err("%s: Verification failed\n", __func__);
		if (reason)
			*reason = "JWT signature verification failed";
		return -1;
	}

	pctx.ja = ja;
	pctx.parsing_grants = 0;
	lejp_construct(&ctx, jwt_auth_lejp_cb, &pctx, auth_paths, LWS_ARRAY_SIZE(auth_paths));
	m = (int)(lejp_parse(&ctx, (uint8_t *)out, (int)out_len));
	lejp_destruct(&ctx);

	if (m < 0 && m != LEJP_REJECT_UNKNOWN) {
		lwsl_err("%s: JSON decode failed\n", __func__);
		if (reason)
			*reason = "Failed to parse JWT payload JSON";
		return -1;
	}

	lws_jwt_auth_schedule(ja);

	return 0;
}

struct lws_jwt_auth *
lws_jwt_auth_create(struct lws *wsi, struct lws_jwk *jwk,
		    const char *cookie_name,
		    lws_jwt_auth_cb_t cb, void *user,
		    const char **reason)
{
	char jwt[8192];
	size_t jwt_len = sizeof(jwt);
	struct lws_jwt_auth *ja;

	if (lws_http_cookie_get(wsi, cookie_name, jwt, &jwt_len)) {
		if (reason)
			*reason = "Cookie not found";
		return NULL;
	}

	ja = malloc(sizeof(*ja));
	if (!ja) {
		if (reason)
			*reason = "OOM";
		return NULL;
	}

	memset(ja, 0, sizeof(*ja));
	ja->cx = lws_get_context(wsi);
	ja->wsi = wsi;
	ja->jwk = jwk;
	ja->cb = cb;
	ja->user = user;
	lws_strncpy(ja->cookie_name, cookie_name, sizeof(ja->cookie_name));

	if (lws_jwt_auth_update(ja, jwt, reason)) {
		free(ja);
		return NULL;
	}

	return ja;
}

int
lws_jwt_auth_query_grant(struct lws_jwt_auth *ja, const char *service_name)
{
	int wildcard_level = -1;

	if (!ja)
		return -1;

	lws_start_foreach_dll_safe(struct lws_dll2 *, d, d1, ja->grants.head) {
		struct lws_jwt_auth_grant *g = lws_container_of(d, struct lws_jwt_auth_grant, list);
		if (!strcmp(g->service_name, service_name))
			return g->grant_level;
		if (!strcmp(g->service_name, "*"))
			wildcard_level = g->grant_level;
	} lws_end_foreach_dll_safe(d, d1);

	return wildcard_level;
}

void
lws_jwt_auth_destroy(struct lws_jwt_auth **ja)
{
	if (!ja || !*ja)
		return;

	lws_sul_cancel(&((*ja)->sul));

	lws_start_foreach_dll_safe(struct lws_dll2 *, d, d1, (*ja)->grants.head) {
		struct lws_jwt_auth_grant *g = lws_container_of(d, struct lws_jwt_auth_grant, list);
		lws_dll2_remove(&g->list);
		free(g);
	} lws_end_foreach_dll_safe(d, d1);

	free(*ja);
	*ja = NULL;
}

const char *
lws_jwt_auth_get_sub(struct lws_jwt_auth *ja)
{
	if (!ja || !ja->sub[0])
		return NULL;
	return ja->sub;
}

const char *
lws_jwt_auth_get_did(struct lws_jwt_auth *ja)
{
	if (!ja || !ja->did[0])
		return NULL;
	return ja->did;
}

uint32_t
lws_jwt_auth_get_uid(struct lws_jwt_auth *ja)
{
	if (!ja)
		return 0;
	return ja->uid;
}

uint64_t
lws_jwt_auth_get_exp(struct lws_jwt_auth *ja)
{
	if (!ja)
		return 0;
	return ja->exp;
}

uint32_t
lws_jwt_auth_count_grants(struct lws_jwt_auth *ja)
{
	if (!ja)
		return 0;
	return ja->grants.count;
}

uint32_t
lws_jwt_auth_get_sec(struct lws_jwt_auth *ja)
{
	if (!ja)
		return 0;
	return ja->session_epoch;
}
