Merge MessageHdr into Message and make Ipv6Address a BasePrimitive

The former makes setting payload_length sensibly much easier, with
no drawbacks that I can see, while the latter lets us override
assign to make comparisons between instances sane.

All good.

lib/btc_wire_proto.rb

@@ -4,41 +4,58 @@ require 'ipaddress'
 # Wraps the IPAddress::IPv6 class found in the ipaddress gem to provide easier
 # handling of binary-format IP addresses
 # @author Nick Thomas <nick@lupine.me.uk>
-class Ipv6Address < BinData::Primitive
+class Ipv6Address < BinData::BasePrimitive
-  endian :big
+  register_self if BinData::VERSION < "1.3.2"
   
-  uint128 :address
+  def value_to_binary_string(v)
+    ip = any_to_ipaddress(v)
+    raise ArgumentError.new("Can't convert #{v.class} to Ipv6Address") unless ip
+    ip.data
+  end
   
-  def get
+  def read_and_return_value(io)
-    if address
+    address = read_uint128(io)
-      if address >= 0xffff00000000 && address <= 0xffffffffffff # v6-mapped v4
+    if address >= 0xffff00000000 && address <= 0xffffffffffff # v6-mapped v4
-        IPAddress::IPv6::Mapped.parse_u128(address)
+      IPAddress::IPv6::Mapped.parse_u128(address)
-      else # v6
+    else # v6
-        IPAddress::IPv6.parse_u128(address)
+      IPAddress::IPv6.parse_u128(address)
-      end
-    else
-      nil # Nothing set
     end
   end
   
-  def set(v)
+  def sensible_default
+    IPAddress("::")
+  end
+  
+  def assign(v)
+    super(any_to_ipaddress(v))
+  end
+
+  protected
+  
+  def any_to_ipaddress(v)
     v = IPAddress(v) if v.is_a?(String)
-    v = IPAddress::IPv6::Mapped.new(v.to_s) if v.is_a?(IPAddress::IPv4)
+
-    if v.is_a?(Fixnum) || v.is_a?(Bignum)
+    if v.is_a?(IPAddress::IPv4)
-      v = if v < 2**32
+      IPAddress::IPv6::Mapped.new("::ffff:#{v.to_s}")
-        IPAddress::IPv6::Mapped.parse_u128(v)
+    elsif v.is_a?(IPAddress::IPv6)
+      v
+    elsif v.is_a?(Fixnum) || v.is_a?(Bignum) || v.respond_to?(:to_u128)
+      v = v.to_u128 if v.respond_to?(:to_u128)
+      if v < 2**32
+        IPAddress::IPv6::Mapped.new("::ffff:#{IPAddress::IPv4.parse_u32(v).to_s}")
       else
-        IPAddress::IPv6::parse_u128(v)
+        IPAddress::IPv6.parse_u128(v)      
       end
-    end
-      
-
-    if v.respond_to?(:to_u128)
-      self.address = v.to_u128
     else
-      raise ArgumentError.new("Can't set #{v.class} to an IPv6Address")      
+      nil
     end
   end
+
+  def read_uint128(io)
+    top, middle1, middle2, bottom = io.readbytes(16).unpack("NNNN")
+    (top << 96) | (middle1 << 64) | (middle2 << 32) | (bottom)
+  end
+
 end
 
 # Implementation of the BitCoin wire protocol, written using bindata.
@@ -395,30 +412,6 @@ module BtcWireProto
 
   ## Top-level message format ##
   
-  # Found at the start of all Bitcoin messages.
-  # @author Nick Thomas <nick@lupine.me.uk>
-  class MessageHdr < BinData::Record
-    endian :little
-    uint32 :magic
-    string :command, :length => 12
-    uint32 :payload_len
-    uint32 :checksum, :onlyif => :has_checksum?
-    
-    def cmd_sym
-      c = command.strip
-      c == "" ? nil : c.to_sym
-    end
-    
-    protected
-    
-    # version and verack messages don't have a checksum. The rest do.
-    # @return[Boolean] does this message header have a checksum field or not?
-    def has_checksum?
-      !%w|version verack|.include?(command.strip)
-    end 
-  end
-
-
   # Everything on the wire is a Message.
   # @author Nick Thomas <nick@lupine.me.uk>
   class Message < BinData::Record
@@ -430,11 +423,11 @@ module BtcWireProto
       @version = v || ::BtcWireProto::CURRENT_VERSION(:main)
     end
 
-    message_hdr :header
+    endian :little
-    
+    uint32 :magic
-    def cmd_sym
+    string :command, :length => 12
-      header ? header.cmd_sym : nil
+    uint32 :payload_len, :value => lambda { payload.num_bytes }
-    end
+    uint32 :checksum, :onlyif => :has_checksum?
     
     choice :payload, :selection => :payload_choice do
       version         "version"
@@ -453,10 +446,20 @@ module BtcWireProto
       null_payload    "null"
     end
 
+    def cmd_sym
+      c = command.strip
+      c == "" ? nil : c.to_sym
+    end
+    # version and verack messages don't have a checksum. The rest do.
+    # @return[Boolean] does this message header have a checksum field or not?
+    def has_checksum?
+      !%w|version verack|.include?(command.to_s.strip)
+    end
+
     # Works out what the payload looks like based on the MessageHdr struct
     # and (potentially) the version
     def payload_choice
-      cmd = header.command.to_s.strip
+      cmd = command.to_s.strip
       return cmd if %w{
         version inv getdata getblocks getheaders tx block headers alert
       }.include?(cmd)

spec/btc_wire_proto_spec.rb

@@ -6,58 +6,50 @@ require 'btc_wire_proto'
 include ::BtcWireProto
 
 describe Ipv6Address do
-  describe "set" do
-    before(:each) { @ip = Ipv6Address.new }
 
-    it "should parse IPv4 strings into IPAddress::IPv6::Mapped objects" do
+  it "should read IPv6 strings into IPv6 addresses" do
-      @ip.set("127.0.0.1")
+    ex_obj = IPAddress("fea1:fea1:fea1:fea1:fea1:fea1:fea1:fea1")
-      @ip.get.should be_a(IPAddress::IPv6::Mapped)
+    ex = ex_obj.data
-      @ip.get.to_s.should == "::ffff:127.0.0.1"
-    end
 
-    it "should parse IPv6 strings into IPAddress::IPv6 objects" do
+    ip = Ipv6Address::read(ex)
-      @ip.set("fe80::1")
+    ip.should == ex_obj
-      @ip.get.should be_a(IPAddress::IPv6)
+    ip.to_binary_s.should == ex
-      @ip.get.to_s.should == "fe80::1"
-    end
     
-    it "should accept IPAddress::IPv(4|6) objects" do
+    ip = Ipv6Address.new(ex_obj)
-      @ip.set(IPAddress("127.0.0.1"))
+    ip.to_binary_s.should == ex
-      @ip.get.should be_a(IPAddress::IPv6::Mapped)
+  end
-      @ip.get.to_s.should == "::ffff:127.0.0.1"
   
-      @ip.set(IPAddress("fe80::1"))
+  it "should read IPv4 strings into IPv6-mapped addresses" do
-      @ip.get.should be_a(IPAddress::IPv6)
+    ip = Ipv6Address.new("127.0.0.1")
-      @ip.get.to_s.should == "fe80::1"
+    ip.should == IPAddress("::ffff:127.0.0.1")
-    end
+  end
 
-    it "should accept arbitrary objects with #to_u128" do
+  it "should read IPAddress objects into IPv6/mapped addresses" do
-      o = Object.new
+    ip = Ipv6Address.new(IPAddress("127.0.0.1"))
-      class << o
+    ip.should == (ip_m = IPAddress("::ffff:127.0.0.1"))
-        def to_u128 ; 666 ; end
+    ip.to_binary_s.should == ip_m.data
-      end
+  end
 
-      @ip.set(o)
+  it "should read arbitrary objects with #to_u128 into IPv6/mapped addresses" do
-      @ip.address.should == 666
+    o = Object.new
+    class << o ; def to_u128 ; 0xffffffffffff ; end ; end  
+    ip = Ipv6Address.new(o)
+    ip.should == IPAddress("::ffff:ffff:ffff")
+  end
 
-    end
+  it "should read Fixnums and Bignums into IPv6/mapped addresses" do
+    ipv4 = IPAddress("127.0.0.1") # Should still be mapped
+    ipv6 = IPAddress("fe80::1")
     
-    it "should treat Fixnums and Bignums as decimal IPs" do
+    ip = Ipv6Address.new(ipv4.to_u32)
-      ipv4 = IPAddress("127.0.0.1") # Should still be mapped
+    ip.should == IPAddress("::ffff:127.0.0.1")
-      ipv6 = IPAddress("fe80::1")
     
-      @ip.set(ipv4.to_u32)
+    ip = Ipv6Address.new(ipv6.to_u128)
-      @ip.get.should be_a(IPAddress::IPv6::Mapped)
+    ip.should == ipv6
-      @ip.get.to_s.should == "::ffff:127.0.0.1"
+  end
 
-      @ip.set(ipv6.to_u128)
+  it "should raise ArgumentError for unparseable objects" do
-      @ip.get.should be_a(IPAddress::IPv6)
+    lambda { Ipv6Address.new("bad IP") }.should raise_error(ArgumentError)
-      @ip.get.to_s.should == "fe80::1"
-    end
-
-    it "should raise an ArgumentError when it can't parse objects" do
-      lambda { @ip.set("bad IP") }.should raise_error(ArgumentError)
-    end
   end
 end
 
@@ -281,34 +273,36 @@ describe ::BtcWireProto do
   
 
   # Messages
-  describe MessageHdr do
+  GOOD_VERSION_DATA = binary(%w{
-  end
+    F9 BE B4 D9 76 65 72 73 69 6F 6E 00 00 00 00 00 55 00 00
+    00 9C 7C 00 00 01 00 00 00 00 00 00 00 E6 15 10 4D 00 00
+    00 00 01 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00
+    00 FF FF 0A 00 00 01 DA F6 01 00 00 00 00 00 00 00 00 00
+    00 00 00 00 00 00 00 00 FF FF 0A 00 00 02 20 8D DD 9D 20
+    2C 3A B4 57 13 00 55 81 01 00
+  })
+  
+  SRV = { :node_network => 1 }
 
   describe Message do
     context "Version message" do
       it "should have a Version payload" do
-        m = Message::new(:header => {:command => 'version'})
+        m = Message::new(:command => 'version')
         m.payload.selection.should == "version"
       end
       
       it "should have a :version command symbol" do
-        m = Message::new(:header => {:command => 'version'})
+        m = Message::new(:command => 'version')
         m.cmd_sym.should == :version      
       end
       
       it "should parse binary data correctly" do
-        m = Message::read(binary(%w{
+        m = Message::read(GOOD_VERSION_DATA)
-          F9 BE B4 D9 76 65 72 73 69 6F 6E 00 00 00 00 00 55 00 00
+
-          00 9C 7C 00 00 01 00 00 00 00 00 00 00 E6 15 10 4D 00 00
+        m.magic.should == BtcWireProto::NETWORKS[:main]
-          00 00 01 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00
+        m.command.should == "version\x00\x00\x00\x00\x00"
-          00 FF FF 0A 00 00 01 DA F6 01 00 00 00 00 00 00 00 00 00
+        m.payload_len.should == 85
-          00 00 00 00 00 00 00 00 FF FF 0A 00 00 02 20 8D DD 9D 20
+        m.has_parameter?(:checksum).should be_false
-          2C 3A B4 57 13 00 55 81 01 00
-        }))
-        m.header.magic.should == BtcWireProto::NETWORKS[:main]
-        m.header.command.should == "version\x00\x00\x00\x00\x00"
-        m.header.payload_len.should == 85
-        m.header.has_parameter?(:checksum).should be_false
         m.payload.version.should == 31900
         m.payload.services.node_network.should == 1
         m.payload.timestamp.should == 1292899814
@@ -317,29 +311,50 @@ describe ::BtcWireProto do
         m.payload.nonce.should == 0x1357B43A2C209DDD
         m.payload.sub_version.should == ""
         m.payload.start_height.should == 98645
+        m.to_binary_s.should == GOOD_VERSION_DATA
+      end
+      
+      it "should generate binary data correctly" do
+        m = Message::read(GOOD_VERSION_DATA)
+        ex = Message::new(
+          :magic => BtcWireProto::NETWORKS[:main], :command => 'version',
+          :payload => {
+            :version => 31900, :services => {:node_network => 1},
+            :timestamp => 1292899814, :nonce => 0x1357B43A2C209DDD,
+            :sub_version => "", :start_height => 98645,
+            :addr_me  => { :services => SRV, :ip => "10.0.0.1", :port => 56054 },
+            :addr_you => { :services => SRV, :ip => "10.0.0.2", :port => 8333 }
+          }
+        )
+        m.should == ex
+        ex.to_binary_s.should == GOOD_VERSION_DATA
 
       end
     end
     
+    GOOD_VERACK_DATA = binary(%w{
+      F9 BE B4 D9 76 65 72 61 63 6B 00 00 00 00 00 00 00 00 00 00
+    })
+
     context "Verack message" do
+
       it "should have no payload" do
-        m = Message.new(:header => {:command => "verack"})
+        m = Message.new(:command => "verack")
         m.payload.selection.should == "null"
       end
       
       it "should have a :verack command symbol" do
-        m = Message.new(:header => {:command => "verack"})
+        m = Message.new(:command => "verack")
         m.cmd_sym.should == :verack
       end
 
       it "should parse the binary data correctly" do
-        m = Message::read(binary(%w{
+        m = Message::read(GOOD_VERACK_DATA)
-          F9 BE B4 D9 76 65 72 61 63 6B 00 00 00 00 00 00 00 00 00 00
+        m.magic.should == BtcWireProto::NETWORKS[:main]
-        }))
+        m.command.should == "verack\x00\x00\x00\x00\x00\x00"
-        m.header.magic.should == BtcWireProto::NETWORKS[:main]
+        m.payload_len.should == 0
-        m.header.command.should == "verack\x00\x00\x00\x00\x00\x00"
+        m.has_parameter?(:chcksum).should be_false
-        m.header.payload_len.should == 0
+        m.to_binary_s.should == GOOD_VERACK_DATA
-        m.header.has_parameter?(:chcksum).should be_false
       end
 
     end