From 25abbc379421f4272e3e9986fd34b5a1028a35f8 Mon Sep 17 00:00:00 2001
From: Louis Thibault <l.thibault@sentimens.com>
Date: Fri, 18 Aug 2023 19:36:56 -0400
Subject: [PATCH] Pass boot socket directly to surveyor

---
 boot/boot.go               |  3 ++-
 boot/survey/survey.go      | 12 ++----------
 boot/survey/survey_test.go |  9 ++++++---
 3 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/boot/boot.go b/boot/boot.go
index a8c44b6a..d6676c9d 100644
--- a/boot/boot.go
+++ b/boot/boot.go
@@ -106,7 +106,8 @@ func DialMulticast(h host.Host, maddr ma.Multiaddr, opt ...socket.Option) (*surv
 		return nil, err
 	}
 
-	return survey.New(h, conn, opt...), nil
+	sock := socket.New(conn, withDefault(h, opt)...)
+	return survey.New(h, sock), nil
 }
 
 func DialPortRange(h host.Host, maddr ma.Multiaddr, opt ...socket.Option) (*crawl.Crawler, error) {
diff --git a/boot/survey/survey.go b/boot/survey/survey.go
index 4c5ae8c4..3ea0d7c9 100644
--- a/boot/survey/survey.go
+++ b/boot/survey/survey.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"encoding/binary"
 	"errors"
-	"net"
 	"time"
 
 	"github.com/libp2p/go-libp2p/core/crypto"
@@ -27,10 +26,10 @@ type Surveyor struct {
 
 // New surveyor.  The supplied PacketConn SHOULD be bound to a multicast
 // group.  Use of JoinMulticastGroup to construct conn is RECOMMENDED.
-func New(h host.Host, conn net.PacketConn, opt ...socket.Option) *Surveyor {
+func New(h host.Host, sock *socket.Socket) *Surveyor {
 	s := &Surveyor{
 		host: h,
-		sock: socket.New(conn, withDefault(h, opt)...),
+		sock: sock,
 	}
 
 	go s.sock.Bind(s.handler())
@@ -38,13 +37,6 @@ func New(h host.Host, conn net.PacketConn, opt ...socket.Option) *Surveyor {
 	return s
 }
 
-func withDefault(h host.Host, opt []socket.Option) []socket.Option {
-	return append([]socket.Option{
-		socket.WithRateLimiter(socket.NewPacketLimiter(16, 8)),
-		socket.WithValidator(socket.BasicValidator(h.ID())),
-	}, opt...)
-}
-
 func (s *Surveyor) Close() error {
 	return s.sock.Close()
 }
diff --git a/boot/survey/survey_test.go b/boot/survey/survey_test.go
index 6ccca6a8..37bf460d 100644
--- a/boot/survey/survey_test.go
+++ b/boot/survey/survey_test.go
@@ -39,7 +39,8 @@ func TestClose(t *testing.T) {
 		ReadFrom(gomock.Any()).
 		AnyTimes()
 
-	err := survey.New(h, conn).Close()
+	sock := socket.New(conn)
+	err := survey.New(h, sock).Close()
 	assert.NoError(t, err, "surveyor should close gracefully")
 }
 
@@ -90,7 +91,8 @@ func TestAdvertise(t *testing.T) {
 		Return(addr).
 		AnyTimes()
 
-	surveyor := survey.New(tt[0].h, conn)
+	sock := socket.New(conn)
+	surveyor := survey.New(tt[0].h, sock)
 	defer surveyor.Close()
 
 	surveyor.Advertise(context.Background(), ns)
@@ -145,7 +147,8 @@ func TestFindPeers(t *testing.T) {
 		Close().
 		Times(1)
 
-	surveyor := survey.New(tt[0].h, conn)
+	sock := socket.New(conn)
+	surveyor := survey.New(tt[0].h, sock)
 	defer surveyor.Close()
 
 	peers, err := surveyor.FindPeers(context.Background(), ns, discovery.Limit(1))