feat: relay: add option for custom filter function

Author: shoriwe

p2p/protocol/circuitv2/relay/options.go

@@ -1,5 +1,9 @@
 package relay
 
+import (
+	"github.com/multiformats/go-multiaddr"
+)
+
 type Option func(*Relay) error
 
 // WithResources is a Relay option that sets specific relay resources for the relay.
@@ -18,6 +22,18 @@ func WithLimit(limit *RelayLimit) Option {
 	}
 }
 
+// Reservation address function used to promote addresses to connected nodes
+type ReservationAddressFilterFunc func(addr multiaddr.Multiaddr) (include bool)
+
+// Overrides the default reservation address filter.
+// This will permit the relay let the client know it have access to non public addresses too.
+func WithReservationAddressFilter(filter ReservationAddressFilterFunc) (option Option) {
+	return func(r *Relay) (err error) {
+		r.reservationAddrFilter = filter
+		return nil
+	}
+}
+
 // WithInfiniteLimits is a Relay option that disables limits.
 func WithInfiniteLimits() Option {
 	return func(r *Relay) error {

p2p/protocol/circuitv2/relay/relay.go

@@ -46,6 +46,8 @@ type Relay struct {
 	ctx    context.Context
 	cancel func()
 
+	reservationAddrFilter ReservationAddressFilterFunc
+
 	host        host.Host
 	rc          Resources
 	acl         ACLFilter
@@ -75,6 +77,8 @@ func New(h host.Host, opts ...Option) (*Relay, error) {
 		acl:    nil,
 		rsvp:   make(map[peer.ID]time.Time),
 		conns:  make(map[peer.ID]int),
+
+		reservationAddrFilter: manet.IsPublicAddr,
 	}
 
 	for _, opt := range opts {
@@ -235,6 +239,7 @@ func (r *Relay) handleReserve(s network.Stream) pbv2.Status {
 	// For example, the stream might be reset or the connection might be closed before the reservation is received.
 	// In that case, the reservation will just be garbage collected later.
 	rsvp := makeReservationMsg(
+		r.reservationAddrFilter,
 		r.host.Peerstore().PrivKey(r.host.ID()),
 		r.host.ID(),
 		r.host.Addrs(),
@@ -612,6 +617,7 @@ func (r *Relay) writeResponse(s network.Stream, status pbv2.Status, rsvp *pbv2.R
 }
 
 func makeReservationMsg(
+	reservationAddrFilter ReservationAddressFilterFunc,
 	signingKey crypto.PrivKey,
 	selfID peer.ID,
 	selfAddrs []ma.Multiaddr,
@@ -630,7 +636,7 @@ func makeReservationMsg(
 
 	addrBytes := make([][]byte, 0, len(selfAddrs))
 	for _, addr := range selfAddrs {
-		if !manet.IsPublicAddr(addr) {
+		if !reservationAddrFilter(addr) {
 			continue
 		}
 

p2p/protocol/circuitv2/relay/relay_priv_test.go

@@ -10,6 +10,8 @@ import (
 	"github.com/stretchr/testify/require"
 
 	ma "github.com/multiformats/go-multiaddr"
+	"github.com/multiformats/go-multiaddr/matest"
+	manet "github.com/multiformats/go-multiaddr/net"
 )
 
 func genKeyAndID(t *testing.T) (crypto.PrivKey, peer.ID) {
@@ -28,26 +30,49 @@ func TestMakeReservationWithP2PAddrs(t *testing.T) {
 	_, otherID := genKeyAndID(t)
 	_, reserverID := genKeyAndID(t)
 
-	addrs := []ma.Multiaddr{
-		ma.StringCast("/ip4/1.2.3.4/tcp/1234"),                         // No p2p part
-		ma.StringCast("/ip4/1.2.3.4/tcp/1235/p2p/" + selfID.String()),  // Already has p2p part
-		ma.StringCast("/ip4/1.2.3.4/tcp/1236/p2p/" + otherID.String()), // Some other peer (?? Not expected, but we could get anything in this func)
-	}
-
-	rsvp := makeReservationMsg(selfKey, selfID, addrs, reserverID, time.Now().Add(time.Minute))
-	require.NotNil(t, rsvp)
+	tcs := []struct {
+		name     string
+		filter   func(ma.Multiaddr) bool
+		input    []ma.Multiaddr
+		expected []ma.Multiaddr
+	}{{
+		name:   "only public",
+		filter: manet.IsPublicAddr,
+		input: []ma.Multiaddr{
+			ma.StringCast("/ip4/1.2.3.4/tcp/1234"),                            // No p2p part
+			ma.StringCast("/ip4/1.2.3.4/tcp/1235/p2p/" + selfID.String()),     // Already has p2p part
+			ma.StringCast("/ip4/192.168.1.9/tcp/1235/p2p/" + selfID.String()), // Already has p2p part
+			ma.StringCast("/ip4/1.2.3.4/tcp/1236/p2p/" + otherID.String()),    // Some other peer (?? Not expected, but we could get anything in this func)
+		},
+		expected: []ma.Multiaddr{
+			ma.StringCast("/ip4/1.2.3.4/tcp/1234/p2p/" + selfID.String()),
+			ma.StringCast("/ip4/1.2.3.4/tcp/1235/p2p/" + selfID.String()),
+		},
+	}, {
+		name:   "only not public",
+		filter: func(m ma.Multiaddr) bool { return !manet.IsPublicAddr(m) },
+		input: []ma.Multiaddr{
+			ma.StringCast("/ip4/1.2.3.4/tcp/1234"),                            // No p2p part
+			ma.StringCast("/ip4/1.2.3.4/tcp/1235/p2p/" + selfID.String()),     // Already has p2p part
+			ma.StringCast("/ip4/192.168.1.9/tcp/1235/p2p/" + selfID.String()), // Already has p2p part
+			ma.StringCast("/ip4/1.2.3.4/tcp/1236/p2p/" + otherID.String()),    // Some other peer (?? Not expected, but we could get anything in this func)
+		},
+		expected: []ma.Multiaddr{
+			ma.StringCast("/ip4/192.168.1.9/tcp/1235/p2p/" + selfID.String()),
+		},
+	}}
+	for _, tc := range tcs {
+		t.Run(tc.name, func(t *testing.T) {
+			rsvp := makeReservationMsg(tc.filter, selfKey, selfID, tc.input, reserverID, time.Now().Add(time.Minute))
+			require.NotNil(t, rsvp)
 
-	expectedAddrs := []string{
-		"/ip4/1.2.3.4/tcp/1234/p2p/" + selfID.String(),
-		"/ip4/1.2.3.4/tcp/1235/p2p/" + selfID.String(),
+			addrsFromRsvp := make([]ma.Multiaddr, 0, len(rsvp.GetAddrs()))
+			for _, addr := range rsvp.GetAddrs() {
+				a, err := ma.NewMultiaddrBytes(addr)
+				require.NoError(t, err)
+				addrsFromRsvp = append(addrsFromRsvp, a)
+			}
+			matest.AssertEqualMultiaddrs(t, tc.expected, addrsFromRsvp)
+		})
 	}
-
-	addrsFromRsvp := make([]string, 0, len(rsvp.GetAddrs()))
-	for _, addr := range rsvp.GetAddrs() {
-		a, err := ma.NewMultiaddrBytes(addr)
-		require.NoError(t, err)
-		addrsFromRsvp = append(addrsFromRsvp, a.String())
-	}
-
-	require.Equal(t, expectedAddrs, addrsFromRsvp)
 }