From 88e7a58016965a787dfae13bab1d54f6d99cfda2 Mon Sep 17 00:00:00 2001
From: Richard Ramos <info@richardramos.me>
Date: Wed, 11 Dec 2024 14:41:37 -0400
Subject: [PATCH] fix: extract timeout from context

---
 wakuv2/nwaku.go | 40 ++++++++++++++++++++++++++--------------
 1 file changed, 26 insertions(+), 14 deletions(-)

diff --git a/wakuv2/nwaku.go b/wakuv2/nwaku.go
index e19b5eb36..72a14d9a8 100644
--- a/wakuv2/nwaku.go
+++ b/wakuv2/nwaku.go
@@ -2724,8 +2724,7 @@ func (n *WakuNode) Version() (string, error) {
 }
 
 func (n *WakuNode) StoreQuery(ctx context.Context, storeRequest *storepb.StoreQueryRequest, peerInfo peer.AddrInfo) (*storepb.StoreQueryResponse, error) {
-	// TODO: extract timeout from context
-	timeoutMs := time.Minute.Milliseconds()
+	timeoutMs := getContextTimeoutMilliseconds(ctx)
 
 	b, err := json.Marshal(storeRequest)
 	if err != nil {
@@ -2766,8 +2765,7 @@ func (n *WakuNode) StoreQuery(ctx context.Context, storeRequest *storepb.StoreQu
 }
 
 func (n *WakuNode) RelayPublish(ctx context.Context, message *pb.WakuMessage, pubsubTopic string) (pb.MessageHash, error) {
-	// TODO: extract timeout from context
-	timeoutMs := time.Minute.Milliseconds()
+	timeoutMs := getContextTimeoutMilliseconds(ctx)
 
 	jsonMsg, err := json.Marshal(message)
 	if err != nil {
@@ -2808,9 +2806,10 @@ func (n *WakuNode) DnsDiscovery(ctx context.Context, enrTreeUrl string, nameDnsS
 	defer C.free(unsafe.Pointer(cEnrTree))
 	defer C.free(unsafe.Pointer(cDnsServer))
 
-	// TODO: extract timeout from context
+	timeoutMs := getContextTimeoutMilliseconds(ctx)
+
 	wg.Add(1)
-	C.cGoWakuDnsDiscovery(n.wakuCtx, cEnrTree, cDnsServer, C.int(time.Minute.Milliseconds()), resp)
+	C.cGoWakuDnsDiscovery(n.wakuCtx, cEnrTree, cDnsServer, C.int(timeoutMs), resp)
 	wg.Wait()
 	if C.getRet(resp) == C.RET_OK {
 		var addrsRet []multiaddr.Multiaddr
@@ -2843,9 +2842,10 @@ func (n *WakuNode) PingPeer(ctx context.Context, peerInfo peer.AddrInfo) (time.D
 	var cPeerId = C.CString(strings.Join(addrs, ","))
 	defer C.free(unsafe.Pointer(cPeerId))
 
-	// TODO: extract timeout from ctx
+	timeoutMs := getContextTimeoutMilliseconds(ctx)
+
 	wg.Add(1)
-	C.cGoWakuPingPeer(n.wakuCtx, cPeerId, C.int(time.Minute.Milliseconds()), resp)
+	C.cGoWakuPingPeer(n.wakuCtx, cPeerId, C.int(timeoutMs), resp)
 	wg.Wait()
 	if C.getRet(resp) == C.RET_OK {
 		rttStr := C.GoStringN(C.getMyCharPtr(resp), C.int(C.getMyCharLen(resp)))
@@ -2943,9 +2943,10 @@ func (n *WakuNode) Connect(ctx context.Context, addr multiaddr.Multiaddr) error
 	defer C.freeResp(resp)
 	defer C.free(unsafe.Pointer(cPeerMultiAddr))
 
-	// TODO: extract timeout from ctx
+	timeoutMs := getContextTimeoutMilliseconds(ctx)
+
 	wg.Add(1)
-	C.cGoWakuConnect(n.wakuCtx, cPeerMultiAddr, C.int(time.Minute.Milliseconds()), resp)
+	C.cGoWakuConnect(n.wakuCtx, cPeerMultiAddr, C.int(timeoutMs), resp)
 	wg.Wait()
 	if C.getRet(resp) == C.RET_OK {
 		return nil
@@ -2965,9 +2966,10 @@ func (n *WakuNode) DialPeerByID(ctx context.Context, peerID peer.ID, protocol li
 	defer C.free(unsafe.Pointer(cPeerId))
 	defer C.free(unsafe.Pointer(cProtocol))
 
-	// TODO: extract timeout from ctx
+	timeoutMs := getContextTimeoutMilliseconds(ctx)
+
 	wg.Add(1)
-	C.cGoWakuDialPeerById(n.wakuCtx, cPeerId, cProtocol, C.int(time.Minute.Milliseconds()), resp)
+	C.cGoWakuDialPeerById(n.wakuCtx, cPeerId, cProtocol, C.int(timeoutMs), resp)
 	wg.Wait()
 	if C.getRet(resp) == C.RET_OK {
 		return nil
@@ -3124,9 +3126,11 @@ func (n *WakuNode) DialPeer(ctx context.Context, peerAddr multiaddr.Multiaddr, p
 	defer C.freeResp(resp)
 	defer C.free(unsafe.Pointer(cPeerMultiAddr))
 	defer C.free(unsafe.Pointer(cProtocol))
-	// TODO: extract timeout from context
+
+	timeoutMs := getContextTimeoutMilliseconds(ctx)
+
 	wg.Add(1)
-	C.cGoWakuDialPeer(n.wakuCtx, cPeerMultiAddr, cProtocol, C.int(requestTimeout.Milliseconds()), resp)
+	C.cGoWakuDialPeer(n.wakuCtx, cPeerMultiAddr, cProtocol, C.int(timeoutMs), resp)
 	wg.Wait()
 	if C.getRet(resp) == C.RET_OK {
 		return nil
@@ -3142,3 +3146,11 @@ func (n *WakuNode) GetNumConnectedPeers() (int, error) {
 	}
 	return len(peers), nil
 }
+
+func getContextTimeoutMilliseconds(ctx context.Context) int {
+	deadline, ok := ctx.Deadline()
+	if ok {
+		return int(time.Until(deadline).Milliseconds())
+	}
+	return 0
+}