ipld-nostd 0.2.1

use {
	super::Error,
	crate::varint::{self, encode as varint_encode},
	alloc::vec::Vec,
	core::{convert::TryInto, fmt::Debug},
	core2::io,
};

/// A Multihash instance that only supports the basic functionality and no
/// hashing.
///
/// With this Multihash implementation you can operate on Multihashes in a
/// generic way, but no hasher implementation is associated with the code.
///
/// # Example
///
/// ```
/// use multihash::Multihash;
///
/// const Sha3_256: u64 = 0x16;
/// let digest_bytes = [
/// 	0x16, 0x20, 0x64, 0x4b, 0xcc, 0x7e, 0x56, 0x43, 0x73, 0x04, 0x09, 0x99,
/// 	0xaa, 0xc8, 0x9e, 0x76, 0x22, 0xf3, 0xca, 0x71, 0xfb, 0xa1, 0xd9, 0x72,
/// 	0xfd, 0x94, 0xa3, 0x1c, 0x3b, 0xfb, 0xf2, 0x4e, 0x39, 0x38,
/// ];
/// let mh = Multihash::<32>::from_bytes(&digest_bytes).unwrap();
/// assert_eq!(mh.code(), Sha3_256);
/// assert_eq!(mh.size(), 32);
/// assert_eq!(mh.digest(), &digest_bytes[2..]);
/// ```
#[derive(Clone, Copy, Debug, Eq, Ord, PartialOrd)]
pub struct Multihash<const S: usize> {
	/// The code of the Multihash.
	code: u64,
	/// The actual size of the digest in bytes (not the allocated size).
	size: u8,
	/// The digest.
	digest: [u8; S],
}

impl<const S: usize> Default for Multihash<S> {
	fn default() -> Self {
		Self {
			code: 0,
			size: 0,
			digest: [0; S],
		}
	}
}

impl<const S: usize> Multihash<S> {
	/// Wraps the digest in a multihash.
	pub const fn wrap(code: u64, input_digest: &[u8]) -> Result<Self, Error> {
		if input_digest.len() > S {
			return Err(Error::invalid_size(input_digest.len() as _));
		}
		let size = input_digest.len();
		let mut digest = [0; S];
		let mut i = 0;
		while i < size {
			digest[i] = input_digest[i];
			i += 1;
		}
		Ok(Self {
			code,
			size: size as u8,
			digest,
		})
	}

	/// Returns the code of the multihash.
	pub const fn code(&self) -> u64 {
		self.code
	}

	/// Returns the size of the digest.
	pub const fn size(&self) -> u8 {
		self.size
	}

	/// Returns the digest.
	pub fn digest(&self) -> &[u8] {
		&self.digest[..self.size as usize]
	}

	/// Reads a multihash from a byte stream.
	pub fn read<R: io::Read>(r: R) -> Result<Self, Error>
	where
		Self: Sized,
	{
		let (code, size, digest) = read_multihash(r)?;
		Ok(Self { code, size, digest })
	}

	/// Parses a multihash from a bytes.
	///
	/// You need to make sure the passed in bytes have the correct length. The
	/// digest length needs to match the `size` value of the multihash.
	pub fn from_bytes(mut bytes: &[u8]) -> Result<Self, Error>
	where
		Self: Sized,
	{
		let result = Self::read(&mut bytes)?;
		// There were more bytes supplied than read
		if !bytes.is_empty() {
			return Err(Error::invalid_size(bytes.len().try_into().expect(
				"Currently the maximum size is 255, therefore always fits into usize",
			)));
		}

		Ok(result)
	}

	/// Writes a multihash to a byte stream, returning the written size.
	pub fn write<W: io::Write>(&self, w: W) -> Result<usize, Error> {
		write_multihash(w, self.code(), self.size(), self.digest())
	}

	/// Returns the length in bytes needed to encode this multihash into bytes.
	pub fn encoded_len(&self) -> usize {
		let mut code_buf = varint_encode::u64_buffer();
		let code = varint_encode::u64(self.code, &mut code_buf);

		let mut size_buf = varint_encode::u8_buffer();
		let size = varint_encode::u8(self.size, &mut size_buf);

		code.len() + size.len() + usize::from(self.size)
	}

	/// Returns the bytes of a multihash.
	pub fn to_bytes(&self) -> Vec<u8> {
		let mut bytes = Vec::with_capacity(self.size().into());
		let written = self
			.write(&mut bytes)
			.expect("writing to a vec should never fail");
		debug_assert_eq!(written, bytes.len());
		bytes
	}

	/// Truncates the multihash to the given size. It's up to the caller to ensure
	/// that the new size is secure (cryptographically) to use.
	///
	/// If the new size is larger than the current size, this method does nothing.
	pub fn truncate(&self, size: u8) -> Self {
		let mut mh = *self;
		mh.size = mh.size.min(size);
		mh
	}

	/// Resizes the backing multihash buffer.
	///
	/// This function fails if the hash digest is larger than the target size.
	pub fn resize<const R: usize>(&self) -> Result<Multihash<R>, Error> {
		let size = self.size as usize;
		if size > R {
			return Err(Error::invalid_size(self.size as u64));
		}
		let mut mh = Multihash {
			code: self.code,
			size: self.size,
			digest: [0; R],
		};
		mh.digest[..size].copy_from_slice(&self.digest[..size]);
		Ok(mh)
	}

	/// Decomposes struct, useful when needing a `Sized` array or moving all the
	/// data into another type
	///
	/// It is recommended to use `digest()` `code()` and `size()` for most cases.
	pub fn into_inner(self) -> (u64, [u8; S], u8) {
		let Self { code, digest, size } = self;
		(code, digest, size)
	}
}

// Don't hash the whole allocated space, but just the actual digest
#[allow(clippy::derived_hash_with_manual_eq)]
impl<const S: usize> core::hash::Hash for Multihash<S> {
	fn hash<T: core::hash::Hasher>(&self, state: &mut T) {
		self.code.hash(state);
		self.digest().hash(state);
	}
}

impl<const S: usize> From<Multihash<S>> for Vec<u8> {
	fn from(multihash: Multihash<S>) -> Self {
		multihash.to_bytes()
	}
}

impl<const A: usize, const B: usize> PartialEq<Multihash<B>> for Multihash<A> {
	fn eq(&self, other: &Multihash<B>) -> bool {
		// NOTE: there's no need to explicitly check the sizes, that's implicit in
		// the digest.
		self.code == other.code && self.digest() == other.digest()
	}
}

impl<const S: usize> scale::Encode for Multihash<S> {
	fn encode_to<EncOut: scale::Output + ?Sized>(&self, dest: &mut EncOut) {
		self.code.encode_to(dest);
		self.size.encode_to(dest);
		// **NOTE** We write the digest directly to dest, since we have known the
		// size of digest.
		//
		// We do not choose to encode &[u8] directly, because it will add extra
		// bytes (the compact length of digest). For a valid multihash, the length
		// of digest must equal to `size`. Therefore, we can only read raw bytes
		// whose length is equal to `size` when decoding.
		dest.write(self.digest());
	}
}

impl<const S: usize> scale::EncodeLike for Multihash<S> {}

impl<const S: usize> scale::Decode for Multihash<S> {
	fn decode<DecIn: scale::Input>(
		input: &mut DecIn,
	) -> Result<Self, scale::Error> {
		let mut mh = Multihash {
			code: scale::Decode::decode(input)?,
			size: scale::Decode::decode(input)?,
			digest: [0; S],
		};
		if mh.size as usize > S {
			return Err(scale::Error::from("invalid size"));
		}
		// For a valid multihash, the length of digest must equal to the size.
		input.read(&mut mh.digest[..mh.size as usize])?;
		Ok(mh)
	}
}

/// Writes the multihash to a byte stream.
fn write_multihash<W>(
	mut w: W,
	code: u64,
	size: u8,
	digest: &[u8],
) -> Result<usize, Error>
where
	W: io::Write,
{
	let mut code_buf = varint_encode::u64_buffer();
	let code = varint_encode::u64(code, &mut code_buf);

	let mut size_buf = varint_encode::u8_buffer();
	let size = varint_encode::u8(size, &mut size_buf);

	let written = code.len() + size.len() + digest.len();

	w.write_all(code)
		.map_err(super::error::io_to_multihash_error)?;
	w.write_all(size)
		.map_err(super::error::io_to_multihash_error)?;
	w.write_all(digest)
		.map_err(super::error::io_to_multihash_error)?;

	Ok(written)
}

/// Reads a multihash from a byte stream that contains a full multihash (code,
/// size and the digest)
///
/// Returns the code, size and the digest. The size is the actual size and not
/// the maximum/allocated size of the digest.
///
/// Currently the maximum size for a digest is 255 bytes.
fn read_multihash<R, const S: usize>(
	mut r: R,
) -> Result<(u64, u8, [u8; S]), Error>
where
	R: io::Read,
{
	let code = read_u64(&mut r)?;
	let size = read_u64(&mut r)?;

	if size > S as u64 || size > u8::MAX as u64 {
		return Err(Error::invalid_size(size));
	}

	let mut digest = [0; S];
	r.read_exact(&mut digest[..size as usize])
		.map_err(super::error::io_to_multihash_error)?;
	Ok((code, size as u8, digest))
}

pub(crate) fn read_u64<R: io::Read>(mut r: R) -> Result<u64, Error> {
	use varint::decode;
	let mut b = varint_encode::u64_buffer();
	for i in 0..b.len() {
		let n = r
			.read(&mut (b[i..i + 1]))
			.map_err(super::error::io_to_multihash_error)?;
		if n == 0 {
			return Err(Error::insufficient_varint_bytes());
		} else if decode::is_last(b[i]) {
			return decode::u64(&b[..=i])
				.map(|decoded| decoded.0)
				.map_err(super::error::varint_decode_to_multihash_error);
		}
	}
	Err(Error::varint_overflow())
}

#[cfg(test)]
mod tests {
	use super::*;

	#[test]
	fn test_scale() {
		use scale::{Decode, Encode};

		let mh1 = Multihash::<32>::wrap(0, b"hello world").unwrap();
		// println!("mh1: code = {}, size = {}, digest = {:?}", mh1.code(),
		// mh1.size(), mh1.digest());
		let mh1_bytes = mh1.encode();
		// println!("Multihash<32>: {}", hex::encode(&mh1_bytes));
		let mh2: Multihash<32> = Decode::decode(&mut &mh1_bytes[..]).unwrap();
		assert_eq!(mh1, mh2);

		let mh3 = Multihash::<64>::wrap(0, b"hello world").unwrap();
		// println!("mh3: code = {}, size = {}, digest = {:?}", mh3.code(),
		// mh3.size(), mh3.digest());
		let mh3_bytes = mh3.encode();
		// println!("Multihash<64>: {}", hex::encode(&mh3_bytes));
		let mh4: Multihash<64> = Decode::decode(&mut &mh3_bytes[..]).unwrap();
		assert_eq!(mh3, mh4);

		assert_eq!(mh1_bytes, mh3_bytes);
	}

	#[test]
	fn test_eq_sizes() {
		let mh1 = Multihash::<32>::default();
		let mh2 = Multihash::<64>::default();
		assert_eq!(mh1, mh2);
	}

	#[test]
	fn decode_non_minimal_error() {
		// This is a non-minimal varint.
		let data = [241, 0, 0, 0, 0, 0, 128, 132, 132, 132, 58];
		let result = read_u64(&data[..]);
		assert!(result.is_err());
	}
}