ms-select2: abbreviation tree and NewMSSelect2

Author: ppopth

abbrevTree.go

@@ -0,0 +1,98 @@
+package multistream
+
+import (
+	"crypto/sha256"
+)
+
+type nodeProtocol[T StringLike] struct {
+	protocolID   T
+	tombstoneBit bool
+}
+
+type abbrevTree[T StringLike] struct {
+	root *abbrevNode[T]
+}
+
+type abbrevNode[T StringLike] struct {
+	p        *nodeProtocol[T]
+	children [256]*abbrevNode[T]
+}
+
+func (at *abbrevTree[T]) Abbreviate(pid T) []byte {
+	var result []byte
+	hash := sha256.Sum256([]byte(pid))
+
+	if at.root == nil {
+		return nil
+	}
+
+	current := at.root
+	// go furthest in the tree
+	for _, b := range hash {
+		if current.children[b] != nil {
+			result = append(result, b)
+			current = current.children[b]
+		}
+	}
+
+	if current.p != nil && current.p.protocolID == pid && !current.p.tombstoneBit {
+		return result
+	}
+	return nil
+}
+
+func (at *abbrevTree[T]) AddProtocol(pid T) {
+	hash := sha256.Sum256([]byte(pid))
+
+	if at.root == nil {
+		at.root = &abbrevNode[T]{}
+	}
+
+	current := at.root
+	for idx, b := range hash {
+		if current.children[b] == nil {
+			current.children[b] = &abbrevNode[T]{
+				p: &nodeProtocol[T]{
+					protocolID:   pid,
+					tombstoneBit: false,
+				},
+			}
+			return
+		}
+		current = current.children[b]
+
+		if current.p != nil {
+			if current.p.protocolID == pid {
+				// Resurrect the protocol ID.
+				current.p.tombstoneBit = false
+			} else if !current.p.tombstoneBit {
+				// There is another protocol in this node, so we need to duplicate it down.
+				h := sha256.Sum256([]byte(current.p.protocolID))
+
+				if current.children[h[idx+1]] == nil {
+					// It should be fine to reference the same nodeProtocol instance.
+					current.children[h[idx+1]] = &abbrevNode[T]{p: current.p}
+				}
+			}
+		}
+	}
+}
+
+func (at *abbrevTree[T]) RemoveProtocol(pid T) {
+	hash := sha256.Sum256([]byte(pid))
+
+	if at.root == nil {
+		return
+	}
+	current := at.root
+	for _, b := range hash {
+		if current.children[b] == nil {
+			break
+		}
+		current = current.children[b]
+
+		if current.p.protocolID == pid {
+			current.p.tombstoneBit = true
+		}
+	}
+}

abbrevTree_test.go

@@ -0,0 +1,155 @@
+package multistream
+
+import (
+	"bytes"
+	"crypto/sha256"
+	"testing"
+)
+
+func TestAbbrevTreeAddProtocol(t *testing.T) {
+	tree := &abbrevTree[string]{}
+
+	proto1 := "protocol1"
+	hash1 := sha256.Sum256([]byte(proto1))
+	proto2 := "protocol2"
+	hash2 := sha256.Sum256([]byte(proto2))
+	proto3 := "protocol251" // this one has the same first byte as "protocol1"
+	hash3 := sha256.Sum256([]byte(proto3))
+
+	// make sure we don't make mistakes on the hashes
+	if hash1[0] == hash2[0] {
+		t.Fatal("the first bytes of hash1 and hash2 should be different")
+	}
+	if hash1[0] != hash3[0] {
+		t.Fatal("the first bytes of hash1 and hash3 should be the same")
+	}
+	if hash1[1] == hash3[1] {
+		t.Fatal("the second bytes of hash1 and hash3 should be different")
+	}
+
+	// add only proto1
+	tree.AddProtocol(proto1)
+
+	if tree.root == nil {
+		t.Fatal("root should not be nil after adding protocol")
+	}
+	if tree.root.children[hash1[0]] == nil || tree.root.children[hash1[0]].p == nil {
+		t.Fatal("the protocol was not added")
+	}
+	if tree.root.children[hash1[0]].p.protocolID != proto1 {
+		t.Fatal("the protocol ID was wrong")
+	}
+	if tree.root.children[hash1[0]].p.tombstoneBit {
+		t.Fatal("tombstoneBit of proto1 must not be set")
+	}
+	if !bytes.Equal(tree.Abbreviate(proto1), []byte{hash1[0]}) {
+		t.Fatal("abbreviation of proto1 is incorrect")
+	}
+
+	// also add proto2
+	tree.AddProtocol(proto2)
+
+	if tree.root.children[hash2[0]] == nil || tree.root.children[hash2[0]].p == nil {
+		t.Fatal("the protocol was not added")
+	}
+	if tree.root.children[hash2[0]].p.protocolID != proto2 {
+		t.Fatal("the protocol ID was wrong")
+	}
+	if tree.root.children[hash2[0]].p.tombstoneBit {
+		t.Fatal("tombstoneBit of proto2 must not be set")
+	}
+	if !bytes.Equal(tree.Abbreviate(proto2), []byte{hash2[0]}) {
+		t.Fatal("abbreviation of proto2 is incorrect")
+	}
+
+	// add proto3 which has the same first byte of the hash as proto1
+	tree.AddProtocol(proto3)
+
+	n1 := tree.root.children[hash1[0]]
+	// the node at the first level should still be proto1
+	if n1.p.protocolID != proto1 {
+		t.Fatal("the node in the first level should not be modified")
+	}
+	// proto1 should be duplicated down
+	if n1.children[hash1[1]] == nil || n1.children[hash1[1]].p == nil {
+		t.Fatal("proto1 was not duplicated")
+	}
+	if n1.children[hash1[1]].p.tombstoneBit {
+		t.Fatal("tombstoneBit of proto1 must not be set")
+	}
+	if !bytes.Equal(tree.Abbreviate(proto1), []byte{hash1[0], hash1[1]}) {
+		t.Fatal("abbreviation of proto1 is incorrect")
+	}
+	// proto3 should be added in the second level
+	if n1.children[hash3[1]] == nil || n1.children[hash3[1]].p == nil {
+		t.Fatal("proto3 was not added")
+	}
+	if n1.children[hash3[1]].p.tombstoneBit {
+		t.Fatal("tombstoneBit of proto3 must not be set")
+	}
+	if !bytes.Equal(tree.Abbreviate(proto3), []byte{hash3[0], hash3[1]}) {
+		t.Fatal("abbreviation of proto3 is incorrect")
+	}
+}
+
+func TestAbbrevTreeRemoveProtocol(t *testing.T) {
+	tree := &abbrevTree[string]{}
+
+	proto1 := "protocol1"
+	hash1 := sha256.Sum256([]byte(proto1))
+	proto2 := "protocol2"
+	proto3 := "protocol251" // this one has the same first byte as "protocol1"
+
+	tree.AddProtocol(proto1)
+	tree.AddProtocol(proto2)
+	tree.AddProtocol(proto3)
+
+	// remove only proto1
+	tree.RemoveProtocol(proto1)
+
+	n1 := tree.root.children[hash1[0]]
+	if !n1.p.tombstoneBit {
+		t.Fatal("tombstoneBit of proto1 must be set")
+	}
+	if !n1.children[hash1[1]].p.tombstoneBit {
+		t.Fatal("tombstoneBit of proto1 must be set")
+	}
+	if tree.Abbreviate(proto1) != nil {
+		t.Fatal("abbreviation of proto1 should be nil")
+	}
+}
+
+func TestAbbrevTreeResurrectProtocol(t *testing.T) {
+	tree := &abbrevTree[string]{}
+
+	proto1 := "protocol1"
+	hash1 := sha256.Sum256([]byte(proto1))
+	proto2 := "protocol2"
+	proto3 := "protocol251" // this one has the same first byte as "protocol1"
+
+	tree.AddProtocol(proto1)
+	tree.AddProtocol(proto2)
+	tree.AddProtocol(proto3)
+	tree.RemoveProtocol(proto1)
+	tree.AddProtocol(proto1)
+
+	n1 := tree.root.children[hash1[0]]
+	if n1.p.tombstoneBit {
+		t.Fatal("tombstoneBit of proto1 must not be set")
+	}
+	if n1.children[hash1[1]].p.tombstoneBit {
+		t.Fatal("tombstoneBit of proto1 must not be set")
+	}
+
+	// There should be another leaf node added for proto1
+	n2 := n1.children[hash1[1]]
+	if n2.children[hash1[2]] == nil || n2.children[hash1[2]].p == nil {
+		t.Fatal("proto1 was not added")
+	}
+	if n2.children[hash1[2]].p.tombstoneBit {
+		t.Fatal("tombstoneBit of proto1 must not be set")
+	}
+	if !bytes.Equal(tree.Abbreviate(proto1), []byte{hash1[0], hash1[1], hash1[2]}) {
+		t.Fatal("abbreviation of proto1 is incorrect")
+	}
+}

lazyClient.go

@@ -1,6 +1,7 @@
 package multistream
 
 import (
+	"encoding/hex"
 	"fmt"
 	"io"
 )
@@ -18,9 +19,15 @@ func NewMSSelect[T StringLike](c io.ReadWriteCloser, proto T) LazyConn {
 }
 
 func NewMSSelect2[T StringLike](c io.ReadWriteCloser, proto T, peerProtos []T) LazyConn {
-	// TODO: put peerProtos into lazyClientConn so that it knows what protocols the other peer supports
+	t := &abbrevTree[T]{}
+	for _, p := range peerProtos {
+		t.AddProtocol(p)
+	}
+
+	// TODO: use a proper varint instead of a hex string later
+	abbrv := T(hex.EncodeToString(t.Abbreviate(proto)))
 	return &lazyClientConn[T]{
-		protos: []T{ProtocolID, proto},
+		protos: []T{ProtocolID, abbrv},
 		con:    c,
 
 		rhandshakeOnce: newOnce(),

multistream.go

@@ -22,6 +22,9 @@ var ErrTooLarge = errors.New("incoming message was too large")
 // the multistream muxers on both sides of a channel can work with each other.
 const ProtocolID = "/multistream/1.0.0"
 
+// Multistream-select version that protocol abbreviation is supported
+const AbbrevSupportedMSSVersion = 2
+
 var writerPool = sync.Pool{
 	New: func() interface{} {
 		return bufio.NewWriter(nil)