Tag the hop relay when creating stop streams

conn.go

@@ -14,6 +14,7 @@ import (
 type Conn struct {
 	stream inet.Stream
 	remote pstore.PeerInfo
+	relay  *Relay
 }
 
 type NetAddr struct {
@@ -30,6 +31,7 @@ func (n *NetAddr) String() string {
 }
 
 func (c *Conn) Close() error {
+	c.untagHop()
 	return c.stream.Reset()
 }
 
@@ -60,6 +62,14 @@ func (c *Conn) RemoteAddr() net.Addr {
 	}
 }
 
+func (c *Conn) tagHop() {
+	c.relay.host.ConnManager().UpsertTag(c.stream.Conn().RemotePeer(), "relay-hop-stream", incrementTag)
+}
+
+func (c *Conn) untagHop() {
+	c.relay.host.ConnManager().UpsertTag(c.stream.Conn().RemotePeer(), "relay-hop-stream", decrementTag)
+}
+
 // TODO: is it okay to cast c.Conn().RemotePeer() into a multiaddr? might be "user input"
 func (c *Conn) RemoteMultiaddr() ma.Multiaddr {
 	proto := ma.ProtocolWithCode(ma.P_P2P).Name

dial.go

@@ -16,6 +16,7 @@ func (d *RelayTransport) Dial(ctx context.Context, a ma.Multiaddr, p peer.ID) (t
 	if err != nil {
 		return nil, err
 	}
+	c.tagHop()
 	return d.upgrader.UpgradeOutbound(ctx, d, c, p)
 }
 

listen.go

@@ -35,6 +35,7 @@ func (l *RelayListener) Accept() (manet.Conn, error) {
 		// TODO: Pretty print.
 		log.Infof("accepted relay connection: %q", c)
 
+		c.tagHop()
 		return c, nil
 	case <-l.ctx.Done():
 		return nil, l.ctx.Err()

relay.go

@@ -118,14 +118,14 @@ func NewRelay(ctx context.Context, h host.Host, upgrader *tptu.Upgrader, opts ..
 
 func (r *Relay) addLiveHop(from, to peer.ID) {
 	atomic.AddInt32(&r.liveHopCount, 1)
-	r.host.ConnManager().UpsertTag(from, "relay-hop-stream", func(v int) int { return v + 1 })
-	r.host.ConnManager().UpsertTag(to, "relay-hop-stream", func(v int) int { return v + 1 })
+	r.host.ConnManager().UpsertTag(from, "relay-hop-stream", incrementTag)
+	r.host.ConnManager().UpsertTag(to, "relay-hop-stream", incrementTag)
 }
 
 func (r *Relay) rmLiveHop(from, to peer.ID) {
 	atomic.AddInt32(&r.liveHopCount, -1)
-	r.host.ConnManager().UpsertTag(from, "relay-hop-stream", func(v int) int { return v - 1 })
-	r.host.ConnManager().UpsertTag(to, "relay-hop-stream", func(v int) int { return v - 1 })
+	r.host.ConnManager().UpsertTag(from, "relay-hop-stream", decrementTag)
+	r.host.ConnManager().UpsertTag(to, "relay-hop-stream", decrementTag)
 
 }
 
@@ -180,7 +180,7 @@ func (r *Relay) DialPeer(ctx context.Context, relay pstore.PeerInfo, dest pstore
 		return nil, RelayError{msg.GetCode()}
 	}
 
-	return &Conn{stream: s, remote: dest}, nil
+	return &Conn{stream: s, remote: dest, relay: r}, nil
 }
 
 func (r *Relay) Matches(addr ma.Multiaddr) bool {
@@ -438,7 +438,7 @@ func (r *Relay) handleStopStream(s inet.Stream, msg *pb.CircuitRelay) {
 	}
 
 	select {
-	case r.incoming <- &Conn{stream: s, remote: src}:
+	case r.incoming <- &Conn{stream: s, remote: src, relay: r}:
 	case <-time.After(RelayAcceptTimeout):
 		r.handleError(s, pb.CircuitRelay_STOP_RELAY_REFUSED)
 	}

util.go

@@ -49,6 +49,18 @@ func peerInfoToPeer(pi pstore.PeerInfo) *pb.CircuitRelay_Peer {
 	return p
 }
 
+func incrementTag(v int) int {
+	return v + 1
+}
+
+func decrementTag(v int) int {
+	if v > 0 {
+		return v - 1
+	} else {
+		return v
+	}
+}
+
 type delimitedReader struct {
 	r   io.Reader
 	buf []byte