network: Add adaptive capabilities message

network/adaptive.go

@@ -0,0 +1,171 @@
+package network
+
+import (
+	"fmt"
+	"sync"
+
+	"github.com/ethereum/go-ethereum/rlp"
+)
+
+// CapabilityID defines a unique type of capability
+type CapabilityID uint64
+
+// Capability contains a bit vector of flags that define what capability a node has in a specific module
+// The module is defined by the Id.
+type Capability struct {
+	Id  CapabilityID
+	Cap []bool
+}
+
+// NewCapability initializes a new Capability with the given id and specified number of bits in the vector
+func NewCapability(id CapabilityID, bitCount int) *Capability {
+	return &Capability{
+		Id:  id,
+		Cap: make([]bool, bitCount),
+	}
+}
+
+// Set switches the bit at the specified index on
+func (c *Capability) Set(idx int) error {
+	l := len(c.Cap)
+	if idx > l-1 {
+		return fmt.Errorf("index %d out of bounds (len=%d)", idx, l)
+	}
+	c.Cap[idx] = true
+	return nil
+}
+
+// Unset switches the bit at the specified index off
+func (c *Capability) Unset(idx int) error {
+	l := len(c.Cap)
+	if idx > l-1 {
+		return fmt.Errorf("index %d out of bounds (len=%d)", idx, l)
+	}
+	c.Cap[idx] = false
+	return nil
+}
+
+// String implements Stringer interface
+func (c *Capability) String() (s string) {
+	s = fmt.Sprintf("%d:", c.Id)
+	for _, b := range c.Cap {
+		if b {
+			s += "1"
+		} else {
+			s += "0"
+		}
+	}
+	return s
+}
+
+// IsSameAs returns true if the given Capability object has the identical bit settings as the receiver
+func (c *Capability) IsSameAs(cp *Capability) bool {
+	if cp == nil {
+		return false
+	}
+	return isSameBools(c.Cap, cp.Cap)
+}
+
+func isSameBools(left []bool, right []bool) bool {
+	if len(left) != len(right) {
+		return false
+	}
+	for i, b := range left {
+		if b != right[i] {
+			return false
+		}
+	}
+	return true
+}
+
+// Capabilities is the collection of capabilities for a Swarm node
+// It is used both to store the capabilities in the node, and
+// to communicate the node capabilities to its peers
+type Capabilities struct {
+	idx  map[CapabilityID]int // maps the CapabilityIDs to their position in the Caps vector
+	Caps []*Capability
+	mu   sync.Mutex
+}
+
+// NewCapabilities initializes a new Capabilities object
+func NewCapabilities() *Capabilities {
+	return &Capabilities{
+		idx: make(map[CapabilityID]int),
+	}
+}
+
+// adds a capability to the Capabilities collection
+func (c *Capabilities) add(cp *Capability) error {
+	if _, ok := c.idx[cp.Id]; ok {
+		return fmt.Errorf("Capability id %d already registered", cp.Id)
+	}
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.Caps = append(c.Caps, cp)
+	c.idx[cp.Id] = len(c.Caps) - 1
+	return nil
+}
+
+// gets the capability with the specified module id
+// returns nil if the id doesn't exist
+func (c *Capabilities) get(id CapabilityID) *Capability {
+	idx, ok := c.idx[id]
+	if !ok {
+		return nil
+	}
+	return c.Caps[idx]
+}
+
+// String Implements Stringer interface
+func (c *Capabilities) String() (s string) {
+	for _, cp := range c.Caps {
+		if s != "" {
+			s += ","
+		}
+		s += cp.String()
+	}
+	return s
+}
+
+// DecodeRLP implements rlp.RLPDecoder
+// this custom deserializer builds the module id to array index map
+// state of receiver is undefined on error
+func (c *Capabilities) DecodeRLP(s *rlp.Stream) error {
+
+	// make sure we have a pristine receiver
+	c.idx = make(map[CapabilityID]int)
+	c.Caps = []*Capability{}
+
+	// discard the Capabilities struct list item
+	_, err := s.List()
+	if err != nil {
+		return err
+	}
+
+	// discard the Capabilities Caps array list item
+	_, err = s.List()
+	if err != nil {
+		return err
+	}
+
+	// All elements in array should be Capability type
+	for {
+		var cap Capability
+
+		// Decode the Capability from the list item
+		// if error means the end of the list we're done
+		// if not then oh-oh spaghettio's
+		err := s.Decode(&cap)
+		if err != nil {
+			if err == rlp.EOL {
+				break
+			}
+			return err
+		}
+
+		// Add the entry to the Capabilities array
+		c.add(&cap)
+	}
+
+	return nil
+}

network/adaptive_test.go

@@ -0,0 +1,158 @@
+package network
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/rlp"
+)
+
+// TestCapabilitySetUnset tests that setting and unsetting bits yield expected results
+func TestCapabilitySetUnset(t *testing.T) {
+	firstSet := []bool{
+		true, false, false, false, false, false, true, true, false,
+	} // 1000 0011 0
+	firstResult := firstSet
+	secondSet := []bool{
+		false, true, false, true, false, false, true, false, true,
+	} // 0101 0010 1
+	secondResult := []bool{
+		true, true, false, true, false, false, true, true, true,
+	} // 1101 0011 1
+	thirdUnset := []bool{
+		true, false, true, true, false, false, true, false, true,
+	} // 1011 0010 1
+	thirdResult := []bool{
+		false, true, false, false, false, false, false, true, false,
+	} // 0100 0001 0
+
+	c := NewCapability(42, 9)
+	for i, b := range firstSet {
+		if b {
+			c.Set(i)
+		}
+	}
+	if !isSameBools(c.Cap, firstResult) {
+		t.Fatalf("first set result mismatch, expected %v, got %v", firstResult, c.Cap)
+	}
+
+	for i, b := range secondSet {
+		if b {
+			c.Set(i)
+		}
+	}
+	if !isSameBools(c.Cap, secondResult) {
+		t.Fatalf("second set result mismatch, expected %v, got %v", secondResult, c.Cap)
+	}
+
+	for i, b := range thirdUnset {
+		if b {
+			c.Unset(i)
+		}
+	}
+	if !isSameBools(c.Cap, thirdResult) {
+		t.Fatalf("second set result mismatch, expected %v, got %v", thirdResult, c.Cap)
+	}
+}
+
+// TestCapabilitiesControl tests that the methods for manipulating the capabilities bitvectors set values correctly and return errors when they should
+func TestCapabilitiesControl(t *testing.T) {
+
+	// Initialize capability
+	caps := NewCapabilities()
+
+	// Register module. Should succeed
+	c1 := NewCapability(1, 16)
+	err := caps.add(c1)
+	if err != nil {
+		t.Fatalf("RegisterCapabilityModule fail: %v", err)
+	}
+
+	// Fail if capability id already exists
+	c2 := NewCapability(1, 1)
+	err = caps.add(c2)
+	if err == nil {
+		t.Fatalf("Expected RegisterCapabilityModule call with existing id to fail")
+	}
+
+	// More than one capabilities flag vector should be possible
+	c3 := NewCapability(2, 1)
+	err = caps.add(c3)
+	if err != nil {
+		t.Fatalf("RegisterCapabilityModule (second) fail: %v", err)
+	}
+}
+
+// TestCapabilitiesString checks that the string representation of the capabilities is correct
+func TestCapabilitiesString(t *testing.T) {
+	sets1 := []bool{
+		false, false, true,
+	}
+	c1 := NewCapability(42, len(sets1))
+	for i, b := range sets1 {
+		if b {
+			c1.Set(i)
+		}
+	}
+	sets2 := []bool{
+		true, false, false, false, true, false, true, false, true,
+	}
+	c2 := NewCapability(666, len(sets2))
+	for i, b := range sets2 {
+		if b {
+			c2.Set(i)
+		}
+	}
+
+	caps := NewCapabilities()
+	caps.add(c1)
+	caps.add(c2)
+
+	correctString := "42:001,666:100010101"
+	if correctString != caps.String() {
+		t.Fatalf("Capabilities string mismatch; expected %s, got %s", correctString, caps)
+	}
+}
+
+// TestCapabilitiesRLP ensures that a round of serialization and deserialization of Capabilities object
+// results in the correct data
+func TestCapabilitiesRLP(t *testing.T) {
+	c := NewCapabilities()
+	cap1 := &Capability{
+		Id:  42,
+		Cap: []bool{true, false, true},
+	}
+	c.add(cap1)
+	cap2 := &Capability{
+		Id:  666,
+		Cap: []bool{true, false, true, false, true, true, false, false, true},
+	}
+	c.add(cap2)
+	buf := bytes.NewBuffer(nil)
+	err := rlp.Encode(buf, &c)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	cRestored := NewCapabilities()
+	err = rlp.Decode(buf, &cRestored)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	cap1Restored := cRestored.get(cap1.Id)
+	if cap1Restored.Id != cap1.Id {
+		t.Fatalf("cap 1 id not correct, expected %d, got %d", cap1.Id, cap1Restored.Id)
+	}
+	if !cap1.IsSameAs(cap1Restored) {
+		t.Fatalf("cap 1 caps not correct, expected %v, got %v", cap1.Cap, cap1Restored.Cap)
+	}
+
+	cap2Restored := cRestored.get(cap2.Id)
+	if cap2Restored.Id != cap2.Id {
+		t.Fatalf("cap 1 id not correct, expected %d, got %d", cap2.Id, cap2Restored.Id)
+	}
+	if !cap2.IsSameAs(cap2Restored) {
+		t.Fatalf("cap 1 caps not correct, expected %v, got %v", cap2.Cap, cap2Restored.Cap)
+	}
+}