p2p: support protocol version negotiation

3 files changed, 110 insertions, 6 deletions

diff --git a/p2p/peer.go b/p2p/peer.go
index 40466cf84..e1bda1d03 100644
--- a/p2p/peer.go
+++ b/p2p/peer.go
@@ -249,15 +249,22 @@ func countMatchingProtocols(protocols []Protocol, caps []Cap) int {
 
 // matchProtocols creates structures for matching named subprotocols.
 func matchProtocols(protocols []Protocol, caps []Cap, rw MsgReadWriter) map[string]*protoRW {
-	sort.Sort(capsByName(caps))
+	sort.Sort(capsByNameAndVersion(caps))
 	offset := baseProtocolLength
 	result := make(map[string]*protoRW)
+
 outer:
 	for _, cap := range caps {
 		for _, proto := range protocols {
-			if proto.Name == cap.Name && proto.Version == cap.Version && result[cap.Name] == nil {
+			if proto.Name == cap.Name && proto.Version == cap.Version {
+				// If an old protocol version matched, revert it
+				if old := result[cap.Name]; old != nil {
+					offset -= old.Length
+				}
+				// Assign the new match
 				result[cap.Name] = &protoRW{Protocol: proto, offset: offset, in: make(chan Msg), w: rw}
 				offset += proto.Length
+
 				continue outer
 			}
 		}
diff --git a/p2p/peer_test.go b/p2p/peer_test.go
index 575d0ff79..95167f4a1 100644
--- a/p2p/peer_test.go
+++ b/p2p/peer_test.go
@@ -196,3 +196,98 @@ func TestNewPeer(t *testing.T) {
 
 	p.Disconnect(DiscAlreadyConnected) // Should not hang
 }
+
+func TestMatchProtocols(t *testing.T) {
+	tests := []struct {
+		Local  []Cap
+		Remote []Protocol
+		Match  map[string]protoRW
+	}{
+		{
+			// No remote protocols
+			Local: []Cap{{Name: "a"}},
+		},
+		{
+			// No local capabilities
+			Remote: []Protocol{{Name: "a"}},
+		},
+		{
+			// No mutual protocols
+			Local:  []Cap{{Name: "a"}},
+			Remote: []Protocol{{Name: "b"}},
+		},
+		{
+			// Some matches, some differences
+			Local:  []Cap{{Name: "local"}, {Name: "match1"}, {Name: "match2"}},
+			Remote: []Protocol{{Name: "match1"}, {Name: "match2"}, {Name: "remote"}},
+			Match:  map[string]protoRW{"match1": {Protocol: Protocol{Name: "match1"}}, "match2": {Protocol: Protocol{Name: "match2"}}},
+		},
+		{
+			// Various alphabetical ordering
+			Local:  []Cap{{Name: "aa"}, {Name: "ab"}, {Name: "bb"}, {Name: "ba"}},
+			Remote: []Protocol{{Name: "ba"}, {Name: "bb"}, {Name: "ab"}, {Name: "aa"}},
+			Match:  map[string]protoRW{"aa": {Protocol: Protocol{Name: "aa"}}, "ab": {Protocol: Protocol{Name: "ab"}}, "ba": {Protocol: Protocol{Name: "ba"}}, "bb": {Protocol: Protocol{Name: "bb"}}},
+		},
+		{
+			// No mutual versions
+			Local:  []Cap{{Version: 1}},
+			Remote: []Protocol{{Version: 2}},
+		},
+		{
+			// Multiple versions, single common
+			Local:  []Cap{{Version: 1}, {Version: 2}},
+			Remote: []Protocol{{Version: 2}, {Version: 3}},
+			Match:  map[string]protoRW{"": {Protocol: Protocol{Version: 2}}},
+		},
+		{
+			// Multiple versions, multiple common
+			Local:  []Cap{{Version: 1}, {Version: 2}, {Version: 3}, {Version: 4}},
+			Remote: []Protocol{{Version: 2}, {Version: 3}},
+			Match:  map[string]protoRW{"": {Protocol: Protocol{Version: 3}}},
+		},
+		{
+			// Various version orderings
+			Local:  []Cap{{Version: 4}, {Version: 1}, {Version: 3}, {Version: 2}},
+			Remote: []Protocol{{Version: 2}, {Version: 3}, {Version: 1}},
+			Match:  map[string]protoRW{"": {Protocol: Protocol{Version: 3}}},
+		},
+		{
+			// Versions overriding sub-protocol lengths
+			Local:  []Cap{{Version: 1}, {Version: 2}, {Version: 3}, {Name: "a"}},
+			Remote: []Protocol{{Version: 1, Length: 1}, {Version: 2, Length: 2}, {Version: 3, Length: 3}, {Name: "a"}},
+			Match:  map[string]protoRW{"": {Protocol: Protocol{Version: 3}}, "a": {Protocol: Protocol{Name: "a"}, offset: 3}},
+		},
+	}
+
+	for i, tt := range tests {
+		result := matchProtocols(tt.Remote, tt.Local, nil)
+		if len(result) != len(tt.Match) {
+			t.Errorf("test %d: negotiation mismatch: have %v, want %v", i, len(result), len(tt.Match))
+			continue
+		}
+		// Make sure all negotiated protocols are needed and correct
+		for name, proto := range result {
+			match, ok := tt.Match[name]
+			if !ok {
+				t.Errorf("test %d, proto '%s': negotiated but shouldn't have", i, name)
+				continue
+			}
+			if proto.Name != match.Name {
+				t.Errorf("test %d, proto '%s': name mismatch: have %v, want %v", i, name, proto.Name, match.Name)
+			}
+			if proto.Version != match.Version {
+				t.Errorf("test %d, proto '%s': version mismatch: have %v, want %v", i, name, proto.Version, match.Version)
+			}
+			if proto.offset-baseProtocolLength != match.offset {
+				t.Errorf("test %d, proto '%s': offset mismatch: have %v, want %v", i, name, proto.offset-baseProtocolLength, match.offset)
+			}
+		}
+		// Make sure no protocols missed negotiation
+		for name, _ := range tt.Match {
+			if _, ok := result[name]; !ok {
+				t.Errorf("test %d, proto '%s': not negotiated, should have", i, name)
+				continue
+			}
+		}
+	}
+}
diff --git a/p2p/protocol.go b/p2p/protocol.go
index 5fa395eda..a229ba911 100644
--- a/p2p/protocol.go
+++ b/p2p/protocol.go
@@ -43,8 +43,10 @@ func (cap Cap) String() string {
 	return fmt.Sprintf("%s/%d", cap.Name, cap.Version)
 }
 
-type capsByName []Cap
+type capsByNameAndVersion []Cap
 
-func (cs capsByName) Len() int           { return len(cs) }
-func (cs capsByName) Less(i, j int) bool { return cs[i].Name < cs[j].Name }
-func (cs capsByName) Swap(i, j int)      { cs[i], cs[j] = cs[j], cs[i] }
+func (cs capsByNameAndVersion) Len() int      { return len(cs) }
+func (cs capsByNameAndVersion) Swap(i, j int) { cs[i], cs[j] = cs[j], cs[i] }
+func (cs capsByNameAndVersion) Less(i, j int) bool {
+	return cs[i].Name < cs[j].Name || (cs[i].Name == cs[j].Name && cs[i].Version < cs[j].Version)
+}