Merge pull request #1188 from libp2p/feat/relay-tags

Tag relay hops in relay implementations
libp2p · Sep 21, 2021 · 8a70f70 · 8a70f70
2 parents d1e2681 + 9674b5c
commit 8a70f70
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 9 deletions.
diff --git a/p2p/protocol/circuitv1/relay/relay.go b/p2p/protocol/circuitv1/relay/relay.go
@@ -29,6 +29,9 @@ const (
  ConnectTimeout = 30 * time.Second
  HandshakeTimeout = time.Minute
 
+ relayHopTag = "relay-v1-hop"
+ relayHopTagValue = 2
+
  maxMessageSize = 4096
 )
 
@@ -158,15 +161,15 @@ func (r *Relay) handleHopStream(s network.Stream, msg *pb.CircuitRelay) {
  }
 
  r.active++
- r.conns[src.ID]++
- r.conns[dest.ID]++
+ r.addConn(src.ID)
+ r.addConn(src.ID)
  r.mx.Unlock()
 
  cleanup := func() {
  r.mx.Lock()
  r.active--
- r.conns[src.ID]--
- r.conns[dest.ID]--
+ r.rmConn(src.ID)
+ r.rmConn(dest.ID)
  r.mx.Unlock()
  }
 
@@ -262,6 +265,26 @@ func (r *Relay) handleHopStream(s network.Stream, msg *pb.CircuitRelay) {
  go r.relayConn(bs, s, dest.ID, src.ID, done)
 }
 
+func (r *Relay) addConn(p peer.ID) {
+ conns := r.conns[p]
+ conns++
+ r.conns[p] = conns
+ if conns == 1 {
+ r.host.ConnManager().TagPeer(p, relayHopTag, relayHopTagValue)
+ }
+}
+
+func (r *Relay) rmConn(p peer.ID) {
+ conns := r.conns[p]
+ conns--
+ if conns > 0 {
+ r.conns[p] = conns
+ } else {
+ delete(r.conns, p)
+ r.host.ConnManager().UntagPeer(p, relayHopTag)
+ }
+}
+
 func (r *Relay) relayConn(src, dest network.Stream, srcID, destID peer.ID, done func()) {
  defer done()
 

diff --git a/p2p/protocol/circuitv2/relay/relay.go b/p2p/protocol/circuitv2/relay/relay.go
@@ -30,6 +30,9 @@ const (
  ConnectTimeout = 30 * time.Second
  HandshakeTimeout = time.Minute
 
+ relayHopTag = "relay-v2-hop"
+ relayHopTagValue = 2
+
  maxMessageSize = 4096
 )
 
@@ -215,23 +218,23 @@ func (r *Relay) handleConnect(s network.Stream, msg *pbv2.HopMessage) {
  r.handleError(s, pbv2.Status_RESOURCE_LIMIT_EXCEEDED)
  return
  }
- r.conns[src]++
 
  destConns := r.conns[dest.ID]
  if destConns >= r.rc.MaxCircuits {
- r.conns[src]--
  r.mx.Unlock()
  log.Debugf("refusing connection from %s to %s; too many connecitons to %s", src, dest.ID, dest.ID)
  r.handleError(s, pbv2.Status_RESOURCE_LIMIT_EXCEEDED)
  return
  }
- r.conns[dest.ID]++
+
+ r.addConn(src)
+ r.addConn(dest.ID)
  r.mx.Unlock()
 
  cleanup := func() {
  r.mx.Lock()
- r.conns[src]--
- r.conns[dest.ID]--
+ r.rmConn(src)
+ r.rmConn(dest.ID)
  r.mx.Unlock()
  }
 
@@ -339,6 +342,26 @@ func (r *Relay) handleConnect(s network.Stream, msg *pbv2.HopMessage) {
  }
 }
 
+func (r *Relay) addConn(p peer.ID) {
+ conns := r.conns[p]
+ conns++
+ r.conns[p] = conns
+ if conns == 1 {
+ r.host.ConnManager().TagPeer(p, relayHopTag, relayHopTagValue)
+ }
+}
+
+func (r *Relay) rmConn(p peer.ID) {
+ conns := r.conns[p]
+ conns--
+ if conns > 0 {
+ r.conns[p] = conns
+ } else {
+ delete(r.conns, p)
+ r.host.ConnManager().UntagPeer(p, relayHopTag)
+ }
+}
+
 func (r *Relay) relayLimited(src, dest network.Stream, srcID, destID peer.ID, limit int64, done func()) {
  defer done()