From 0d8a1852cfd31feb6f8d951ad0a8df258f9ae4d7 Mon Sep 17 00:00:00 2001
From: vyzo <vyzo@hackzen.org>
Date: Thu, 16 Sep 2021 14:06:52 +0300
Subject: [PATCH] make tag value fixed, tagged in first circuit, untagged in
 last.

---
 p2p/protocol/circuitv1/relay/relay.go | 13 +++++--------
 p2p/protocol/circuitv2/relay/relay.go | 12 +++++-------
 2 files changed, 10 insertions(+), 15 deletions(-)

diff --git a/p2p/protocol/circuitv1/relay/relay.go b/p2p/protocol/circuitv1/relay/relay.go
index 35ebaae9..11102617 100644
--- a/p2p/protocol/circuitv1/relay/relay.go
+++ b/p2p/protocol/circuitv1/relay/relay.go
@@ -29,8 +29,8 @@ const (
 	ConnectTimeout   = 30 * time.Second
 	HandshakeTimeout = time.Minute
 
-	RelayHopTag    = "relay-hop"
-	MaxRelayHopTag = 5
+	RelayHopTag      = "relay-v1-hop"
+	RelayHopTagValue = 2
 
 	maxMessageSize = 4096
 )
@@ -269,8 +269,8 @@ func (r *Relay) addConn(p peer.ID) {
 	conns := r.conns[p]
 	conns++
 	r.conns[p] = conns
-	if conns < MaxRelayHopTag {
-		r.host.ConnManager().UpsertTag(p, RelayHopTag, func(v int) int { return v + 1 })
+	if conns == 1 {
+		r.host.ConnManager().TagPeer(p, RelayHopTag, RelayHopTagValue)
 	}
 }
 
@@ -281,11 +281,8 @@ func (r *Relay) rmConn(p peer.ID) {
 		r.conns[p] = conns
 	} else {
 		delete(r.conns, p)
+		r.host.ConnManager().UntagPeer(p, RelayHopTag)
 	}
-	if conns <= MaxRelayHopTag {
-		r.host.ConnManager().UpsertTag(p, RelayHopTag, func(v int) int { return v - 1 })
-	}
-
 }
 
 func (r *Relay) relayConn(src, dest network.Stream, srcID, destID peer.ID, done func()) {
diff --git a/p2p/protocol/circuitv2/relay/relay.go b/p2p/protocol/circuitv2/relay/relay.go
index 5ea569d9..a50c32fd 100644
--- a/p2p/protocol/circuitv2/relay/relay.go
+++ b/p2p/protocol/circuitv2/relay/relay.go
@@ -30,8 +30,8 @@ const (
 	ConnectTimeout   = 30 * time.Second
 	HandshakeTimeout = time.Minute
 
-	RelayHopTag    = "relay-hop"
-	MaxRelayHopTag = 5
+	RelayHopTag      = "relay-v2-hop"
+	RelayHopTagValue = 2
 
 	maxMessageSize = 4096
 )
@@ -346,8 +346,8 @@ func (r *Relay) addConn(p peer.ID) {
 	conns := r.conns[p]
 	conns++
 	r.conns[p] = conns
-	if conns <= MaxRelayHopTag {
-		r.host.ConnManager().UpsertTag(p, RelayHopTag, func(v int) int { return v + 1 })
+	if conns == 1 {
+		r.host.ConnManager().TagPeer(p, RelayHopTag, RelayHopTagValue)
 	}
 }
 
@@ -358,9 +358,7 @@ func (r *Relay) rmConn(p peer.ID) {
 		r.conns[p] = conns
 	} else {
 		delete(r.conns, p)
-	}
-	if conns < MaxRelayHopTag {
-		r.host.ConnManager().UpsertTag(p, RelayHopTag, func(v int) int { return v - 1 })
+		r.host.ConnManager().UntagPeer(p, RelayHopTag)
 	}
 }