From d43352dd6e3ab607430ae6116840beb37ae851b4 Mon Sep 17 00:00:00 2001
From: Jernej Kos <jernej@kos.mx>
Date: Thu, 1 Feb 2024 14:44:48 +0100
Subject: [PATCH] go/worker/keymanager/p2p: Explicitly configure timeout for
 CallEnclave

---
 go/worker/keymanager/p2p/client.go   | 4 +++-
 go/worker/keymanager/p2p/protocol.go | 9 +++++++--
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/go/worker/keymanager/p2p/client.go b/go/worker/keymanager/p2p/client.go
index babfe9fa00b..b7972d1168e 100644
--- a/go/worker/keymanager/p2p/client.go
+++ b/go/worker/keymanager/p2p/client.go
@@ -35,7 +35,9 @@ type client struct {
 
 func (c *client) CallEnclave(ctx context.Context, request *CallEnclaveRequest, peers []core.PeerID) (*CallEnclaveResponse, rpc.PeerFeedback, error) {
 	var rsp CallEnclaveResponse
-	pf, err := c.rc.CallOne(ctx, c.mgr.GetBestPeers(rpc.WithLimitPeers(peers)), MethodCallEnclave, request, &rsp)
+	pf, err := c.rc.CallOne(ctx, c.mgr.GetBestPeers(rpc.WithLimitPeers(peers)), MethodCallEnclave, request, &rsp,
+		rpc.WithMaxPeerResponseTime(MethodCallEnclaveTimeout),
+	)
 	if err != nil {
 		return nil, nil, err
 	}
diff --git a/go/worker/keymanager/p2p/protocol.go b/go/worker/keymanager/p2p/protocol.go
index ef09af50588..1eeb6a32af1 100644
--- a/go/worker/keymanager/p2p/protocol.go
+++ b/go/worker/keymanager/p2p/protocol.go
@@ -1,6 +1,8 @@
 package p2p
 
 import (
+	"time"
+
 	"github.com/libp2p/go-libp2p/core"
 
 	"github.com/oasisprotocol/oasis-core/go/common/node"
@@ -16,8 +18,11 @@ const KeyManagerProtocolID = "keymanager"
 // KeyManagerProtocolVersion is the supported version of the keymanager protocol.
 var KeyManagerProtocolVersion = version.Version{Major: 2, Minor: 0, Patch: 0}
 
-// MethodCallEnclave is the name of the CallEnclave method.
-const MethodCallEnclave = "CallEnclave"
+// Constants related to the CallEnclave method.
+const (
+	MethodCallEnclave        = "CallEnclave"
+	MethodCallEnclaveTimeout = 3 * time.Second
+)
 
 // CallEnclaveRequest is a CallEnclave request.
 type CallEnclaveRequest struct {