jam-std-common 0.1.16

use super::concat;
use crate::{availability::AvailabilityStatementHash, simple::TrancheIndex};
use ed25519_consensus::{Error as SignatureError, VerificationKey};
use jam_types::{hex::HexDisplay, CoreIndex, HeaderHash, OpaqueEd25519Public, WorkReportHash};
use scale::{ConstEncodedLen, Decode, Encode, MaxEncodedLen};

pub const SIGNATURE_LEN: usize = 64;
#[derive(Clone, Copy, Eq, PartialEq)]
pub struct Signature(pub ed25519_consensus::Signature);

pub const PUBLIC_LEN: usize = 32;
#[derive(Clone, Copy, Eq, Ord, PartialEq, PartialOrd)]
pub struct Public(pub ed25519_consensus::VerificationKeyBytes);
impl Default for Public {
	fn default() -> Self {
		Self([0u8; PUBLIC_LEN].into())
	}
}
impl From<OpaqueEd25519Public> for Public {
	fn from(opaque: OpaqueEd25519Public) -> Self {
		opaque.0.into()
	}
}
impl From<Public> for OpaqueEd25519Public {
	fn from(real: Public) -> Self {
		Self(real.to_bytes())
	}
}

pub const SECRET_LEN: usize = 32;
#[derive(Clone)]
pub struct Secret(pub ed25519_consensus::SigningKey);

impl Signature {
	pub fn null() -> Self {
		Self([0u8; SIGNATURE_LEN].into())
	}
	pub fn to_bytes(self) -> [u8; SIGNATURE_LEN] {
		self.0.to_bytes()
	}
	pub fn to_vec(self) -> Vec<u8> {
		self.to_bytes().to_vec()
	}
}

impl Encode for Signature {
	fn encode_to<T: scale::Output + ?Sized>(&self, dest: &mut T) {
		self.0.to_bytes().encode_to(dest)
	}
	fn size_hint(&self) -> usize {
		self.0.to_bytes().size_hint()
	}
}
impl MaxEncodedLen for Signature {
	fn max_encoded_len() -> usize {
		SIGNATURE_LEN
	}
}
impl ConstEncodedLen for Signature {}
impl Decode for Signature {
	fn decode<I: scale::Input>(input: &mut I) -> Result<Self, scale::Error> {
		Ok(Self(<[u8; SIGNATURE_LEN]>::decode(input)?.into()))
	}
}
impl std::fmt::Debug for Signature {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		write!(f, "Signature")
	}
}
impl From<Signature> for [u8; SIGNATURE_LEN] {
	fn from(sig: Signature) -> Self {
		sig.0.to_bytes()
	}
}

impl Signature {
	pub fn verify(&self, message: &Message, signer: &Public) -> Result<(), SignatureError> {
		signer.verify(message, self)
	}
}

impl std::fmt::Debug for Public {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		write!(f, "Pub({})", HexDisplay::from(self.as_bytes()))
	}
}
impl std::hash::Hash for Public {
	fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
		self.0.as_ref().hash(state)
	}
}
impl Encode for Public {
	fn encode_to<T: scale::Output + ?Sized>(&self, dest: &mut T) {
		AsRef::<[u8; PUBLIC_LEN]>::as_ref(&self).encode_to(dest)
	}
	fn size_hint(&self) -> usize {
		PUBLIC_LEN
	}
}
impl MaxEncodedLen for Public {
	fn max_encoded_len() -> usize {
		PUBLIC_LEN
	}
}
impl ConstEncodedLen for Public {}

impl Decode for Public {
	fn decode<I: scale::Input>(input: &mut I) -> Result<Self, scale::Error> {
		Ok(<[u8; PUBLIC_LEN]>::decode(input)?.into())
	}
}
impl From<[u8; PUBLIC_LEN]> for Public {
	fn from(bytes: [u8; PUBLIC_LEN]) -> Self {
		Self(ed25519_consensus::VerificationKeyBytes::from(bytes))
	}
}
impl AsRef<[u8; PUBLIC_LEN]> for Public {
	fn as_ref(&self) -> &[u8; PUBLIC_LEN] {
		self.as_bytes()
	}
}
impl From<Public> for [u8; PUBLIC_LEN] {
	fn from(p: Public) -> Self {
		p.0.into()
	}
}
impl Public {
	pub fn as_bytes(&self) -> &[u8; PUBLIC_LEN] {
		self.0.as_bytes()
	}

	pub fn to_bytes(self) -> [u8; PUBLIC_LEN] {
		*self.as_bytes()
	}

	pub fn verify(&self, message: &Message, signature: &Signature) -> Result<(), SignatureError> {
		let key: VerificationKey = self.0.try_into()?;
		message.using_encoded(|enc| key.verify(&signature.0, enc))
	}
}

/// A message signed by a node's Ed25519 key.
pub enum Message<'a> {
	/// Availability assurance.
	Assurance(AvailabilityStatementHash),
	/// Guarantee statement.
	Guarantee(WorkReportHash),
	/// Audit announcement statement.
	Announcement(HeaderHash, TrancheIndex, &'a [(CoreIndex, WorkReportHash)]),
	/// Work-report judgement.
	Judgement(bool, WorkReportHash),
	/// X.509 certificate.
	X509(&'a [u8]),
	/// TLS handshake transcript.
	Tls(&'a [u8]),
}

impl Message<'_> {
	/// Call `f` with the encoded message.
	pub fn using_encoded<R>(&self, f: impl FnOnce(&[u8]) -> R) -> R {
		match self {
			Self::Assurance(hash) => f(concat(&mut [0; 13 + 32], &[b"jam_available", &hash.0])),
			Self::Guarantee(hash) => f(concat(&mut [0; 13 + 32], &[b"jam_guarantee", &hash.0])),
			Self::Announcement(anchor, tranche, reports) => {
				let mut buf = Vec::with_capacity(13 + reports.len() * 34 + 32);
				buf.extend_from_slice(b"jam_announce");
				buf.extend_from_slice(&[*tranche]);
				reports.iter().for_each(|t| {
					t.encode_to(&mut buf);
				});
				buf.extend_from_slice(&anchor.0);
				f(&buf)
			},
			Self::Judgement(true, hash) => f(concat(&mut [0; 9 + 32], &[b"jam_valid", &hash.0])),
			Self::Judgement(false, hash) =>
				f(concat(&mut [0; 11 + 32], &[b"jam_invalid", &hash.0])),
			Self::X509(enc) | Self::Tls(enc) => {
				// These messages never begin with "jam":
				//
				// - TLS messages begin with a context string consisting of 64 spaces followed by
				//   "TLS 1.3, ..."; see https://www.rfc-editor.org/rfc/rfc8446.html#section-4.4.3.
				// - X.509 certificates are not signed with a context string, however they always
				//   begin with the byte 0x30 ("0").
				assert!(!enc.starts_with(b"jam"));
				f(enc)
			},
		}
	}
}

#[cfg(any(test, feature = "rand"))]
impl rand::distr::Distribution<Public> for rand::distr::StandardUniform {
	fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> Public {
		let bytes: [u8; PUBLIC_LEN] = rng.random();
		bytes.into()
	}
}

impl Eq for Secret {}
impl PartialEq for Secret {
	fn eq(&self, other: &Self) -> bool {
		self.0.as_ref() == other.0.as_ref()
	}
}
impl Encode for Secret {
	fn encode_to<T: scale::Output + ?Sized>(&self, dest: &mut T) {
		self.0.as_bytes().encode_to(dest)
	}
	fn size_hint(&self) -> usize {
		SECRET_LEN
	}
}
impl MaxEncodedLen for Secret {
	fn max_encoded_len() -> usize {
		SECRET_LEN
	}
}
impl Decode for Secret {
	fn decode<I: scale::Input>(input: &mut I) -> Result<Self, scale::Error> {
		Ok(Self(ed25519_consensus::SigningKey::from(<[u8; SECRET_LEN]>::decode(input)?)))
	}
}
impl AsRef<[u8; SECRET_LEN]> for Secret {
	fn as_ref(&self) -> &[u8; SECRET_LEN] {
		self.as_bytes()
	}
}
impl From<Secret> for [u8; SECRET_LEN] {
	fn from(s: Secret) -> Self {
		s.to_bytes()
	}
}

impl Secret {
	pub fn new(rng: &mut impl rand::CryptoRng) -> Self {
		Self::from_seed(rand::Rng::random(rng))
	}
	pub fn from_seed(seed: [u8; SECRET_LEN]) -> Self {
		Self(seed.into())
	}
	pub fn try_from_seed_slice(seed: &[u8]) -> Result<Self, ed25519_consensus::Error> {
		ed25519_consensus::SigningKey::try_from(seed).map(Self)
	}
	pub fn as_bytes(&self) -> &[u8; SECRET_LEN] {
		self.0.as_bytes()
	}
	pub fn to_bytes(&self) -> [u8; SECRET_LEN] {
		*self.as_bytes()
	}
	pub fn generate() -> Self {
		Self::new(&mut rand::rng())
	}
	pub fn trivial(id: u32) -> Self {
		[id; 8].using_encoded(|data| Self(data.try_into().expect("32 bytes is enough")))
	}
	pub fn sign(&self, message: &Message) -> Signature {
		message.using_encoded(|enc| Signature(self.0.sign(enc)))
	}
	pub fn public(&self) -> Public {
		Public((&self.0).into())
	}
}