Feat/chapter 13 segwit

lib/bitcoin/script.rb

@@ -1,5 +1,6 @@
 require_relative '../bitcoin_data_io'
 require_relative '../encoding_helper'
+require_relative '../hash_helper'
 require_relative './op'
 
 module Bitcoin
@@ -88,14 +89,14 @@ def serialize
       encode_varint(raw.length) + raw
     end
 
-    def evaluate(z) # rubocop:disable Metrics/MethodLength
+    def evaluate(z, witness: nil)
       cmds = @cmds.clone
       stack = []
       altstack = []
 
       while cmds.any?
         cmd = cmds.shift
-        return false unless resolve_cmd(cmd, cmds, stack, altstack, z)
+        return false unless resolve_cmd(cmd, cmds, stack, altstack, z, witness: witness)
       end
 
       return false if stack.empty? || stack.pop.empty?
@@ -110,16 +111,44 @@ def p2sh?(cmds = @cmds)
       && cmds[2] == 135
     end
 
+    def p2wpkh?(cmds = @cmds)
+      cmds.length == 2 \
+      && cmds[0] == 0 \
+      && cmds[1].is_a?(String) && cmds[1].length == 20
+    end
+
+    def p2wsh?(cmds = @cmds)
+      cmds.length == 2 \
+      && cmds[0] == 0 \
+      && cmds[1].is_a?(String) && cmds[1].length == 32
+    end
+
+    def self.p2pkh(hash160)
+      Script.new([118, 169, hash160, 136, 172])
+    end
+
+    def self.p2wpkh(hash160)
+      Script.new([0, hash160])
+    end
+
+    def self.p2wsh(hash256)
+      Script.new([0, hash256])
+    end
+
     private
 
-    def resolve_cmd(cmd, cmds, stack, altstack, z)
+    def resolve_cmd(cmd, cmds, stack, altstack, z, witness: nil)
       if cmd.is_a? Integer
         return execute_operation(cmd, cmds, stack, altstack, z)
       else
         stack.append(cmd)
 
         if p2sh?(cmds)
           return execute_p2sh(cmd, cmds, stack)
+        elsif p2wpkh?
+          return execute_p2wpkh(cmds, stack, witness)
+        elsif p2wsh?
+          return execute_p2wsh(cmds, stack, witness)
         end
       end
 
@@ -142,6 +171,30 @@ def execute_p2sh(cmd, cmds, stack)
       cmds.concat self.class.parse(stream).cmds
     end
 
+    def execute_p2wpkh(cmds, stack, witness)
+      h160 = stack.pop
+      stack.pop
+      cmds.concat witness
+      cmds.concat self.class.p2wpkh(h160).cmds
+    end
+
+    def execute_p2wsh(cmds, stack, witness)
+      s256_stack = stack.pop
+      stack.pop
+      cmds.concat witness[0...-1]
+      witness_script = witness.last
+      s256_script = HashHelper.hash256(witness_script)
+
+      unless s256_stack == s256_script
+        raise "Witness script hash mismatch: \n"\
+          "stack: #{bytes_to_hex(s256_stack)} \n"\
+          "script: #{bytes_to_hex(s256_script)}"
+      end
+
+      stream = encode_varint(witness_script.size) + witness_script
+      cmds.concat parse(stream).cmds
+    end
+
     def raw_serialize
       @cmds.map do |cmd|
         cmd.is_a?(Integer) ? int_to_little_endian(cmd, 1) : serialized_element_prefix(cmd) + cmd

lib/bitcoin/tx.rb

@@ -17,12 +17,13 @@ class Tx
     class TxIn
       include EncodingHelper
 
-      def initialize(prev_tx, prev_index, script_sig = nil, sequence = 0xffffffff)
+      def initialize(prev_tx, prev_index, script_sig = nil, sequence = 0xffffffff, witness = nil)
         @prev_tx = prev_tx
         @prev_index = prev_index
         @script_sig = script_sig || Script.new
         @sequence = sequence
         @tx_fetcher = UriFetcher.new
+        @witness = witness
       end
 
       def self.parse(_io)
@@ -42,6 +43,11 @@ def fetch_tx(testnet: false)
         @tx_fetcher.fetch tx_id, testnet: testnet
       end
 
+      def value(testnet: false)
+        tx = fetch_tx testnet: testnet
+        tx.outs[prev_index].amount
+      end
+
       def script_pubkey(testnet: false)
         tx = fetch_tx testnet: testnet
         tx.outs[prev_index].script_pubkey
@@ -54,7 +60,19 @@ def serialize
         result << to_bytes(sequence, 4, 'little')
       end
 
-      attr_accessor :prev_tx, :prev_index, :script_sig, :sequence
+      def serialize_witness
+        result = int_to_little_endian(@witness.size, 1)
+        @witness.each do |wt|
+          result << if wt.is_a?(Integer)
+                      int_to_little_endian(wt, 1)
+                    else
+                      encode_varint(wt.size) + wt
+                    end
+        end
+        result
+      end
+
+      attr_accessor :prev_tx, :prev_index, :script_sig, :sequence, :witness
     end
 
     class TxOut
@@ -77,6 +95,10 @@ def serialize
     end
 
     def self.parse(_io, _options = {})
+      segwit?(_io) ? parse_segwit(_io, _options) : parse_legacy(_io, _options)
+    end
+
+    def self.parse_legacy(_io, _options)
       io = BitcoinDataIO(_io)
 
       new(_options).tap do |tx|
@@ -87,11 +109,69 @@ def self.parse(_io, _options = {})
       end
     end
 
+    def self.parse_segwit(_io, _options)
+      io = BitcoinDataIO(_io)
+
+      new(_options).tap do |tx|
+        tx.version = io.read_le_int32
+        marker = io.read(2)
+        raise "Not a segwit transaction #{marker}" unless marker == "\x00\x01"
+
+        io.read_varint.times { tx.ins << TxIn.parse(io) }
+        io.read_varint.times { tx.outs << TxOut.parse(io) }
+        tx.read_witness_items(io)
+        tx.locktime = io.read_le_int32
+        tx.segwit = true
+      end
+    end
+
+    def self.segwit?(_io)
+      _io.read(4)
+      flag_byte = _io.read(1)
+      _io.rewind
+
+      flag_byte == "\x00"
+    end
+
+    def read_witness_items(_io)
+      @ins.each do |tx_in|
+        items = []
+        _io.read_varint.times do
+          item_len = _io.read_varint
+          items << if item_len.zero?
+                     0
+                   else
+                     _io.read(item_len)
+                   end
+        end
+        tx_in.witness = items
+      end
+    end
+
     def id
-      HashHelper.hash256(serialize).reverse.unpack('H*')
+      HashHelper.hash256(serialize_legacy).reverse.unpack('H*')
     end
 
     def serialize
+      segwit ? serialize_segwit : serialize_legacy
+    end
+
+    # rubocop:disable Metrics/AbcSize
+    def serialize_segwit
+      result = to_bytes(version, 4, 'little')
+      result << "\x00\x01"
+      result << encode_varint(ins.size)
+      result << ins.map(&:serialize).join
+      result << encode_varint(outs.size)
+      result << outs.map(&:serialize).join
+      result << ins.map(&:serialize_witness).join
+      result << to_bytes(locktime, 4, 'little')
+
+      result
+    end
+    # rubocop:enable Metrics/AbcSize
+
+    def serialize_legacy
       result = to_bytes(version, 4, 'little')
       result << encode_varint(ins.size)
       result << ins.map(&:serialize).join
@@ -102,13 +182,18 @@ def serialize
       result
     end
 
-    attr_accessor :version, :locktime, :ins, :outs
+    attr_accessor :version, :locktime, :ins, :outs, :segwit,
+                  :_hash_prevouts, :_hash_sequence, :_hash_outputs
 
-    def initialize(tx_fetcher: nil, testnet: false)
+    def initialize(tx_fetcher: nil, testnet: false, segwit: false)
       @tx_fetcher = tx_fetcher
       @ins = []
       @outs = []
       @testnet = testnet
+      @segwit = segwit
+      @_hash_prevouts = nil
+      @_hash_sequence = nil
+      @_hash_outputs = nil
     end
 
     def fee
@@ -127,20 +212,92 @@ def sig_hash(input_index, redeem_script = nil)
       from_bytes hash256, 'big'
     end
 
+    # rubocop:disable Metrics/AbcSize
+    def sig_hash_bip143(input_index, redeem_script: nil, witness_script: nil)
+      tx_in = @ins[input_index]
+      result = int_to_little_endian(version, 4)
+
+      result += hash_prevouts + hash_sequence
+      result += tx_in.prev_tx.reverse + int_to_little_endian(tx_in.prev_index, 4)
+      result += build_script_raw(redeem_script, witness_script, tx_in)
+      result += int_to_little_endian(tx_in.value, 8)
+      result += int_to_little_endian(tx_in.sequence, 4)
+      result += hash_outputs
+      result << int_to_little_endian(locktime, 4)
+      result << int_to_little_endian(SIGHASH_ALL, 4)
+
+      hash256 = HashHelper.hash256 result
+
+      from_bytes hash256, 'big'
+    end
+    # rubocop:enable Metrics/AbcSize
+
+    def hash_prevouts
+      unless @_hash_prevouts
+        all_prevouts = ''
+        all_sequence = ''
+        @ins.each do |tx_in|
+          all_prevouts += tx_in.prev_tx.reverse + int_to_little_endian(tx_in.prev_index, 4)
+          all_sequence += int_to_little_endian(tx_in.sequence, 4)
+        end
+        @_hash_prevouts = HashHelper.hash256(all_prevouts)
+        @_hash_sequence = HashHelper.hash256(all_sequence)
+      end
+      @_hash_prevouts
+    end
+
+    def hash_sequence
+      hash_prevouts unless @_hash_sequence
+      @_hash_sequence
+    end
+
+    def hash_outputs
+      unless @_hash_outputs
+        all_outputs = ''
+        @outs.each { |tx_out| all_outputs += tx_out.serialize }
+        @_hash_outputs = HashHelper.hash256(all_outputs)
+      end
+      @_hash_outputs
+    end
+
     def verify_input(input_index)
       tx_in = ins[input_index]
       script_pubkey = tx_in.script_pubkey testnet: @testnet
+      z, witness = build_z_and_witness(script_pubkey, tx_in, input_index)
 
+      combined = tx_in.script_sig + script_pubkey
+      combined.evaluate(z, witness: witness)
+    end
+
+    def build_z_and_witness(script_pubkey, tx_in, input_index) # rubocop:disable Metrics/MethodLength
       if script_pubkey.p2sh?
         cmd = tx_in.script_sig.cmds[-1]
         raw_redeem = encode_varint(cmd.length) + cmd
         redeem_script = Script.parse(StringIO.new(raw_redeem))
+
+        if redeem_script.p2wpkh?
+          [sig_hash_bip143(input_index, redeem_script), tx_in.witness]
+        elsif redeem_script.p2wsh?
+          build_z_from_witness(tx_in, input_index)
+        else
+          [sig_hash(input_index, redeem_script), nil]
+        end
+
+      elsif script_pubkey.p2wpkh?
+        [sig_hash_bip143(input_index, redeem_script: redeem_script), tx_in.witness]
+      elsif script_pubkey.p2wsh?
+        build_z_from_witness(tx_in, input_index)
       else
-        redeem_script = nil
+        [sig_hash(input_index), nil]
       end
-      z = sig_hash(input_index, redeem_script)
-      combined = tx_in.script_sig + script_pubkey
-      combined.evaluate(z)
+    end
+
+    def build_z_from_witness(tx_in, input_index)
+      cmd = tx_in.witness.last
+      raw_witness = encode_varint(cmd.size) + cmd
+      witness_script = Script.parse(StringIO.new(raw_witness))
+
+      [sig_hash_bip143(input_index, witness_script: witness_script), tx_in.witness]
     end
 
     def verify?
@@ -204,6 +361,16 @@ def encode_outs
       encode_varint(outs.size) + outs.map(&:serialize).join
     end
 
+    def build_script_raw(redeem_script, witness_script, tx_in)
+      if witness_script
+        witness_script.serialize
+      elsif redeem_script
+        Script.p2pkh(redeem_script.cms[1]).serialize
+      else
+        Script.p2pkh(tx_in.script_pubkey(testnet: @testnet).cmds[1]).serialize
+      end
+    end
+
     def calculate_fee
       raise 'transaction fetcher not provided' if @tx_fetcher.nil?