use linked_hash_map::LinkedHashMap;
use std::{
collections::{BTreeSet, HashMap},
fmt::Debug,
};
use crate::{
byron::ByronAddress,
certs::Credential,
crypto::{hash::hash_plutus_data, BootstrapWitness, Vkey, Vkeywitness},
plutus::{PlutusData, PlutusScript, PlutusV1Script, PlutusV2Script, Redeemer},
transaction::{RequiredSigners, TransactionWitnessSet},
NativeScript, Script,
};
use cml_crypto::{
DatumHash, Ed25519KeyHash, Ed25519Signature, PublicKey, RawBytesEncoding, ScriptHash,
};
use super::redeemer_builder::{MissingExunitError, RedeemerBuilderError, RedeemerWitnessKey};
#[derive(Debug, thiserror::Error)]
pub enum WitnessBuilderError {
#[error("Missing the following witnesses: {0:?}")]
MissingWitnesses(RequiredWitnessSet),
#[error("Missing ExUnit: {0}")]
MissingExUnit(#[from] MissingExunitError),
#[error("Redeemer build failed: {0}")]
RedeemBuildFailed(#[from] RedeemerBuilderError),
}
#[derive(Debug, Clone)] pub enum PlutusScriptWitness {
Ref(ScriptHash),
Script(PlutusScript),
}
impl PlutusScriptWitness {
pub fn hash(&self) -> ScriptHash {
match self {
Self::Ref(hash) => *hash,
Self::Script(script) => script.hash(),
}
}
}
impl From<PlutusScript> for PlutusScriptWitness {
fn from(script: PlutusScript) -> Self {
PlutusScriptWitness::Script(script)
}
}
impl From<ScriptHash> for PlutusScriptWitness {
fn from(hash: ScriptHash) -> Self {
PlutusScriptWitness::Ref(hash)
}
}
#[derive(Clone, Debug)]
pub struct PartialPlutusWitness {
pub script: PlutusScriptWitness,
pub redeemer: PlutusData,
}
impl PartialPlutusWitness {
pub fn new(script: PlutusScriptWitness, redeemer: PlutusData) -> Self {
Self { script, redeemer }
}
}
#[derive(Clone, Debug)]
pub enum InputAggregateWitnessData {
NativeScript(NativeScript, NativeScriptWitnessInfo),
PlutusScript(PartialPlutusWitness, RequiredSigners, Option<PlutusData>),
}
impl InputAggregateWitnessData {
pub fn redeemer_plutus_data(&self) -> Option<&PlutusData> {
match self {
InputAggregateWitnessData::PlutusScript(witness, _, _) => Some(&witness.redeemer),
_ => None,
}
}
}
#[derive(Clone, Debug, Default)]
pub struct RequiredWitnessSet {
pub vkeys: BTreeSet<Ed25519KeyHash>,
pub bootstraps: BTreeSet<ByronAddress>,
pub scripts: BTreeSet<ScriptHash>,
pub plutus_data: BTreeSet<DatumHash>,
pub redeemers: BTreeSet<RedeemerWitnessKey>,
pub script_refs: BTreeSet<ScriptHash>,
}
impl RequiredWitnessSet {
pub fn add_vkey_key_hash(&mut self, hash: Ed25519KeyHash) {
self.vkeys.insert(hash);
}
pub fn add_bootstrap(&mut self, address: ByronAddress) {
self.bootstraps.insert(address);
}
pub fn add_script_ref(&mut self, script_hash: ScriptHash) {
self.scripts.remove(&script_hash);
self.script_refs.insert(script_hash);
}
pub fn add_script_hash(&mut self, script_hash: ScriptHash) {
match self.script_refs.get(&script_hash) {
None => {
self.scripts.insert(script_hash);
}
Some(_) => {}
}
}
pub(crate) fn add_from_credential(&mut self, credential: Credential) {
match credential {
Credential::PubKey { hash, .. } => self.add_vkey_key_hash(hash),
Credential::Script { hash, .. } => self.add_script_hash(hash),
}
}
pub fn add_plutus_datum_hash(&mut self, plutus_datum: DatumHash) {
self.plutus_data.insert(plutus_datum);
}
pub fn add_redeemer_tag(&mut self, redeemer: RedeemerWitnessKey) {
self.redeemers.insert(redeemer);
}
pub fn add_all(&mut self, requirements: RequiredWitnessSet) {
self.vkeys.extend(requirements.vkeys);
self.bootstraps.extend(requirements.bootstraps);
self.scripts.extend(requirements.scripts);
self.plutus_data.extend(requirements.plutus_data);
self.redeemers.extend(requirements.redeemers);
}
pub(crate) fn len(&self) -> usize {
self.vkeys.len()
+ self.bootstraps.len()
+ self.scripts.len()
+ self.plutus_data.len()
+ self.redeemers.len()
}
pub(crate) fn add_fake_vkey_witnesses_by_num(&mut self, num: usize) {
for _ in 0..num {
self.add_vkey_key_hash(fake_key_hash(self.vkeys.len() as u8));
}
}
pub(crate) fn add_input_aggregate_fake_witness_data(
&mut self,
data: &InputAggregateWitnessData,
) {
match data {
InputAggregateWitnessData::NativeScript(script, info) => {
match info {
NativeScriptWitnessInfo::Count(num) => {
self.add_fake_vkey_witnesses_by_num(*num)
}
NativeScriptWitnessInfo::Vkeys(ref vkeys) => {
vkeys
.iter()
.cloned()
.for_each(|vkey| self.add_vkey_key_hash(vkey));
}
NativeScriptWitnessInfo::AssumeWorst => {
let num = script.get_required_signers().len();
self.add_fake_vkey_witnesses_by_num(num);
}
}
}
InputAggregateWitnessData::PlutusScript(_witness, required_signers, _option) => {
required_signers
.iter()
.cloned()
.for_each(|vkey| self.add_vkey_key_hash(vkey));
}
}
}
pub fn new() -> Self {
Self::default()
}
}
#[derive(Clone, Default, Debug)]
pub struct TransactionWitnessSetBuilder {
pub vkeys: HashMap<Vkey, Vkeywitness>,
pub bootstraps: HashMap<Vkey, BootstrapWitness>,
pub scripts: HashMap<ScriptHash, Script>,
pub plutus_data: LinkedHashMap<DatumHash, PlutusData>,
pub redeemers: LinkedHashMap<RedeemerWitnessKey, Redeemer>,
pub required_wits: RequiredWitnessSet,
}
impl TransactionWitnessSetBuilder {
pub fn add_vkey(&mut self, vkey_witness: Vkeywitness) {
let vkey = vkey_witness.vkey.clone();
self.vkeys.insert(vkey, vkey_witness);
}
pub fn add_bootstrap(&mut self, bootstrap: BootstrapWitness) {
self.bootstraps
.insert(bootstrap.public_key.clone(), bootstrap);
}
pub fn add_script(&mut self, script: Script) {
self.scripts.insert(script.hash(), script);
}
pub fn get_native_script(&self) -> Vec<NativeScript> {
self.scripts
.iter()
.filter(|entry| self.required_wits.script_refs.get(entry.0).is_none())
.fold(
Vec::<NativeScript>::new(),
|mut acc, script| match &script.1 {
Script::Native { script, .. } => {
acc.push(script.clone());
acc
}
_ => acc,
},
)
}
pub fn get_plutus_v1_script(&self) -> Vec<PlutusV1Script> {
self.scripts
.iter()
.filter(|entry| self.required_wits.script_refs.get(entry.0).is_none())
.fold(
Vec::<PlutusV1Script>::new(),
|mut acc, script| match &script.1 {
Script::PlutusV1 { script, .. } => {
acc.push(script.clone());
acc
}
_ => acc,
},
)
}
pub fn get_plutus_v2_script(&self) -> Vec<PlutusV2Script> {
self.scripts
.iter()
.filter(|entry| self.required_wits.script_refs.get(entry.0).is_none())
.fold(
Vec::<PlutusV2Script>::new(),
|mut acc, script| match &script.1 {
&Script::PlutusV2 { script, .. } => {
acc.push(script.clone());
acc
}
_ => acc,
},
)
}
pub fn add_plutus_datum(&mut self, plutus_datum: PlutusData) {
self.plutus_data
.insert(hash_plutus_data(&plutus_datum), plutus_datum);
}
pub fn get_plutus_datum(&self) -> Vec<PlutusData> {
self.plutus_data.values().cloned().collect()
}
pub fn add_redeemer(&mut self, redeemer: Redeemer) {
self.redeemers
.insert(RedeemerWitnessKey::from(&redeemer), redeemer);
}
pub fn get_redeemer(&self) -> Vec<Redeemer> {
self.redeemers.values().cloned().collect()
}
pub fn add_required_wits(&mut self, required_wits: RequiredWitnessSet) {
self.required_wits.add_all(required_wits)
}
pub fn new() -> Self {
Self::default()
}
pub fn add_existing(&mut self, wit_set: TransactionWitnessSet) {
if let Some(vkeys) = wit_set.vkeywitnesses {
vkeys.into_iter().for_each(|vkey| {
self.add_vkey(vkey);
});
}
if let Some(bootstraps) = wit_set.bootstrap_witnesses {
bootstraps.into_iter().for_each(|bootstrap| {
self.add_bootstrap(bootstrap);
});
}
if let Some(native_scripts) = wit_set.native_scripts {
native_scripts.into_iter().for_each(|native_script| {
self.add_script(native_script.into());
});
}
if let Some(plutus_scripts) = wit_set.plutus_v1_scripts {
plutus_scripts.into_iter().for_each(|plutus_script| {
self.add_script(plutus_script.into());
});
}
if let Some(plutus_scripts) = wit_set.plutus_v2_scripts {
plutus_scripts.into_iter().for_each(|plutus_script| {
self.add_script(plutus_script.into());
});
}
if let Some(redeemers) = wit_set.redeemers {
redeemers.into_iter().for_each(|redeemer| {
self.add_redeemer(redeemer);
});
}
}
pub(crate) fn add_input_aggregate_real_witness_data(
&mut self,
data: &InputAggregateWitnessData,
) {
match data {
InputAggregateWitnessData::NativeScript(script, _info) => {
self.add_script(script.clone().into());
}
InputAggregateWitnessData::PlutusScript(witness, _info, option) => {
match &witness.script {
PlutusScriptWitness::Script(plutus_script) => {
self.add_script(plutus_script.clone().into());
}
PlutusScriptWitness::Ref(_) => {
}
}
if let Some(data) = option {
self.add_plutus_datum(data.clone());
}
}
}
}
pub fn build(self) -> TransactionWitnessSet {
let mut result = TransactionWitnessSet::new();
let native_scripts = self.get_native_script();
let plutus_v1_scripts = self.get_plutus_v1_script();
let plutus_v2_scripts = self.get_plutus_v2_script();
let plutus_datums = self.get_plutus_datum();
if !self.vkeys.is_empty() {
result.vkeywitnesses = Some(self.vkeys.into_values().collect());
}
if !self.bootstraps.is_empty() {
result.bootstrap_witnesses = Some(self.bootstraps.into_values().collect());
}
if !native_scripts.is_empty() {
result.native_scripts = Some(native_scripts);
}
if !plutus_v1_scripts.is_empty() {
result.plutus_v1_scripts = Some(plutus_v1_scripts);
}
if !plutus_v2_scripts.is_empty() {
result.plutus_v2_scripts = Some(plutus_v2_scripts);
}
if !self.plutus_data.is_empty() {
result.plutus_datums = Some(plutus_datums);
}
if !self.redeemers.is_empty() {
result.redeemers = Some(self.redeemers.values().cloned().collect());
}
result
}
pub fn remaining_wits(&self) -> RequiredWitnessSet {
let mut remaining_wits = self.required_wits.clone();
self.vkeys.keys().for_each(|key| {
remaining_wits.vkeys.remove(&key.hash());
});
self.bootstraps.values().for_each(|wit| {
remaining_wits
.bootstraps
.remove(&wit.to_address().unwrap().to_address());
});
self.scripts.keys().for_each(|hash| {
remaining_wits.scripts.remove(hash);
});
self.plutus_data.keys().for_each(|hash| {
remaining_wits.plutus_data.remove(hash);
});
self.redeemers.keys().for_each(|key| {
remaining_wits.redeemers.remove(key);
});
remaining_wits
}
pub fn try_build(&self) -> Result<TransactionWitnessSet, WitnessBuilderError> {
let remaining_wits = self.remaining_wits();
if remaining_wits.len() > 0 {
return Err(WitnessBuilderError::MissingWitnesses(remaining_wits));
}
Ok(self.clone().build())
}
}
pub fn merge_fake_witness(
builder: &mut TransactionWitnessSetBuilder,
required_wits: &RequiredWitnessSet,
) {
let mut remaining_wits = required_wits.clone();
builder.vkeys.keys().for_each(|key| {
remaining_wits.vkeys.remove(&key.hash());
});
builder.bootstraps.values().for_each(|wit| {
remaining_wits
.bootstraps
.remove(&wit.to_address().unwrap().to_address());
});
let fake_prefix = [0u8; 4];
for remaining_vkey in remaining_wits.vkeys.iter() {
let fake_vkey =
PublicKey::from_raw_bytes(&[&fake_prefix, remaining_vkey.to_raw_bytes()].concat())
.unwrap();
let fake_sig = fake_raw_key_sig(0);
let fake_vkey_witness = Vkeywitness::new(fake_vkey, fake_sig);
if !builder.vkeys.contains_key(&fake_vkey_witness.vkey) {
builder.add_vkey(fake_vkey_witness);
}
}
for remaining_bootstrap in remaining_wits.bootstraps.iter() {
let address_content = &remaining_bootstrap.content;
let fake_vkey = PublicKey::from_raw_bytes(
&[&fake_prefix, address_content.address_id.to_raw_bytes()].concat(),
)
.unwrap();
let fake_sig = fake_raw_key_sig(0);
let fake_chaincode = [0u8; 32]; let fake_witness = BootstrapWitness::new(
fake_vkey,
fake_sig,
fake_chaincode.to_vec(),
address_content.addr_attributes.clone(),
)
.unwrap();
if !builder.bootstraps.contains_key(&fake_witness.public_key) {
builder.add_bootstrap(fake_witness);
}
}
}
fn fake_raw_key_sig(id: u8) -> Ed25519Signature {
Ed25519Signature::from_raw_bytes(&[
id, 248, 153, 211, 155, 23, 253, 93, 102, 193, 146, 196, 181, 13, 52, 62, 66, 247, 35, 91,
48, 80, 76, 138, 231, 97, 159, 147, 200, 40, 220, 109, 206, 69, 104, 221, 105, 23, 124, 85,
24, 40, 73, 45, 119, 122, 103, 39, 253, 102, 194, 251, 204, 189, 168, 194, 174, 237, 146,
3, 44, 153, 121, 10,
])
.unwrap()
}
fn fake_key_hash(x: u8) -> Ed25519KeyHash {
Ed25519KeyHash::from_raw_bytes(&[
x, 239, 181, 120, 142, 135, 19, 200, 68, 223, 211, 43, 46, 145, 222, 30, 48, 159, 239, 255,
213, 85, 248, 39, 204, 158, 225, 100,
])
.unwrap()
}
#[derive(Clone, Debug)]
pub enum NativeScriptWitnessInfo {
Count(usize),
Vkeys(Vec<Ed25519KeyHash>),
AssumeWorst,
}
impl NativeScriptWitnessInfo {
pub fn num_signatures(num: usize) -> Self {
NativeScriptWitnessInfo::Count(num)
}
pub fn vkeys(vkeys: Vec<Ed25519KeyHash>) -> Self {
NativeScriptWitnessInfo::Vkeys(vkeys)
}
pub fn assume_signature_count() -> Self {
NativeScriptWitnessInfo::AssumeWorst
}
}
#[cfg(test)]
mod tests {
use cml_crypto::{Bip32PrivateKey, Deserialize, Serialize, TransactionHash};
use crate::byron::make_icarus_bootstrap_witness;
use super::*;
fn fake_raw_key_public(id: u8) -> PublicKey {
PublicKey::from_raw_bytes(&[
id, 118, 57, 154, 33, 13, 232, 114, 14, 159, 168, 148, 228, 94, 65, 226, 154, 181, 37,
227, 11, 196, 2, 128, 28, 7, 98, 80, 209, 88, 91, 205,
])
.unwrap()
}
fn fake_private_key1() -> Bip32PrivateKey {
Bip32PrivateKey::from_raw_bytes(&[
0xb8, 0xf2, 0xbe, 0xce, 0x9b, 0xdf, 0xe2, 0xb0, 0x28, 0x2f, 0x5b, 0xad, 0x70, 0x55,
0x62, 0xac, 0x99, 0x6e, 0xfb, 0x6a, 0xf9, 0x6b, 0x64, 0x8f, 0x44, 0x45, 0xec, 0x44,
0xf4, 0x7a, 0xd9, 0x5c, 0x10, 0xe3, 0xd7, 0x2f, 0x26, 0xed, 0x07, 0x54, 0x22, 0xa3,
0x6e, 0xd8, 0x58, 0x5c, 0x74, 0x5a, 0x0e, 0x11, 0x50, 0xbc, 0xce, 0xba, 0x23, 0x57,
0xd0, 0x58, 0x63, 0x69, 0x91, 0xf3, 0x8a, 0x37, 0x91, 0xe2, 0x48, 0xde, 0x50, 0x9c,
0x07, 0x0d, 0x81, 0x2a, 0xb2, 0xfd, 0xa5, 0x78, 0x60, 0xac, 0x87, 0x6b, 0xc4, 0x89,
0x19, 0x2c, 0x1e, 0xf4, 0xce, 0x25, 0x3c, 0x19, 0x7e, 0xe2, 0x19, 0xa4,
])
.unwrap()
}
fn fake_private_key2() -> Bip32PrivateKey {
Bip32PrivateKey::from_raw_bytes(
&hex::decode("d84c65426109a36edda5375ea67f1b738e1dacf8629f2bb5a2b0b20f3cd5075873bf5cdfa7e533482677219ac7d639e30a38e2e645ea9140855f44ff09e60c52c8b95d0d35fe75a70f9f5633a3e2439b2994b9e2bc851c49e9f91d1a5dcbb1a3").unwrap()
).unwrap()
}
#[test]
fn test_add_fake_vkey_witnesses_by_num() {
let mut builder = RequiredWitnessSet::new();
builder.add_fake_vkey_witnesses_by_num(2);
assert_eq!(builder.vkeys.len(), 2);
builder.add_fake_vkey_witnesses_by_num(1);
assert_eq!(builder.vkeys.len(), 3);
}
#[test]
fn test_add_input_aggregate_witness_data() {
let mut required_wits = RequiredWitnessSet::new();
let data = {
let witness = {
let script = PlutusScript::PlutusV1(PlutusV1Script::new(vec![0]));
PartialPlutusWitness {
script: PlutusScriptWitness::Script(script),
redeemer: PlutusData::new_integer(0u64.into()),
}
};
let missing_signers = vec![fake_raw_key_public(0).hash()];
InputAggregateWitnessData::PlutusScript(witness, missing_signers, None)
};
assert_eq!(required_wits.vkeys.len(), 0);
required_wits.add_input_aggregate_fake_witness_data(&data);
assert_eq!(required_wits.vkeys.len(), 1);
}
#[test]
fn test_add_input_aggregate_witness_data_with_existing_key_hash() {
let mut required_wits = RequiredWitnessSet::new();
let key = fake_raw_key_public(0);
let hash = key.hash();
required_wits.add_vkey_key_hash(hash);
let data = {
let witness = {
let script = PlutusScript::PlutusV1(PlutusV1Script::new(vec![0]));
PartialPlutusWitness {
script: PlutusScriptWitness::Script(script),
redeemer: PlutusData::new_integer(0u64.into()),
}
};
let missing_signers = vec![hash];
InputAggregateWitnessData::PlutusScript(witness, missing_signers, None)
};
assert_eq!(required_wits.vkeys.len(), 1);
required_wits.add_input_aggregate_fake_witness_data(&data);
assert_eq!(required_wits.vkeys.len(), 1);
}
#[test]
fn vkey_test() {
let mut builder = TransactionWitnessSetBuilder::new();
let raw_key_public = fake_raw_key_public(0);
let fake_sig = fake_raw_key_sig(0);
builder.add_vkey(Vkeywitness::new(raw_key_public.clone(), fake_sig.clone()));
builder.add_vkey(Vkeywitness::new(raw_key_public, fake_sig));
builder.add_vkey(Vkeywitness::new(
fake_raw_key_public(1),
fake_raw_key_sig(1),
));
let wit_set = builder.build();
assert_eq!(wit_set.vkeywitnesses.unwrap().len(), 2);
}
#[test]
fn bootstrap_test() {
let mut builder = TransactionWitnessSetBuilder::new();
let wit1 = make_icarus_bootstrap_witness(
TransactionHash::from([0u8; TransactionHash::BYTE_COUNT]),
ByronAddress::from_base58(
"Ae2tdPwUPEZGUEsuMAhvDcy94LKsZxDjCbgaiBBMgYpR8sKf96xJmit7Eho",
)
.unwrap(),
&fake_private_key1(),
);
builder.add_bootstrap(wit1.clone());
builder.add_bootstrap(wit1);
builder.add_bootstrap(make_icarus_bootstrap_witness(
TransactionHash::from([0u8; TransactionHash::BYTE_COUNT]),
ByronAddress::from_base58(
"Ae2tdPwUPEZGUEsuMAhvDcy94LKsZxDjCbgaiBBMgYpR8sKf96xJmit7Eho",
)
.unwrap(),
&fake_private_key2(),
));
let wit_set = builder.build();
assert_eq!(wit_set.bootstrap_witnesses.unwrap().len(), 2);
}
#[test]
fn native_script_test() {
let mut builder = TransactionWitnessSetBuilder::new();
let wit1: Script = NativeScript::new_script_invalid_before(1).into();
builder.add_script(wit1.clone());
builder.add_script(wit1);
builder.add_script(NativeScript::new_script_invalid_before(2).into());
let wit_set = builder.build();
assert_eq!(wit_set.native_scripts.unwrap().len(), 2);
}
#[test]
fn requirement_test_fail() {
let mut builder = TransactionWitnessSetBuilder::new();
let mut required_wits = RequiredWitnessSet::new();
required_wits.add_vkey_key_hash(fake_raw_key_public(0).hash());
required_wits.add_script_hash(NativeScript::new_script_invalid_before(2).hash());
builder.add_required_wits(required_wits);
builder.add_vkey(Vkeywitness::new(
fake_raw_key_public(1),
fake_raw_key_sig(1),
));
assert!(builder.try_build().is_err());
}
#[test]
fn requirement_test_pass() {
let mut builder = TransactionWitnessSetBuilder::new();
let mut required_wits = RequiredWitnessSet::new();
required_wits.add_vkey_key_hash(fake_raw_key_public(0).hash());
builder.add_required_wits(required_wits);
builder.add_vkey(Vkeywitness::new(
fake_raw_key_public(0),
fake_raw_key_sig(0),
));
assert!(builder.try_build().is_ok());
}
#[test]
fn tx_witness_set_roundtrip_test() {
let data = "a102818458205e8379f58f0838234af67f73738f0fee0d8185232e200b8e42887f4f06544a9a5840f5cfea560d2f8645ed624b65bf08cf83346eb5168ee4df0f63ce2d0d5f677db88fef2d5d9f032f09223889b5e85504ab44dd0a0cde1f1fd8f57deefde8c2080658202d3b7d9b806f88f10f1193e94ef97e5c02370c1464f61a30a8f1ac1a46115b2d5829a201581e581c072931653330243cf126aea85d39e73c6bd04601fe77424efb9e371002451a4170cb17";
let witness_set =
TransactionWitnessSet::from_cbor_bytes(&hex::decode(data).unwrap()).unwrap();
let round_trip = witness_set.to_cbor_bytes();
assert_eq!(data, hex::encode(round_trip));
}
}