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

snap_coin/core/
utxo.rs
