mirror of
https://github.com/libp2p/go-libp2p.git
synced 2025-09-26 20:21:26 +08:00
feat: relay: add option for custom filter function
This commit is contained in:
@@ -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 {
|
||||
|
@@ -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
|
||||
}
|
||||
|
||||
|
@@ -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)
|
||||
}
|
||||
|
Reference in New Issue
Block a user