Merge pull request #77 from libp2p/feat/stop-tags

Tag the hop relay when creating stop streams

conn.go

@@ -5,6 +5,7 @@ import (
 	"net"
 	"time"
 
+	host "github.com/libp2p/go-libp2p-host"
 	inet "github.com/libp2p/go-libp2p-net"
 	pstore "github.com/libp2p/go-libp2p-peerstore"
 	ma "github.com/multiformats/go-multiaddr"
@@ -14,6 +15,7 @@ import (
 type Conn struct {
 	stream inet.Stream
 	remote pstore.PeerInfo
+	host   host.Host
 }
 
 type NetAddr struct {
@@ -30,6 +32,7 @@ func (n *NetAddr) String() string {
 }
 
 func (c *Conn) Close() error {
+	c.untagHop()
 	return c.stream.Reset()
 }
 
@@ -60,6 +63,20 @@ func (c *Conn) RemoteAddr() net.Addr {
 	}
 }
 
+// Increment the underlying relay connection tag by 1, thus increasing its protection from
+// connection pruning. This ensures that connections to relays are not accidentally closed,
+// by the connection manager, taking with them all the relayed connections (that may themselves
+// be protected).
+func (c *Conn) tagHop() {
+	c.host.ConnManager().UpsertTag(c.stream.Conn().RemotePeer(), "relay-hop-stream", incrementTag)
+}
+
+// Decrement the underlying relay connection tag by 1; this is performed when we close the
+// relayed connection.
+func (c *Conn) untagHop() {
+	c.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

@@ -116,16 +116,21 @@ func NewRelay(ctx context.Context, h host.Host, upgrader *tptu.Upgrader, opts ..
 	return r, nil
 }
 
+// Increment the live hop count and increment the connection manager tags by 1 for the two
+// sides of the hop stream. This ensures that connections with many hop streams will be protected
+// from pruning, thus minimizing disruption from connection trimming in a relay node.
 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)
 }
 
+// Decrement the live hpo count and decrement the connection manager tags for the two sides
+// of the hop stream.
 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 +185,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, host: r.host}, nil
 }
 
 func (r *Relay) Matches(addr ma.Multiaddr) bool {
@@ -438,7 +443,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, host: r.host}:
 	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