feat: relay: add option for custom filter function

This commit is contained in:
Antonio
2025-08-08 19:36:33 -03:00
committed by Marco Munizaga
parent 72894e318a
commit 95cad5091f
3 changed files with 69 additions and 22 deletions

View File

@@ -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 {

View File

@@ -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
}

View File

@@ -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)
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)
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)
})
}
rsvp := makeReservationMsg(selfKey, selfID, addrs, 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([]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)
}