/* Copyright (C) 2015-2025 maClara, LLC <info@maclara-llc.com>
   This file is part of the JWT C Library

   SPDX-License-Identifier:  MPL-2.0
   This Source Code Form is subject to the terms of the Mozilla Public
   License, v. 2.0. If a copy of the MPL was not distributed with this
   file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#include <stdlib.h>
#include <string.h>
#include <stdio.h>

#include <jwt.h>

#ifdef USE_KCAPI_MD
#include <kcapi.h>
#define KCAPI_MAX_MD_SIZE 64
#endif

#include "jwt-private.h"

const char *jwt_alg_str(jwt_alg_t alg)
{
	switch (alg) {
	case JWT_ALG_NONE:
		return "none";
	case JWT_ALG_HS256:
		return "HS256";
	case JWT_ALG_HS384:
		return "HS384";
	case JWT_ALG_HS512:
		return "HS512";
	case JWT_ALG_RS256:
		return "RS256";
	case JWT_ALG_RS384:
		return "RS384";
	case JWT_ALG_RS512:
		return "RS512";
	case JWT_ALG_ES256:
		return "ES256";
	case JWT_ALG_ES256K:
		return "ES256K";
	case JWT_ALG_ES384:
		return "ES384";
	case JWT_ALG_ES512:
		return "ES512";
	case JWT_ALG_PS256:
		return "PS256";
	case JWT_ALG_PS384:
		return "PS384";
	case JWT_ALG_PS512:
		return "PS512";
	case JWT_ALG_EDDSA:
		return "EdDSA";
	default:
		return NULL;
	}
}

jwt_alg_t jwt_str_alg(const char *alg)
{
	if (alg == NULL)
		return JWT_ALG_INVAL;

	if (!strcmp(alg, "none"))
		return JWT_ALG_NONE;
	else if (!strcmp(alg, "HS256"))
		return JWT_ALG_HS256;
	else if (!strcmp(alg, "HS384"))
		return JWT_ALG_HS384;
	else if (!strcmp(alg, "HS512"))
		return JWT_ALG_HS512;
	else if (!strcmp(alg, "RS256"))
		return JWT_ALG_RS256;
	else if (!strcmp(alg, "RS384"))
		return JWT_ALG_RS384;
	else if (!strcmp(alg, "RS512"))
		return JWT_ALG_RS512;
	else if (!strcmp(alg, "ES256"))
		return JWT_ALG_ES256;
	else if (!strcmp(alg, "ES256K"))
		return JWT_ALG_ES256K;
	else if (!strcmp(alg, "ES384"))
		return JWT_ALG_ES384;
	else if (!strcmp(alg, "ES512"))
		return JWT_ALG_ES512;
	else if (!strcmp(alg, "PS256"))
		return JWT_ALG_PS256;
	else if (!strcmp(alg, "PS384"))
		return JWT_ALG_PS384;
	else if (!strcmp(alg, "PS512"))
		return JWT_ALG_PS512;
	else if (!strcmp(alg, "EdDSA"))
		return JWT_ALG_EDDSA;

	return JWT_ALG_INVAL;
}

JWT_NO_EXPORT
jwt_t *jwt_new(void)
{
	jwt_t *jwt = jwt_malloc(sizeof(*jwt));

	if (!jwt)
		return NULL; // LCOV_EXCL_LINE

	memset(jwt, 0, sizeof(*jwt));

	jwt->claims = jwt_json_create();
	jwt->headers = jwt_json_create();

	if (!jwt->claims || !jwt->headers)
		jwt_freep(&jwt); // LCOV_EXCL_LINE

	return jwt;
}

jwt_alg_t jwt_get_alg(const jwt_t *jwt)
{
	if (jwt == NULL)
		return JWT_ALG_INVAL;

	return jwt->alg;
}

void jwt_free(jwt_t *jwt)
{
	if (!jwt)
		return;

	jwt_json_release(jwt->claims);
	jwt_json_release(jwt->headers);

	memset(jwt, 0, sizeof(*jwt));

	jwt_freemem(jwt);
}

static int __check_hmac(jwt_t *jwt)
{
	int key_bits = jwt->key->bits;

	switch (jwt->alg) {
	case JWT_ALG_HS256:
		if (key_bits >= 256)
			return 0;
		jwt_write_error(jwt, "Key too short for HS256: %d bits",
				key_bits);
		break;

	case JWT_ALG_HS384:
		if (key_bits >= 384)
			return 0;
		jwt_write_error(jwt, "Key too short for HS384: %d bits",
				key_bits);
		break;

	case JWT_ALG_HS512:
		if (key_bits >= 512)
			return 0;
		jwt_write_error(jwt, "Key too short for HS512: %d bits",
				key_bits);
		break;
	// LCOV_EXCL_START
	default:
		return 1;
	// LCOV_EXCL_STOP
	}

	return 1;
}

static int __check_key_bits(jwt_t *jwt)
{
	int key_bits = jwt->key->bits;

	switch (jwt->alg) {
	case JWT_ALG_RS256:
	case JWT_ALG_RS384:
	case JWT_ALG_RS512:

	case JWT_ALG_PS256:
	case JWT_ALG_PS384:
	case JWT_ALG_PS512:
		if (key_bits >= 2048)
			return 0;
		jwt_write_error(jwt, "Key too short for RSA algs: %d bits",
				key_bits);
		break;

	case JWT_ALG_EDDSA:
		if (key_bits == 256 || key_bits == 456)
			return 0;
		jwt_write_error(jwt, "Key needs to be 256 or 456 bits: %d bits",
				key_bits);
		break;

	case JWT_ALG_ES256K:
	case JWT_ALG_ES256:
		if (key_bits == 256)
			return 0;
		jwt_write_error(jwt, "Key needs to be 256 bits: %d bits",
				key_bits);
		break;

	case JWT_ALG_ES384:
		if (key_bits == 384)
			return 0;
		jwt_write_error(jwt, "Key needs to be 384 bits: %d bits",
				key_bits);
		break;

	case JWT_ALG_ES512:
		if (key_bits == 521)
			return 0;
		jwt_write_error(jwt, "Key needs to be 521 bits: %d bits",
				key_bits);
		break;
	// LCOV_EXCL_START
	default:
		break;
	// LCOV_EXCL_STOP
	}

	return 1; // LCOV_EXCL_LINE
}

static int sign_sha_hmac(jwt_t *jwt, char **out, unsigned int *len,
			 const char *str, unsigned int str_len)
{
#ifdef USE_KCAPI_MD
	static int skip_kcapi = 0;
	size_t key_len;
	void *key;
	int ret;

	/* If kcapi fails once, we don't try again. */
	if (skip_kcapi)
		goto fallback_ops; // LCOV_EXCL_LINE

	key = jwt->key->oct.key;
	key_len = jwt->key->oct.len;

	*out = jwt_malloc(KCAPI_MAX_MD_SIZE);
	if (*out == NULL)
		return 1; // LCOV_EXCL_LINE

	switch (jwt->alg) {
	/* HMAC */
	case JWT_ALG_HS256:
		ret = kcapi_md_hmac_sha256(key, key_len, (uint8_t *)str,
					   str_len, (uint8_t *)*out,
					   KCAPI_MAX_MD_SIZE);
		break;
	case JWT_ALG_HS384:
		ret = kcapi_md_hmac_sha384(key, key_len, (uint8_t *)str,
					   str_len, (uint8_t *)*out,
					   KCAPI_MAX_MD_SIZE);
		break;
	case JWT_ALG_HS512:
		ret = kcapi_md_hmac_sha512(key, key_len, (uint8_t *)str,
					   str_len, (uint8_t *)*out,
					   KCAPI_MAX_MD_SIZE);
		break;
	// LCOV_EXCL_START
	default:
		/* This isn't a failure in kcapi, so just error out */
		jwt_freemem(out);
		return 1;
	// LCOV_EXCL_STOP
	}

	if (ret > 0) {
		*len = ret;
		return 0;
	}

	/* Fallthrough to normal ops */
	// LCOV_EXCL_START
	jwt_freemem(*out);
	skip_kcapi = 1;
	// LCOV_EXCL_STOP

fallback_ops:
#endif
	return jwt_ops->sign_sha_hmac(jwt, out, len, str, str_len);
}

int jwt_sign(jwt_t *jwt, char **out, unsigned int *len, const char *str,
	     unsigned int str_len)
{
	switch (jwt->alg) {
	/* HMAC */
	case JWT_ALG_HS256:
	case JWT_ALG_HS384:
	case JWT_ALG_HS512:
		if (__check_hmac(jwt))
			return 1;
		if (sign_sha_hmac(jwt, out, len, str, str_len)) {
			/* There's not really a way to induce failure here,
			 * and there's not really much of a chance this can fail
			 * other than an internal fatal error in the crypto
			 * library. */
			// LCOV_EXCL_START
			jwt_write_error(jwt, "Token failed signing");
			return 1;
			// LCOV_EXCL_STOP
		} else {
			return 0;
		}

	/* RSA */
	case JWT_ALG_RS256:
	case JWT_ALG_RS384:
	case JWT_ALG_RS512:

	/* RSA-PSS */
	case JWT_ALG_PS256:
	case JWT_ALG_PS384:
	case JWT_ALG_PS512:

	/* ECC */
	case JWT_ALG_ES256:
	case JWT_ALG_ES256K:
	case JWT_ALG_ES384:
	case JWT_ALG_ES512:

	/* EdDSA */
	case JWT_ALG_EDDSA:
		if (__check_key_bits(jwt))
			return 1;
		if (jwt_ops->sign_sha_pem(jwt, out, len, str, str_len)) {
			jwt_write_error(jwt, "Token failed signing");
			return 1;
		} else {
			return 0;
		}

	/* You wut, mate? */
	// LCOV_EXCL_START
	default:
		jwt_write_error(jwt, "Unknown algorigthm");
		return 1;
	// LCOV_EXCL_STOP
	}
}

/* A time-safe strcmp function */
static int _crypto_strcmp(const char *str1, const char *str2)
{
	/* Get the LONGEST length */
	int len1 = strlen(str1);
	int len2 = strlen(str2);
	int len_max = len1 >= len2 ? len1 : len2;

	int i;
	volatile int ret = 0;

	/* Iterate the entire longest string no matter what. Only testing
	 * the shortest string would still allow attacks for
	 * "a" == "aKJSDHkjashaaHJASJ", adding a character each time one
	 * is found. */
	for (i = 0; i < len_max; i++) {
		char c1, c2;

		c1 = (i < len1) ? str1[i] : 0;
		c2 = (i < len2) ? str2[i] : 0;

		ret |= c1 ^ c2;
	}

	/* Don't forget to check length */
	ret |= len1 ^ len2;

	return ret;
}

static int _verify_sha_hmac(jwt_t *jwt, const char *head,
			    unsigned int head_len, const char *sig)
{
	char_auto *res = NULL;
	char_auto *buf = NULL;
	unsigned int res_len;
	int ret;

	ret = jwt_sign(jwt, &res, &res_len, head, head_len);
	if (ret)
		return 1; // LCOV_EXCL_LINE

	ret = jwt_base64uri_encode(&buf, res, res_len);
	if (ret <= 0)
		return 1; // LCOV_EXCL_LINE

	return _crypto_strcmp(buf, sig) ? 1 : 0;
}

jwt_t *jwt_verify_sig(jwt_t *jwt, const char *head, unsigned int head_len,
		      const char *sig_b64)
{
	int sig_len;
	char_auto *sig = NULL;

	switch (jwt->alg) {
	/* HMAC */
	case JWT_ALG_HS256:
	case JWT_ALG_HS384:
	case JWT_ALG_HS512:
		if (_verify_sha_hmac(jwt, head, head_len, sig_b64))
			jwt_write_error(jwt, "Token failed verification");
		break;

	/* RSA */
	case JWT_ALG_RS256:
	case JWT_ALG_RS384:
	case JWT_ALG_RS512:

	/* RSA-PSS */
	case JWT_ALG_PS256:
	case JWT_ALG_PS384:
	case JWT_ALG_PS512:

	/* ECC */
	case JWT_ALG_ES256:
	case JWT_ALG_ES256K:
	case JWT_ALG_ES384:
	case JWT_ALG_ES512:

	/* EdDSA */
	case JWT_ALG_EDDSA:
		if (__check_key_bits(jwt))
			break;

		sig = jwt_base64uri_decode(sig_b64, &sig_len);
		if (sig == NULL) {
			jwt_write_error(jwt, "Error decoding signature");
			break;
		}

		if (jwt_ops->verify_sha_pem(jwt, head, head_len,
					    (unsigned char *)sig, sig_len))
			jwt_write_error(jwt, "Token failed verification");
		break;

	/* You wut, mate? */
	// LCOV_EXCL_START
	default:
		jwt_write_error(jwt, "Unknown algorigthm");
	} // LCOV_EXCL_STOP

	return jwt;
}

/* This is a public domain base64 implementation written by WEI Zhicheng.
   https://github.com/zhicheng/base64 */
#define BASE64_ENCODE_OUT_SIZE(s) ((unsigned int)((((s) + 2) / 3) * 4 + 1))
#define BASE64_DECODE_OUT_SIZE(s) ((unsigned int)(((s) / 4) * 3))
#define BASE64_PAD '='
#define BASE64DE_FIRST '+'
#define BASE64DE_LAST 'z'

/* BASE 64 encode table */
static const char base64en[] = {
	'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H',
	'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P',
	'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X',
	'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f',
	'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n',
	'o', 'p', 'q', 'r', 's', 't', 'u', 'v',
	'w', 'x', 'y', 'z', '0', '1', '2', '3',
	'4', '5', '6', '7', '8', '9', '+', '/',
};

/* ASCII order for BASE 64 decode, 255 in unused character */
static const unsigned char base64de[] = {
	/* nul, soh, stx, etx, eot, enq, ack, bel, */
	   255, 255, 255, 255, 255, 255, 255, 255,

	/*  bs,  ht,  nl,  vt,  np,  cr,  so,  si, */
	   255, 255, 255, 255, 255, 255, 255, 255,

	/* dle, dc1, dc2, dc3, dc4, nak, syn, etb, */
	   255, 255, 255, 255, 255, 255, 255, 255,

	/* can,  em, sub, esc,  fs,  gs,  rs,  us, */
	   255, 255, 255, 255, 255, 255, 255, 255,

	/*  sp, '!', '"', '#', '$', '%', '&', ''', */
	   255, 255, 255, 255, 255, 255, 255, 255,

	/* '(', ')', '*', '+', ',', '-', '.', '/', */
	   255, 255, 255,  62, 255, 255, 255,  63,

	/* '0', '1', '2', '3', '4', '5', '6', '7', */
	    52,  53,  54,  55,  56,  57,  58,  59,

	/* '8', '9', ':', ';', '<', '=', '>', '?', */
	    60,  61, 255, 255, 255, 255, 255, 255,

	/* '@', 'A', 'B', 'C', 'D', 'E', 'F', 'G', */
	   255,   0,   1,  2,   3,   4,   5,    6,

	/* 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', */
	     7,   8,   9,  10,  11,  12,  13,  14,

	/* 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', */
	    15,  16,  17,  18,  19,  20,  21,  22,

	/* 'X', 'Y', 'Z', '[', '\', ']', '^', '_', */
	    23,  24,  25, 255, 255, 255, 255, 255,

	/* '`', 'a', 'b', 'c', 'd', 'e', 'f', 'g', */
	   255,  26,  27,  28,  29,  30,  31,  32,

	/* 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', */
	    33,  34,  35,  36,  37,  38,  39,  40,

	/* 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', */
	    41,  42,  43,  44,  45,  46,  47,  48,

	/* 'x', 'y', 'z', '{', '|', '}', '~', del, */
	    49,  50,  51, 255, 255, 255, 255, 255
};

static unsigned int
base64_encode(const unsigned char *in, unsigned int inlen, char *out)
{
	int s;
	unsigned int i;
	unsigned int j;
	unsigned char c;
	unsigned char l;

	s = 0;
	l = 0;
	for (i = j = 0; i < inlen; i++) {
		c = in[i];

		switch (s) {
		case 0:
			s = 1;
			out[j++] = base64en[(c >> 2) & 0x3F];
			break;
		case 1:
			s = 2;
			out[j++] = base64en[((l & 0x3) << 4) | ((c >> 4) & 0xF)];
			break;
		case 2:
			s = 0;
			out[j++] = base64en[((l & 0xF) << 2) | ((c >> 6) & 0x3)];
			out[j++] = base64en[c & 0x3F];
			break;
		}
		l = c;
	}

	switch (s) {
	case 1:
		out[j++] = base64en[(l & 0x3) << 4];
		out[j++] = BASE64_PAD;
		out[j++] = BASE64_PAD;
		break;
	case 2:
		out[j++] = base64en[(l & 0xF) << 2];
		out[j++] = BASE64_PAD;
		break;
	}

	out[j] = 0;

	return j;
}

static unsigned int
base64_decode(const char *in, unsigned int inlen, unsigned char *out)
{
	unsigned int i;
	unsigned int j;
	unsigned char c;

	if (inlen & 0x3) {
		/* We sanitize the input to where this never gets hit */
		return 0; // LCOV_EXCL_LINE
	}

	for (i = j = 0; i < inlen; i++) {
		if (in[i] == BASE64_PAD) {
			break;
		}
		if (in[i] < BASE64DE_FIRST || in[i] > BASE64DE_LAST) {
			return 0;
		}

		c = base64de[(unsigned char)in[i]];
		if (c == 255) {
			return 0;
		}

		switch (i & 0x3) {
		case 0:
			out[j] = (c << 2) & 0xFF;
			break;
		case 1:
			out[j++] |= (c >> 4) & 0x3;
			out[j] = (c & 0xF) << 4; 
			break;
		case 2:
			out[j++] |= (c >> 2) & 0xF;
			out[j] = (c & 0x3) << 6;
			break;
		case 3:
			out[j++] |= c;
			break;
		}
	}

	return j;
}

void *jwt_base64uri_decode(const char *src, int *ret_len)
{
	void *buf;
	char *new;
	int len, i, z;

	if (src == NULL || ret_len == NULL)
		return NULL; // LCOV_EXCL_LINE
			     // Should really be an abort

	/* Decode based on RFC-4648 URI safe encoding. */
	len = (int)strlen(src);

	/* Validate length */
	z = (len % 4);
	switch (z) {
	case 0:
		/* No added padding */
		break;

	case 2:
		/* 2 added = for padding */
		break;

	case 3:
		/* 1 added = for padding */
		z = 1;
		break;

	default:
		/* Something bad */
		return NULL;
	}

	new = jwt_malloc(len + z + 1);
	if (!new)
		return NULL; // LCOV_EXCL_LINE

	for (i = 0; i < len; i++) {
		switch (src[i]) {
		case '-':
			new[i] = '+';
			break;
		case '_':
			new[i] = '/';
			break;
		default:
			new[i] = src[i];
		}
	}

	for (; z > 0; z--)
		new[i++] = '=';

	new[i] = '\0';
	len = i;

	/* Now we have a standard base64 encoded string. */
	buf = jwt_malloc(BASE64_DECODE_OUT_SIZE(len) + 1);
	if (buf == NULL) {
		// LCOV_EXCL_START
		jwt_freemem(new);
		return NULL;
		// LCOV_EXCL_STOP
	}

	*ret_len = base64_decode(new, len, buf);
	jwt_freemem(new);

	if (*ret_len <= 0)
		jwt_freemem(buf);

	return buf;
}

int jwt_base64uri_encode(char **_dst, const char *plain, int plain_len)
{
	int len, i;
	char *dst;

	len = BASE64_ENCODE_OUT_SIZE(plain_len);
	dst = jwt_malloc(len + 1);
	if (dst == NULL)
		return -1; // LCOV_EXCL_LINE
	*_dst = dst;

	/* First, a normal base64 encoding */
	len = base64_encode((const unsigned char *)plain, plain_len, dst);

	/* Now for the URI encoding */
	for (i = 0; i < len; i++) {
		switch (dst[i]) {
		case '+':
			dst[i] = '-';
			break;
		case '/':
			dst[i] = '_';
			break;
		case '=':
			dst[i] = '\0';
			break;
		}
	}

	/* Just in case there's no padding. */
	dst[i] = '\0';

	return i;
}
