snap_coin/core/
utxo.rs

1use bincode::{Decode, Encode};
2use chrono::Utc;
3use num_bigint::BigUint;
4use std::{
5    collections::{HashMap, HashSet},
6    ops::Deref,
7};
8use thiserror::Error;
9
10use crate::{
11    core::transaction::{self, Transaction, TransactionId, TransactionOutput},
12    crypto::keys::Public,
13};
14
15#[derive(Error, Debug)]
16pub enum TransactionError {
17    #[error("Transaction timestamp is in the future: {0}")]
18    FutureTimestamp(u64),
19
20    #[error("Transaction missing ID")]
21    MissingId,
22
23    #[error("Transaction hash is invalid: {0}")]
24    InvalidHash(String),
25
26    #[error("Transaction hash does not meet required difficulty: {0}")]
27    InsufficientDifficulty(String),
28
29    #[error("Transaction has no inputs")]
30    NoInputs,
31
32    #[error("Transaction input not found in UTXOs: {0}")]
33    InputNotFound(String),
34
35    #[error("Transaction input index invalid for transaction {tx_id}: {input_tx_id}")]
36    InvalidInputIndex { tx_id: String, input_tx_id: String },
37
38    #[error("Referenced transaction input is already spent")]
39    SpentInputIndex,
40
41    #[error("Transaction input signature is invalid for transaction {0}")]
42    InvalidSignature(String),
43
44    #[error("Double spending detected in the same transaction {0}")]
45    DoubleSpend(String),
46
47    #[error("Transaction output amount cannot be zero for transaction {0}")]
48    ZeroOutput(String),
49
50    #[error("Transaction inputs and outputs don't sum up to same amount for transaction {0}")]
51    SumMismatch(String),
52
53    #[error("Transaction has too many inputs or outputs")]
54    TooManyIO,
55}
56
57/// This represents a singular transaction undo, or a whole block, essentially we need to extend each of these lists to combine with other utxo diffs (prob. from other tx's)
58#[derive(Encode, Decode, Debug, Clone)]
59pub struct UTXODiff {
60    spent: Vec<(TransactionId, usize, TransactionOutput)>,
61    created: Vec<(TransactionId, usize, TransactionOutput)>,
62}
63
64impl UTXODiff {
65    pub fn new_empty() -> Self {
66        Self {
67            spent: Vec::new(),
68            created: Vec::new(),
69        }
70    }
71
72    pub fn extend(&mut self, other: &mut Self) {
73        self.spent.append(&mut other.spent);
74        self.created.append(&mut other.created);
75    }
76}
77
78#[derive(Encode, Decode, Clone, Debug)]
79pub struct UTXOs {
80    /// Map of transaction (id's) and its outputs
81    pub utxos: HashMap<TransactionId, Vec<Option<TransactionOutput>>>,
82}
83
84impl UTXOs {
85    /// Create a new empty UTXO set
86    pub fn new() -> Self {
87        UTXOs {
88            utxos: HashMap::new(),
89        }
90    }
91
92    /// Validate a certain transaction in the context of these UTXOs
93    /// WARNING: Timestamp validation is not checked here!
94    pub fn validate_transaction(
95        &self,
96        transaction: &Transaction,
97        tx_hashing_difficulty: &BigUint,
98    ) -> Result<(), TransactionError> {
99        let tx_id = transaction
100            .transaction_id
101            .ok_or(TransactionError::MissingId)?;
102        let input_signing_buf = transaction
103            .get_input_signing_buf()
104            .map_err(|_| TransactionError::InvalidHash(tx_id.dump_base36()))?;
105        let transaction_hashing_buf = transaction
106            .get_tx_hashing_buf()
107            .map_err(|_| TransactionError::InvalidHash(tx_id.dump_base36()))?;
108
109        if transaction.timestamp > Utc::now().timestamp() as u64 {
110            return Err(TransactionError::FutureTimestamp(transaction.timestamp));
111        }
112
113        if !tx_id.compare_with_data(&transaction_hashing_buf) {
114            return Err(TransactionError::InvalidHash(tx_id.dump_base36()));
115        }
116
117        if BigUint::from_bytes_be(tx_id.deref()) > *tx_hashing_difficulty {
118            return Err(TransactionError::InsufficientDifficulty(
119                tx_id.dump_base36(),
120            ));
121        }
122
123        if transaction.inputs.is_empty() {
124            return Err(TransactionError::NoInputs);
125        }
126
127        if transaction.inputs.len() + transaction.outputs.len() > transaction::MAX_TRANSACTION_IO {
128            return Err(TransactionError::TooManyIO);
129        }
130
131        let mut used_utxos: HashSet<(TransactionId, usize)> = HashSet::new();
132        let mut input_sum: u64 = 0;
133        let mut output_sum: u64 = 0;
134
135        for input in &transaction.inputs {
136            let prev_outputs = self
137                .utxos
138                .get(&input.transaction_id)
139                .ok_or(TransactionError::InputNotFound(tx_id.dump_base36()))?;
140
141            let output = prev_outputs.get(input.output_index).ok_or(
142                TransactionError::InvalidInputIndex {
143                    tx_id: tx_id.dump_base36(),
144                    input_tx_id: input.transaction_id.dump_base36(),
145                },
146            )?;
147
148            if output.is_none() {
149                return Err(TransactionError::SpentInputIndex);
150            }
151
152            if input.signature.is_none()
153                || input
154                    .signature
155                    .unwrap()
156                    .validate_with_public(&output.unwrap().receiver, &input_signing_buf)
157                    .map_or(true, |valid| !valid)
158            {
159                return Err(TransactionError::InvalidSignature(tx_id.dump_base36()));
160            }
161
162            let utxo_key = (input.transaction_id, input.output_index);
163            if used_utxos.contains(&utxo_key) {
164                return Err(TransactionError::DoubleSpend(tx_id.dump_base36()));
165            }
166            used_utxos.insert(utxo_key);
167
168            input_sum += output.unwrap().amount;
169        }
170
171        for output in &transaction.outputs {
172            if output.amount == 0 {
173                return Err(TransactionError::ZeroOutput(tx_id.dump_base36()));
174            }
175            output_sum += output.amount;
176        }
177
178        if input_sum != output_sum {
179            return Err(TransactionError::SumMismatch(tx_id.dump_base36()));
180        }
181
182        Ok(())
183    }
184
185    /// Execute a already valid transaction.
186    /// Returns a UTXODiff object that can be used to undo these changes
187    /// WARNING: Transaction validity must be checked before calling this function
188    pub fn execute_transaction(&mut self, transaction: &Transaction) -> UTXODiff {
189        // Keep track of created and removed utxos
190        let mut spent_utxos: Vec<(TransactionId, usize, TransactionOutput)> = Vec::new();
191        let mut created_utxos: Vec<(TransactionId, usize, TransactionOutput)> = Vec::new();
192        for input in &transaction.inputs {
193            if let Some(outputs) = self.utxos.get_mut(&input.transaction_id) {
194                let spent = outputs[input.output_index].clone().unwrap();
195                outputs[input.output_index] = None;
196
197                if outputs.iter().all(|o| o.is_none()) {
198                    self.utxos.remove(&input.transaction_id);
199                }
200                spent_utxos.push((input.transaction_id, input.output_index, spent));
201            }
202        }
203
204        self.utxos.insert(
205            transaction.transaction_id.unwrap(),
206            transaction.outputs.iter().map(|o| Some(*o)).collect(),
207        );
208
209        for (output_index, output) in transaction.outputs.iter().enumerate() {
210            created_utxos.push((transaction.transaction_id.unwrap(), output_index, *output));
211        }
212
213        // println!("Executing transaction! added {:#?}", created_utxos);
214        UTXODiff {
215            spent: spent_utxos,
216            created: created_utxos,
217        }
218    }
219
220    /// Undo a certain transactions with its UTXODiff
221    pub fn recall_block_utxos(&mut self, diffs: UTXODiff) {
222        for spent in diffs.spent {
223            if self.utxos.get(&spent.0).is_some() {
224                self.utxos.get_mut(&spent.0).unwrap()[spent.1] = Some(spent.2);
225            } else {
226                let mut outputs = vec![None; spent.1 + 1];
227                outputs[spent.1] = Some(spent.2);
228                self.utxos.insert(spent.0, outputs);
229            }
230        }
231
232        for created in diffs.created {
233            if let Some(vec) = self.utxos.get_mut(&created.0) {
234                vec[created.1] = None;
235                if vec.iter().all(|x| x.is_none()) {
236                    self.utxos.remove(&created.0);
237                }
238            }
239        }
240    }
241
242    pub fn calculate_confirmed_balance(&self, address: Public) -> u64 {
243        let mut balance = 0u64;
244
245        for transaction_out in self.utxos.values().flat_map(|utxos| utxos) {
246            if let Some(transaction_out) = transaction_out {
247                if transaction_out.receiver == address {
248                    balance += transaction_out.amount;
249                }
250            }
251        }
252
253        balance
254    }
255
256    pub fn get_utxos(&self, address: Public) -> Vec<TransactionId> {
257        let mut utxos: Vec<TransactionId> = vec![];
258
259        for (transaction_id, outputs) in &self.utxos {
260            if outputs
261                .iter()
262                .any(|output| output.is_some_and(|output| output.receiver == address))
263            {
264                utxos.push(*transaction_id);
265            }
266        }
267
268        utxos
269    }
270}