From 2166786b0ead0fda4f037568aea7c0a06997c3a2 Mon Sep 17 00:00:00 2001
From: Andrea Maria Piana <andrea.maria.piana@gmail.com>
Date: Tue, 29 Nov 2022 12:16:01 +0000
Subject: [PATCH] Setup autorelay

---
 VERSION        |  2 +-
 wakuv2/waku.go | 51 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/VERSION b/VERSION
index efdc3890a..55e6414ce 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-0.115.3
+0.115.4
diff --git a/wakuv2/waku.go b/wakuv2/waku.go
index d628bfe11..042a86c79 100644
--- a/wakuv2/waku.go
+++ b/wakuv2/waku.go
@@ -25,6 +25,7 @@ import (
 	"database/sql"
 	"errors"
 	"fmt"
+	"math/rand"
 	"net"
 	"runtime"
 	"strings"
@@ -32,6 +33,8 @@ import (
 	"time"
 
 	"github.com/libp2p/go-libp2p/core/peer"
+	"github.com/libp2p/go-libp2p/p2p/host/autorelay"
+	"github.com/libp2p/go-libp2p/p2p/protocol/circuitv2/proto"
 	"github.com/libp2p/go-libp2p/p2p/protocol/identify"
 	"github.com/multiformats/go-multiaddr"
 
@@ -70,6 +73,7 @@ import (
 
 const messageQueueLimit = 1024
 const requestTimeout = 30 * time.Second
+const autoRelayMinInterval = 2 * time.Second
 
 type settings struct {
 	LightClient         bool   // Indicates if the node is a light client
@@ -190,6 +194,10 @@ func New(nodeKey string, cfg *Config, logger *zap.Logger, appDB *sql.DB) (*Waku,
 	libp2pOpts := node.DefaultLibP2POptions
 	libp2pOpts = append(libp2pOpts, libp2p.BandwidthReporter(waku.bandwidthCounter))
 	libp2pOpts = append(libp2pOpts, libp2p.NATPortMap())
+	libp2pOpts = append(libp2pOpts, libp2p.EnableHolePunching())
+	libp2pOpts = append(libp2pOpts, libp2p.EnableAutoRelay(
+		autorelay.WithPeerSource(waku.autoRelayPeerSource, autoRelayMinInterval),
+	))
 
 	opts := []node.WakuNodeOption{
 		node.WithLibP2POptions(libp2pOpts...),
@@ -1173,6 +1181,49 @@ func (w *Waku) AddStorePeer(address string) (string, error) {
 	return string(*peerID), nil
 }
 
+func (w *Waku) autoRelayPeerSource(ctx context.Context, numPeers int) <-chan peer.AddrInfo {
+
+	w.logger.Debug("auto-relay asking for peers", zap.Int("num-peers", numPeers))
+
+	output := make(chan peer.AddrInfo, numPeers)
+	go func() {
+		peers, err := w.node.Peers()
+		if err != nil {
+			w.logger.Error("failed to fetch peers", zap.Error(err))
+			close(output)
+		}
+
+		// Shuffle peers
+		rand.Seed(time.Now().UnixNano())
+		rand.Shuffle(len(peers), func(i, j int) { peers[i], peers[j] = peers[j], peers[i] })
+
+		for _, p := range peers {
+			info := w.node.Host().Peerstore().PeerInfo(p.ID)
+
+			supportedProtocols, err := w.node.Host().Peerstore().SupportsProtocols(p.ID, proto.ProtoIDv2Hop)
+			if err != nil {
+				w.logger.Error("could not check supported protocols", zap.Error(err))
+			}
+
+			if len(supportedProtocols) == 0 {
+				continue
+			}
+
+			select {
+			case <-ctx.Done():
+				w.logger.Debug("context done, auto-relay has enough peers")
+				close(output)
+
+			case output <- info:
+				w.logger.Debug("published auto-relay peer info", zap.Any("peer-id", p.ID))
+
+			}
+		}
+		close(output)
+	}()
+	return output
+}
+
 func (w *Waku) AddRelayPeer(address string) (string, error) {
 	addr, err := multiaddr.NewMultiaddr(address)
 	if err != nil {