// Copyright © 2025 blackshirt.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
//
// This module implements building block for elliptic-curve diffie-helman
// key exchange (ECDH) mechanism through curve25519 curve.
module curve25519

import crypto.rand
import crypto.internal.subtle
import crypto.ed25519.internal.edwards25519

// scalar_size is the size of the Curve25519 key
const scalar_size = 32

// point_size is the size of the Curve25519 point
const point_size = 32

// zero_point is point with 32 bytes length of zeros bytes
const zero_point = []u8{len: 32, init: u8(0x00)}

// base_point is the canonical Curve25519 generator, encoded as a byte with value 9,
// followed by 31 zero bytes
const base_point = [u8(9), 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
	0, 0, 0, 0, 0, 0, 0, 0]

// PrivateKey represents Curve25519 private key
@[noinit]
pub struct PrivateKey {
mut:
	// boolean flag that tells this key should not be
	// used again when its true. its set after .free call
	done bool
	// clamped key, 32 bytes length
	key []u8
}

// new creates a new Curve25519 key using randomly generated bytes from `crypto.rand`.
@[direct_array_access]
pub fn PrivateKey.new() !&PrivateKey {
	mut bytes := rand.read(scalar_size)!
	// do bytes clamping
	clamp(mut bytes)!

	if is_zero_point(bytes) || is_base_point(bytes) {
		return error('PrivateKey.new: bytes zeros or base point')
	}

	return &PrivateKey{
		key: bytes
	}
}

// new_from_seed creates a new Curve25519 key from provided seed bytes.
@[direct_array_access]
pub fn PrivateKey.new_from_seed(seed []u8) !&PrivateKey {
	if seed.len != scalar_size {
		return error('invalid scalar size')
	}
	mut bytes := seed.clone()
	clamp(mut bytes)!

	if is_zero_point(bytes) || is_base_point(bytes) {
		return error('PrivateKey.new_from_seed: bytes was zeros or base point')
	}
	return &PrivateKey{
		key: bytes
	}
}

// public_key returns the associated public key part of the PrivateKey.
@[direct_array_access]
pub fn (mut pv PrivateKey) public_key() !&PublicKey {
	if pv.done {
		return error('your key has been marked as freed')
	}
	out := x25519(mut pv.key, base_point)!
	return &PublicKey{
		key: out
	}
}

// equal returns whether two private keys are equal.
@[direct_array_access]
pub fn (pv PrivateKey) equal(oth PrivateKey) bool {
	if pv.done || oth.done {
		panic('the key has been marked as freed')
	}
	if pv.key.len != scalar_size || oth.key.len != scalar_size {
		return false
	}
	return subtle.constant_time_compare(pv.key, oth.key) == 1
}

// x25519 performs scalar multiplication between key and point and return another bytes (point).
@[direct_array_access]
pub fn (mut pv PrivateKey) x25519(point []u8) ![]u8 {
	if pv.done {
		return error('PrivateKey has been marked as freed')
	}
	if point.len != point_size {
		return error('bad point size, should be 32')
	}
	// We reject and disallow zero-bytes point to be passed
	// and check it here as a quick exit before heavy math
	// calculation on `x25519` call
	if is_zero(point) {
		return error('x25519: get zeros point')
	}
	// even technically its possible, but we limit to unallow it
	if subtle.constant_time_compare(pv.key, point) == 1 {
		return error('pv.key identical with point')
	}
	out := x25519(mut pv.key, point)!

	return out
}

// bytes return a clone of the bytes of the underlying PrivateKey
pub fn (pv PrivateKey) bytes() ![]u8 {
	if pv.done {
		return error('PrivateKey has been marked as freed')
	}
	return pv.key.clone()
}

// free releases underlying key. Dont use the key after calling .free
@[unsafe]
pub fn (mut pv PrivateKey) free() {
	// when private key has been marked as done (freed),
	// calling free on already freed key would lead to undefined behavior.
	// so, we check it
	if pv.done {
		return
	}
	unsafe { pv.key.free() }
	// sets flag to finish
	pv.done = true
}

// PublicKey represent Curve25519 key.
@[noinit]
pub struct PublicKey {
mut:
	// 32 bytes length of scalar * point
	key []u8
}

// new_from_bytes creates a new Curve25519 public key from provided bytes.
pub fn PublicKey.new_from_bytes(bytes []u8) !&PublicKey {
	if bytes.len != point_size {
		return error('PublicKey.new: bad bytes length')
	}
	// Refers to the D.J. Bernstein, the designer of the curve25519, public key validation
	// in curve25519 is generally not needed for Diffie-Hellman key exchange.
	// See https://cr.yp.to/ecdh.html#validate
	// But there are availables suggestion to do validation on them spreads on the internet, likes
	// - blacklisting the known bad public keys
	// - check the shared value and to raise exception if it is zero.
	// - You can also bind the exchanged public keys to the shared keys, i.e.,
	//   instead of using H(abG) as the shared keys, you should use H(aG || bG || abG)
	//
	// We only, check for zeros public key
	if is_zero(bytes) {
		return error('PublicKey.new: get zeros bytes')
	}

	// otherwise, we can return it
	return &PublicKey{
		key: bytes
	}
}

// equal tells whether two public keys are equal
pub fn (pb PublicKey) equal(other PublicKey) bool {
	// different length, should not happen
	if pb.key.len != point_size || other.key.len != point_size {
		return false
	}
	return subtle.constant_time_compare(pb.key, other.key) == 1
}

// bytes return the clone of the bytes of the underlying PublicKey
pub fn (pb PublicKey) bytes() ![]u8 {
	if pb.key.len != point_size {
		return error('bad public key size')
	}
	return pb.key.clone()
}

// SharedOpts is the configuration options to `derive_shared_secret` routine
@[params]
pub struct SharedOpts {
pub mut:
	should_derive bool
	derivator     Derivator = RawDerivator{}
	drv_opts      DeriveOpts
}

// DeriveOpts is config to drive the Derivator's derive operation.
@[params]
pub struct DeriveOpts {}

// Derivator represent key derivation function
pub interface Derivator {
	// derive transforms bytes in sec into another form of bytes.
	derive(sec []u8, opt DeriveOpts) ![]u8
}

// RawDerivator was a simple derivator with no derivation behaviour.
struct RawDerivator {}

fn (rd RawDerivator) derive(sec []u8, opt DeriveOpts) ![]u8 {
	return sec
}

// derive_shared_secret derives a shared secret between two peer's
// between first private key's peer and the second PublicKey's peer.
// Its accepts SharedOpts options to advance supports for other key derivation mechanism.
//
// 6.  Diffie-Hellman with Curve25519
// See https://datatracker.ietf.org/doc/html/rfc7748#section-6
pub fn derive_shared_secret(mut local PrivateKey, remote PublicKey, opt SharedOpts) ![]u8 {
	// TODO: should this check be relaxed ?
	// check for safety
	local_pubkey := local.public_key()!
	if local_pubkey.equal(remote) {
		return error('unallowed equal public key between peer')
	}
	// The local peer generates local private key, local_privkey, generates local public key, local_pubkey.
	// and remote peer generates remote private key, ie, remote_privkey,
	// with generated remote public key, remote_pubkey.
	// Both now share shared = X25519(local_privkey, remote_pubkey) = X25519(remote_privkey, local_pubkey)
	// as a shared secret, which is then used as a key or input to a key derivation function.
	sec := local.x25519(remote.key)!
	// Internally, x25519 has builtin check for zeros result
	// but only for non base point branch on x25519_generic routine
	if is_zero(sec) {
		return error('zeroes shared secret')
	}
	// you can choose this sec as an input into other key derivator, and pass this sec
	// into provided derivator
	if opt.should_derive {
		// While the shared secret can be used directly, it's often recommended to apply
		// a key derivation function (KDF), like HKDF to derive a more robust key for cryptographic operations.
		new_sec := opt.derivator.derive(sec, opt.drv_opts)!
		if is_zero(new_sec) {
			return error('zeroes shared secret after derivation')
		}
		return new_sec
	}
	// otherwise, just return the sec as is.
	return sec
}

// x25519 returns the result of the scalar multiplication (`scalar` * `point`),
// according to RFC 7748, Section 5. scalar, point and the return value are slices of 32 bytes.
// The functions take a scalar and a `u-coordinate` as inputs and produce a `u-coordinate` as output.
// Although the functions work internally with integers, the inputs and
// outputs are 32-bytes length (for X25519).
// scalar can be generated at random, for example with `crypto.rand` and point should
// be either `base_point` or the output of another `x25519` call.
@[direct_array_access]
pub fn x25519(mut scalar []u8, point []u8) ![]u8 {
	// likes the previous comment, we add zeroes point check here
	// and reject if it happen.
	if is_zero(point) || is_zero(scalar) {
		return error('x25519: unallowed zeros/scalar point')
	}
	mut dst := []u8{len: 32}
	// we do bytes clamping here, to make sure scalar was ready to use
	clamp(mut scalar)!
	return x25519_generic(mut dst, mut scalar, point)
}

@[direct_array_access; inline]
fn x25519_generic(mut dst []u8, mut scalar []u8, point []u8) ![]u8 {
	// we dont check arrays length here, its has been checked
	// on the underlying scalar_mult routine
	if is_base_point(point) {
		scalar_base_mult(mut dst, mut scalar)!
		// check for base_point result
		if is_base_point(dst) {
			return error('dst: get base_point')
		}
	} else {
		scalar_mult(mut dst, mut scalar, point)!
		// check for zeros point result
		if is_zero_point(dst) {
			return error('bad input point: low order point')
		}
	}
	return dst
}

// scalar_base_mult performs scalar * base_point
@[direct_array_access]
fn scalar_base_mult(mut dst []u8, mut scalar []u8) ! {
	scalar_mult(mut dst, mut scalar, base_point)!
}

// scalar_mult performs scalar multiplicatipn (scalar * point) on the curve25519 curve.
// for performance reason, scalar marked as mutable.
// scalar_mult is the main routine to perform scalar multiplication
// between scalar key and point on the curve25519 curve.
@[direct_array_access; inline]
fn scalar_mult(mut dst []u8, mut scalar []u8, point []u8) ! {
	if dst.len != point_size {
		return error('bad dst length')
	}
	if scalar.len != scalar_size {
		return error('scalar.lenght != 32')
	}
	if point.len != point_size {
		return error('point.lenght != 32')
	}
	// Note: we  dont clamping scalar here, and responsible to the caller
	// to do the clamping. Its assumed scalar has been clamped.
	mut x1 := edwards25519.Element{}
	mut x2 := edwards25519.Element{}
	mut z2 := edwards25519.Element{}
	mut x3 := edwards25519.Element{}
	mut z3 := edwards25519.Element{}
	mut tmp0 := edwards25519.Element{}
	mut tmp1 := edwards25519.Element{}

	x1.set_bytes(point[..])!
	x2.one()
	x3.set(x1)
	z3.one()

	mut swap := 0
	for pos := 254; pos >= 0; pos-- {
		mut b := scalar[pos / 8] >> u32(pos & 7)
		b &= 1
		swap = swap ^ int(b)
		x2.swap(mut x3, swap)
		z2.swap(mut z3, swap)
		swap = int(b)

		tmp0.subtract(x3, z3)
		tmp1.subtract(x2, z2)
		x2.add(x2, z2)
		z2.add(x3, z3)
		z3.multiply(tmp0, x2)
		z2.multiply(z2, tmp1)
		tmp0.square(tmp1)
		tmp1.square(x2)
		x3.add(z3, z2)
		z2.subtract(z3, z2)
		x2.multiply(tmp1, tmp0)
		tmp1.subtract(tmp1, tmp0)
		z2.square(z2)

		z3.mult_32(tmp1, 121666)
		x3.square(x3)
		tmp0.add(tmp0, z3)
		z3.multiply(x1, z2)
		z2.multiply(tmp1, tmp0)
	}

	x2.swap(mut x3, swap)
	z2.swap(mut z3, swap)

	z2.invert(z2)
	x2.multiply(x2, z2)
	copy(mut dst, x2.bytes())
}

// Utility helpers
//

@[direct_array_access; inline]
fn is_zero_point(point []u8) bool {
	if point.len != point_size {
		return false
	}
	return subtle.constant_time_compare(point, zero_point) == 1
}

@[direct_array_access; inline]
fn is_base_point(point []u8) bool {
	if point.len != point_size {
		return false
	}
	return subtle.constant_time_compare(point, base_point) == 1
}

// clamp clears out some bits of seed bytes
@[direct_array_access; inline]
fn clamp(mut seed []u8) ! {
	if seed.len != scalar_size {
		return error('bad seed sizes for clamp')
	}
	// According to RFC 7748, for x25519, in order to decode 32 random bytes
	// as an integer scalar, set the three least significant bits of the first byte
	// and the most significant bit of the last to zero,
	// set the second most significant bit of the last byte to 1
	//
	seed[0] &= 248
	seed[31] &= 127
	seed[31] |= 64
}

// is_zero returns whether seed is all zeroes in constant time.
fn is_zero(seed []u8) bool {
	mut acc := u8(0)
	for b in seed {
		acc |= b
	}
	return acc == 0
}