From abf76fdbbaf803549fa6bfbe340a0c459151adb0 Mon Sep 17 00:00:00 2001
From: Ceyhun Onur <ceyhun.onur@avalabs.org>
Date: Thu, 17 Oct 2024 17:49:51 +0300
Subject: [PATCH] cache signatures only in acp118 handler (#3474)

---
 network/p2p/acp118/handler.go      | 25 ++++++++++++++-----------
 network/p2p/acp118/handler_test.go |  6 ++++++
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/network/p2p/acp118/handler.go b/network/p2p/acp118/handler.go
index fd7ac8dc4f3..971f29af9b8 100644
--- a/network/p2p/acp118/handler.go
+++ b/network/p2p/acp118/handler.go
@@ -39,16 +39,16 @@ func NewHandler(verifier Verifier, signer warp.Signer) *Handler {
 }
 
 // NewCachedHandler returns an instance of Handler that caches successful
-// requests.
+// signatures.
 func NewCachedHandler(
 	cacher cache.Cacher[ids.ID, []byte],
 	verifier Verifier,
 	signer warp.Signer,
 ) *Handler {
 	return &Handler{
-		cacher:   cacher,
-		verifier: verifier,
-		signer:   signer,
+		signatureCache: cacher,
+		verifier:       verifier,
+		signer:         signer,
 	}
 }
 
@@ -56,9 +56,9 @@ func NewCachedHandler(
 type Handler struct {
 	p2p.NoOpHandler
 
-	cacher   cache.Cacher[ids.ID, []byte]
-	verifier Verifier
-	signer   warp.Signer
+	signatureCache cache.Cacher[ids.ID, []byte]
+	verifier       Verifier
+	signer         warp.Signer
 }
 
 func (h *Handler) AppRequest(
@@ -84,8 +84,8 @@ func (h *Handler) AppRequest(
 	}
 
 	msgID := msg.ID()
-	if responseBytes, ok := h.cacher.Get(msgID); ok {
-		return responseBytes, nil
+	if signatureBytes, ok := h.signatureCache.Get(msgID); ok {
+		return signatureToResponse(signatureBytes)
 	}
 
 	if err := h.verifier.Verify(ctx, msg, request.Justification); err != nil {
@@ -100,6 +100,11 @@ func (h *Handler) AppRequest(
 		}
 	}
 
+	h.signatureCache.Put(msgID, signature)
+	return signatureToResponse(signature)
+}
+
+func signatureToResponse(signature []byte) ([]byte, *common.AppError) {
 	response := &sdk.SignatureResponse{
 		Signature: signature,
 	}
@@ -111,7 +116,5 @@ func (h *Handler) AppRequest(
 			Message: fmt.Sprintf("failed to marshal response: %s", err),
 		}
 	}
-
-	h.cacher.Put(msgID, responseBytes)
 	return responseBytes, nil
 }
diff --git a/network/p2p/acp118/handler_test.go b/network/p2p/acp118/handler_test.go
index 081ecd1f035..e58d61a8f6a 100644
--- a/network/p2p/acp118/handler_test.go
+++ b/network/p2p/acp118/handler_test.go
@@ -125,6 +125,12 @@ func TestHandler(t *testing.T) {
 				require.NoError(err)
 
 				require.True(bls.Verify(pk, signature, request.Message))
+
+				// Ensure the cache is populated with correct signature
+				sig, ok := tt.cacher.Get(unsignedMessage.ID())
+				if ok {
+					require.Equal(sig, response.Signature)
+				}
 			}
 
 			for _, expectedErr = range tt.expectedErrs {