From 7c2b5a39b370ceb9276a01b53cf9fb966bfe9d5b Mon Sep 17 00:00:00 2001
From: Andrea Maria Piana <andrea.maria.piana@gmail.com>
Date: Fri, 24 Mar 2023 12:05:42 +0000
Subject: [PATCH] Use correct timesource in waku-2

---
 eth-node/bridge/geth/waku.go   |  5 -----
 eth-node/bridge/geth/wakuv2.go |  5 -----
 eth-node/types/waku.go         |  3 +--
 wakuv2/waku.go                 | 23 ++++++++---------------
 4 files changed, 9 insertions(+), 27 deletions(-)

diff --git a/eth-node/bridge/geth/waku.go b/eth-node/bridge/geth/waku.go
index a54a2d164..049675e0b 100644
--- a/eth-node/bridge/geth/waku.go
+++ b/eth-node/bridge/geth/waku.go
@@ -121,11 +121,6 @@ func (w *gethWakuWrapper) GetCurrentTime() time.Time {
 	return w.waku.CurrentTime()
 }
 
-// SetTimeSource assigns a particular source of time to a whisper object.
-func (w *gethWakuWrapper) SetTimeSource(timesource func() time.Time) {
-	w.waku.SetTimeSource(timesource)
-}
-
 func (w *gethWakuWrapper) SubscribeEnvelopeEvents(eventsProxy chan<- types.EnvelopeEvent) types.Subscription {
 	events := make(chan wakucommon.EnvelopeEvent, 100) // must be buffered to prevent blocking whisper
 	go func() {
diff --git a/eth-node/bridge/geth/wakuv2.go b/eth-node/bridge/geth/wakuv2.go
index f1660a63a..070e02315 100644
--- a/eth-node/bridge/geth/wakuv2.go
+++ b/eth-node/bridge/geth/wakuv2.go
@@ -69,11 +69,6 @@ func (w *gethWakuV2Wrapper) GetCurrentTime() time.Time {
 	return w.waku.CurrentTime()
 }
 
-// SetTimeSource assigns a particular source of time to a whisper object.
-func (w *gethWakuV2Wrapper) SetTimeSource(timesource func() time.Time) {
-	w.waku.SetTimeSource(timesource)
-}
-
 func (w *gethWakuV2Wrapper) SubscribeEnvelopeEvents(eventsProxy chan<- types.EnvelopeEvent) types.Subscription {
 	events := make(chan wakucommon.EnvelopeEvent, 100) // must be buffered to prevent blocking whisper
 	go func() {
diff --git a/eth-node/types/waku.go b/eth-node/types/waku.go
index 7eedc6c1e..70f0ee4ce 100644
--- a/eth-node/types/waku.go
+++ b/eth-node/types/waku.go
@@ -107,8 +107,7 @@ type Waku interface {
 	// If a message does not match the bloom, it will tantamount to spam, and the peer will
 	// be disconnected.
 	BloomFilter() []byte
-	// SetTimeSource assigns a particular source of time to a whisper object.
-	SetTimeSource(timesource func() time.Time)
+
 	// GetCurrentTime returns current time.
 	GetCurrentTime() time.Time
 
diff --git a/wakuv2/waku.go b/wakuv2/waku.go
index 17838549f..c68df5773 100644
--- a/wakuv2/waku.go
+++ b/wakuv2/waku.go
@@ -61,7 +61,6 @@ import (
 	"github.com/waku-org/go-waku/waku/v2/protocol/filter"
 	"github.com/waku-org/go-waku/waku/v2/protocol/peer_exchange"
 	"github.com/waku-org/go-waku/waku/v2/protocol/relay"
-	"github.com/waku-org/go-waku/waku/v2/utils"
 
 	"github.com/status-im/status-go/connection"
 	"github.com/status-im/status-go/eth-node/types"
@@ -132,8 +131,6 @@ type Waku struct {
 	connStatusSubscriptions map[string]*types.ConnStatusSubscription
 	connStatusMu            sync.Mutex
 
-	timeSource func() time.Time // source of time for waku
-
 	logger *zap.Logger
 
 	// NTP Synced timesource
@@ -166,7 +163,7 @@ func getUsableUDPPort() (int, error) {
 }
 
 // New creates a WakuV2 client ready to communicate through the LibP2P network.
-func New(nodeKey string, fleet string, cfg *Config, logger *zap.Logger, appDB *sql.DB, timesource *timesource.NTPTimeSource) (*Waku, error) {
+func New(nodeKey string, fleet string, cfg *Config, logger *zap.Logger, appDB *sql.DB, ts *timesource.NTPTimeSource) (*Waku, error) {
 	var err error
 	if logger == nil {
 		logger, err = zap.NewDevelopment()
@@ -175,6 +172,10 @@ func New(nodeKey string, fleet string, cfg *Config, logger *zap.Logger, appDB *s
 		}
 	}
 
+	if ts == nil {
+		ts = timesource.Default()
+	}
+
 	cfg = setDefaults(cfg)
 
 	if cfg.UDPPort == 0 {
@@ -201,8 +202,8 @@ func New(nodeKey string, fleet string, cfg *Config, logger *zap.Logger, appDB *s
 		dnsAddressCache:         make(map[string][]dnsdisc.DiscoveredNode),
 		dnsAddressCacheLock:     &sync.RWMutex{},
 		storeMsgIDs:             make(map[gethcommon.Hash]bool),
+		timesource:              ts,
 		storeMsgIDsMu:           sync.RWMutex{},
-		timeSource:              time.Now,
 		logger:                  logger,
 		discV5BootstrapNodes:    cfg.DiscV5BootstrapNodes,
 	}
@@ -764,12 +765,7 @@ func (w *Waku) ConfirmationsEnabled() bool {
 
 // CurrentTime returns current time.
 func (w *Waku) CurrentTime() time.Time {
-	return w.timeSource()
-}
-
-// SetTimeSource assigns a particular source of time to a waku object.
-func (w *Waku) SetTimeSource(timesource func() time.Time) {
-	w.timeSource = timesource
+	return w.timesource.Now()
 }
 
 // APIs returns the RPC descriptors the Waku implementation offers
@@ -1489,10 +1485,7 @@ func (w *Waku) AddStorePeer(address string) (peer.ID, error) {
 }
 
 func (w *Waku) timestamp() int64 {
-	if w.timesource != nil {
-		return w.timesource.Now().UnixNano()
-	}
-	return utils.GetUnixEpoch()
+	return w.timesource.Now().UnixNano()
 }
 
 func (w *Waku) autoRelayPeerSource(ctx context.Context, numPeers int) <-chan peer.AddrInfo {