Toy BTC transaction verification

ValidateBTCTransaction.swift

// MARK: - Primitives

typealias Address = String
typealias Signature = String
typealias Coins = String

// MARK: - Transaction

enum Transaction {
    case coinbase(id: ID, outputs: [Output])
    case standard(id: ID, inputs: [Input], outputs: [Output])
    
    typealias ID = String
    
    struct Input {
        let previousTxID: ID
        let previousTxOutputIndex: Int
        let signature: Signature
    }
    
    struct Output {
        let recipientAddress: Address
        let value: Coins
    }
    
    struct UnspentTransactionOutputID: Hashable {
        let txID: ID
        let txOutputIndex: Int
    }
}

extension Transaction {
    var inputs: [Input] {
        switch self {
        case .coinbase: return []
        case let .standard(_,inputs,_): return inputs
        }
    }
    
    var outputs: [Output] {
        switch self {
        case let .coinbase(_,outputs): return outputs
        case let .standard(_,_,outputs): return outputs
        }
    }
    
    var id: ID {
        switch self {
        case let .coinbase(id,_): return id
        case let .standard(id,_,_): return id
        }
    }
}

extension Transaction.Input {
    var uxtoID: Transaction.UnspentTransactionOutputID {
        return .init(
            txID: previousTxID,
            txOutputIndex: previousTxOutputIndex
        )
    }
}

// MARK: - CustomStringConvertible

extension Transaction.Output: CustomStringConvertible {
    var description: String {
        return "owner: \(recipientAddress.prefix { $0 != "📮" }) | amount: \(value)"
    }
}

extension Transaction.UnspentTransactionOutputID: CustomStringConvertible {
    var description: String {
        return "\(txID)(\(txOutputIndex))"
    }
}

extension Node.UnspentTransactionOutputPool: CustomStringConvertible {
    var description: String {
        return self.map {"\($0.key) | \($0.value) \n"}.reduce("", +)
    }
}

extension Node.TransactionValidationError: CustomStringConvertible {
    var description: String {
        switch self {
        case .inputNotFromUTXOPool:
            return "Input not found in UTXO pool"
        case .inputInvalidSignature:
            return "Invalid input signature"
        case .inputClaimMultipleSameUTXO:
            return "Same UTXO claimed multiple times in inputs"
        case let .totalInputLessThanTotalOutput(inputValue, outputValue):
            return "Total input \(inputValue) is less than total output \(outputValue)"
        }
    }
}

// MARK: - Node

final class Node {
    
    typealias UnspentTransactionOutputPool = [Transaction.UnspentTransactionOutputID: Transaction.Output]
    
    enum TransactionValidationError: Error {
        case inputNotFromUTXOPool
        case inputInvalidSignature
        case inputClaimMultipleSameUTXO
        case totalInputLessThanTotalOutput(inputValue: Coins, outputValue: Coins)
    }
    
    private var utxoPool: UnspentTransactionOutputPool = [:]
    
    init() {}
    
    func addTransaction(_ transaction: Transaction) {
        print("Try adding a transaction...")
        
        switch isValidTransaction(transaction) {
        case .success:
            for (index, output) in transaction.outputs.enumerated() {
                let uxtoID = Transaction.UnspentTransactionOutputID(
                    txID: transaction.id,
                    txOutputIndex: index)
                utxoPool[uxtoID] = output
            }
            
            for input in transaction.inputs {
                utxoPool.removeValue(forKey: input.uxtoID)
            }
            
            print("🤠 hooray your transaction is validated!")
        
        case let .failure(error):
            print("😏 uh oh your transaction is denied. \nReason: \(error)")
        }
        
        print("Current UTXO pool:\n\(utxoPool.description)")
    }
    
    private func isValidTransaction(_ transaction: Transaction) -> Result<Void, TransactionValidationError> {
        switch transaction{
        case .coinbase:
            return .success(())
            
        case let .standard(_, inputs, outputs):
            var claimedUTXOIDs: Set<Transaction.UnspentTransactionOutputID> = .init()
            var totalInputValue: Coins = ""
            
            for input in inputs {
                // Input must be in current UTXO pool
                guard let utxoToClaim = utxoPool[input.uxtoID] else {
                    return .failure(.inputNotFromUTXOPool)
                }
                
                // Same UTXO can't be claimed multiple times
                if claimedUTXOIDs.contains(input.uxtoID) {
                    return .failure(.inputClaimMultipleSameUTXO)
                }
                
                // Input signature must be valid
                let isSignatureValid = verify(
                    signature: input.signature,
                    address: utxoToClaim.recipientAddress,
                    coins: utxoToClaim.value)
                guard isSignatureValid else {
                    return .failure(.inputInvalidSignature)
                }
                
                claimedUTXOIDs.insert(input.uxtoID)
                totalInputValue += utxoToClaim.value
            }
            
            // Total output can't be more than total input
            let totalOutputValue = outputs.map {$0.value}.reduce("", +)
            if totalInputValue >= totalOutputValue {
                return .success(())
            } else {
                return .failure(
                    .totalInputLessThanTotalOutput(inputValue: totalInputValue, outputValue: totalOutputValue)
                )
            }
        }
    }
    
    private func verify(signature: Signature, address: Address, coins: Coins) -> Bool {
        let expectedRecipient = address.prefix { $0 != "📮" }
        let signer =  String(signature.split(separator: "🔐").first!)
        return expectedRecipient == signer
    }
}

// MARK: - Profit!

let node = Node()

/**
 Case 1 - 👩🏻‍🌾 and 🦹🏻‍♂️ both get initial 💰💰💰 from a coinbase transaction
 */
let initialCoins: Transaction = .coinbase(
    id: "🔮",
    outputs: [
        .init(recipientAddress: "👩🏻‍🌾📮", value: "💰💰💰"),
        .init(recipientAddress: "🦹🏻‍♂️📮", value: "💰💰💰"),
    ]
)
node.addTransaction(initialCoins)

/**
 Case 2 - 👩🏻‍🌾 pays 💰💰 to 🦹🏻‍♂️
 */
let realAlicePayBob: Transaction = .standard(
    id: "🍹",
    inputs: [
        .init(
            previousTxID: "🔮",
            previousTxOutputIndex: 0,
            signature: "👩🏻‍🌾🔐🍹"
        ),
    ],
    outputs: [
        .init(recipientAddress: "🦹🏻‍♂️📮", value: "💰💰"),
        .init(recipientAddress: "👩🏻‍🌾📮", value: "💰"), // change
    ]
)
node.addTransaction(realAlicePayBob)

/**
 Case 3 - 🦹🏻‍♂️ counterfeits 💰💰💰 out of thin air
*/
let bobDayDream: Transaction = .standard(
    id: "🍸",
    inputs: [
        .init(
            previousTxID: "🥃",
            previousTxOutputIndex: 1,
            signature: "🦹🏻‍♂️🔐🍸"
        ),
    ],
    outputs: [
        .init(recipientAddress: "🦹🏻‍♂️📮", value: "💰💰💰"),
    ]
)
node.addTransaction(bobDayDream)

/**
 Case 4 - 🦹🏻‍♂️ counterfeits 💰 paid by 👩🏻‍🌾
*/
let fakeAlicePayBob: Transaction = .standard(
    id: "🍺",
    inputs: [
        .init(
            previousTxID: "🍹",
            previousTxOutputIndex: 1,
            signature: "🦹🏻‍♂️🔐🍺"
        ),
    ],
    outputs: [
        .init(recipientAddress: "🦹🏻‍♂️📮", value: "💰"),
    ]
)
node.addTransaction(fakeAlicePayBob)

/**
 Case 5 - 👩🏻‍🌾 tries paying 💰💰 to 🦹🏻‍♂️ by redeeming 💰 twice
 */
let alicePayBobDoubleSpend: Transaction = .standard(
    id: "🥂",
    inputs: [
        .init(
            previousTxID: "🍹",
            previousTxOutputIndex: 1,
            signature: "👩🏻‍🌾🔐🥂"
        ),
        .init(
            previousTxID: "🍹",
            previousTxOutputIndex: 1,
            signature: "👩🏻‍🌾🔐🥂"
        ),
    ],
    outputs: [
        .init(recipientAddress: "🦹🏻‍♂️📮", value: "💰💰"),
    ]
)
node.addTransaction(alicePayBobDoubleSpend)

/**
 Case 6 - 🦹🏻‍♂️ tries aggregating 💰💰💰 + 💰💰 into 💰💰💰💰💰💰
 */
let bobAggregateChangesMoreThanHeOwn: Transaction = .standard(
    id: "🍻",
    inputs: [
        .init(
            previousTxID: "🔮",
            previousTxOutputIndex: 1,
            signature: "🦹🏻‍♂️🔐🍻"
        ),
        .init(
            previousTxID: "🍹",
            previousTxOutputIndex: 0,
            signature: "🦹🏻‍♂️🔐🍻"
        ),
    ],
    outputs: [
        .init(recipientAddress: "🦹🏻‍♂️📮", value: "💰💰💰💰💰💰"),
    ]
)
node.addTransaction(bobAggregateChangesMoreThanHeOwn)