serai/processor/bitcoin/src/primitives/output.rs

use std::io;

use ciphersuite::*;
use ciphersuite_kp256::Secp256k1;

use bitcoin_serai::{
  bitcoin::{
    hashes::Hash as HashTrait, consensus::Encodable, script::Instruction, transaction::Transaction,
  },
  wallet::ReceivedOutput as WalletOutput,
};

use borsh::{BorshSerialize, BorshDeserialize};
use serai_db::Get;

use serai_client::{
  primitives::{
    coin::ExternalCoin,
    balance::{Amount, ExternalBalance},
    address::ExternalAddress,
  },
  networks::bitcoin::Address,
};

use primitives::{OutputType, ReceivedOutput};

use crate::{
  primitives::x_coord_to_even_point,
  scan::{offsets_for_key, presumed_origin, extract_serai_data},
};

#[derive(Clone, PartialEq, Eq, Hash, Debug, BorshSerialize, BorshDeserialize)]
pub(crate) struct OutputId([u8; 36]);
impl Default for OutputId {
  fn default() -> Self {
    Self([0; 36])
  }
}
impl AsRef<[u8]> for OutputId {
  fn as_ref(&self) -> &[u8] {
    self.0.as_ref()
  }
}
impl AsMut<[u8]> for OutputId {
  fn as_mut(&mut self) -> &mut [u8] {
    self.0.as_mut()
  }
}

#[derive(Clone, PartialEq, Eq, Debug)]
pub(crate) struct Output {
  kind: OutputType,
  presumed_origin: Option<Address>,
  pub(crate) output: WalletOutput,
  data: Vec<u8>,
}

impl Output {
  pub(crate) fn new(
    getter: &impl Get,
    key: <Secp256k1 as WrappedGroup>::G,
    tx: &Transaction,
    output: WalletOutput,
  ) -> Self {
    Self {
      kind: offsets_for_key(key)
        .into_iter()
        .find_map(|(kind, offset)| (offset == output.offset()).then_some(kind))
        .expect("scanned output for unknown offset"),
      presumed_origin: presumed_origin(getter, tx),
      output,
      data: extract_serai_data(tx),
    }
  }

  pub(crate) fn new_with_presumed_origin(
    key: <Secp256k1 as WrappedGroup>::G,
    tx: &Transaction,
    presumed_origin: Option<Address>,
    output: WalletOutput,
  ) -> Self {
    Self {
      kind: offsets_for_key(key)
        .into_iter()
        .find_map(|(kind, offset)| (offset == output.offset()).then_some(kind))
        .expect("scanned output for unknown offset"),
      presumed_origin,
      output,
      data: extract_serai_data(tx),
    }
  }
}

impl ReceivedOutput<<Secp256k1 as WrappedGroup>::G, Address> for Output {
  type Id = OutputId;
  type TransactionId = [u8; 32];

  fn kind(&self) -> OutputType {
    self.kind
  }

  fn id(&self) -> Self::Id {
    let mut id = OutputId::default();
    self.output.outpoint().consensus_encode(&mut id.as_mut()).unwrap();
    id
  }

  fn transaction_id(&self) -> Self::TransactionId {
    let mut res = self.output.outpoint().txid.to_raw_hash().to_byte_array();
    res.reverse();
    res
  }

  fn key(&self) -> <Secp256k1 as WrappedGroup>::G {
    // We read the key from the script pubkey so we don't have to independently store it
    let script = &self.output.output().script_pubkey;

    // These assumptions are safe since it's an output we successfully scanned
    assert!(script.is_p2tr());
    let Instruction::PushBytes(key) = script.instructions_minimal().last().unwrap().unwrap() else {
      panic!("last item in v1 Taproot script wasn't bytes")
    };
    let key = x_coord_to_even_point(key.as_ref())
      .expect("last item in scanned v1 Taproot script wasn't a valid x-only public key");

    // The output's key minus the output's offset is the root key
    key - (<Secp256k1 as WrappedGroup>::G::GENERATOR * self.output.offset())
  }

  fn presumed_origin(&self) -> Option<Address> {
    self.presumed_origin.clone()
  }

  fn balance(&self) -> ExternalBalance {
    ExternalBalance { coin: ExternalCoin::Bitcoin, amount: Amount(self.output.value()) }
  }

  fn data(&self) -> &[u8] {
    &self.data
  }

  fn write<W: io::Write>(&self, writer: &mut W) -> io::Result<()> {
    self.kind.write(writer)?;
    let presumed_origin: Option<ExternalAddress> = self.presumed_origin.clone().map(Into::into);
    presumed_origin.serialize(writer)?;
    self.output.write(writer)?;
    writer.write_all(&u16::try_from(self.data.len()).unwrap().to_le_bytes())?;
    writer.write_all(&self.data)
  }

  fn read<R: io::Read>(mut reader: &mut R) -> io::Result<Self> {
    Ok(Output {
      kind: OutputType::read(reader)?,
      presumed_origin: {
        Option::<ExternalAddress>::deserialize_reader(&mut reader)
          .map_err(|e| io::Error::other(format!("couldn't decode ExternalAddress: {e:?}")))?
          .map(|address| {
            Address::try_from(address)
              .map_err(|()| io::Error::other("couldn't decode Address from ExternalAddress"))
          })
          .transpose()?
      },
      output: WalletOutput::read(reader)?,
      data: {
        let mut data_len = [0; 2];
        reader.read_exact(&mut data_len)?;

        let mut data = vec![0; usize::from(u16::from_le_bytes(data_len))];
        reader.read_exact(&mut data)?;
        data
      },
    })
  }
}