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#[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 pub utxos: HashMap<TransactionId, Vec<Option<TransactionOutput>>>,
82}
83
84impl UTXOs {
85 pub fn new() -> Self {
87 UTXOs {
88 utxos: HashMap::new(),
89 }
90 }
91
92 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 pub fn execute_transaction(&mut self, transaction: &Transaction) -> UTXODiff {
189 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 UTXODiff {
215 spent: spent_utxos,
216 created: created_utxos,
217 }
218 }
219
220 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}