Make the merkle proof path refresh safe against concurrent invalidation

ensureFreshMerkleProofPath returns the fetched path instead of leaving
callers to re-read merkleProofCache. generateProof uses the returned
value, so an invalidate landing during the refresh's updateMemberCache
suspension can no longer empty the cache under an in-flight proof-gen
and trip its empty-cache guard. The cache field is now only a hint for
the next call.

invalidateMerkleProofCache also bumps merkleProofCacheGeneration. The
refresh stamps the generation before each fetchMerkleProofElements and
only caches a result whose generation is unchanged, otherwise it fetches
again, bounded by MerkleProofRefetchMaxAttempts. A rejection-driven
invalidate is therefore never absorbed by a fetch that predates it, so
the path a caller receives always postdates the invalidate.

refreshRoots awaits rootsRefreshInFlightFut via join() at both await
sites, giving the roots refresh the same cancellation shield the proof
path refetch already had: a cancelled validateRoot caller no longer
cancels the refresh its coalesced peers are waiting on.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

logos_delivery/waku/rln/group_manager/on_chain/group_manager.nim

@@ -44,7 +44,8 @@ type
     registrationHandler*: Option[RegistrationHandler]
     latestProcessedBlock*: BlockNumber
     merkleProofCache*: seq[byte]
-    proofPathRefreshInFlightFut*: Future[void]
+    merkleProofCacheGeneration: uint64
+    proofPathRefreshInFlightFut*: Future[seq[byte]]
     lastRootsRefreshMoment*: Moment
     rootsRefreshInFlightFut*: Future[void]
 
@@ -235,12 +236,13 @@ proc updateMemberCount*(
 proc refreshRoots(g: OnchainGroupManager): Future[void] {.async.} =
   ## On-demand refresh of validRoots from the on-chain root cache.
   ## Throttled to at most one refresh per RootsRefreshMinInterval; concurrent
-  ## callers outside the throttle window coalesce onto a single in-flight refresh.
+  ## callers outside the throttle window coalesce onto a single in-flight
+  ## refresh, awaited via `join` so cancelling one caller never cancels it.
   if Moment.now() - g.lastRootsRefreshMoment < RootsRefreshMinInterval:
     return
 
   if not g.rootsRefreshInFlightFut.isNil() and not g.rootsRefreshInFlightFut.finished():
-    await g.rootsRefreshInFlightFut
+    await g.rootsRefreshInFlightFut.join()
     return
 
   proc doRefresh(): Future[void] {.async.} =
@@ -250,7 +252,7 @@ proc refreshRoots(g: OnchainGroupManager): Future[void] {.async.} =
     g.lastRootsRefreshMoment = Moment.now()
 
   g.rootsRefreshInFlightFut = doRefresh()
-  await g.rootsRefreshInFlightFut
+  await g.rootsRefreshInFlightFut.join()
 
 method validateRoot*(g: OnchainGroupManager, root: MerkleNode): Future[bool] {.async.} =
   if g.indexOfRoot(root) >= 0:
@@ -259,54 +261,71 @@ method validateRoot*(g: OnchainGroupManager, root: MerkleNode): Future[bool] {.a
   await g.refreshRoots()
   return g.indexOfRoot(root) >= 0
 
+# Bounds refetches when invalidations keep landing mid-fetch, so a rejection
+# storm cannot pin the refresh loop to the eth client.
+const MerkleProofRefetchMaxAttempts = 3
+
 proc ensureFreshMerkleProofPath*(
     g: OnchainGroupManager
-): Future[Result[void, string]] {.async.} =
-  ## Refetches `merkleProofCache` only when empty; suspected-stale callers
-  ## invalidate first. Concurrent callers coalesce onto one refetch, awaited
-  ## via `join` so cancelling one caller never cancels the shared refetch.
+): Future[Result[seq[byte], string]] {.async.} =
+  ## Returns the merkle proof path, refetching from chain when the cache is
+  ## empty; suspected-stale callers invalidate first. Concurrent callers
+  ## coalesce onto one refetch, awaited via `join` so cancelling one caller
+  ## never cancels the shared refetch. Callers use the returned path, not the
+  ## cache field, so a concurrent invalidate cannot fail an in-flight proof-gen.
   if g.membershipIndex.isNone():
     return err("membership index is not set")
 
   if g.merkleProofCache.len > 0:
-    return ok()
+    return ok(g.merkleProofCache)
 
-  if not g.proofPathRefreshInFlightFut.isNil() and
-      not g.proofPathRefreshInFlightFut.finished():
-    await g.proofPathRefreshInFlightFut.join()
-    if g.merkleProofCache.len > 0:
-      return ok()
+  proc doRefresh(): Future[seq[byte]] {.async.} =
+    var pathBytes: seq[byte]
+    for _ in 0 ..< MerkleProofRefetchMaxAttempts:
+      let generation = g.merkleProofCacheGeneration
+      pathBytes = (await g.fetchMerkleProofElements()).valueOr:
+        error "Failed to refresh merkle proof path", error = error
+        return @[]
+      if g.merkleProofCacheGeneration == generation:
+        g.merkleProofCache = pathBytes
+        break
+      # An invalidate raced this fetch, so the fetched path may predate the
+      # change that triggered it: fetch again. If attempts run out the last
+      # path is returned uncached and the next publish refetches.
+    if pathBytes.len > 0:
+      # Best-effort metric refresh - if there's a failure, it will update with the next root change
+      discard await g.updateMemberCount()
+    return pathBytes
+
+  if g.proofPathRefreshInFlightFut.isNil() or g.proofPathRefreshInFlightFut.finished():
+    g.proofPathRefreshInFlightFut = doRefresh()
+
+  let refreshFut = g.proofPathRefreshInFlightFut
+  await refreshFut.join()
+  if not refreshFut.completed():
     return err("merkle proof path refresh failed")
 
-  var fetchOk = false
-  proc doRefresh(): Future[void] {.async.} =
-    let pathBytes = (await g.fetchMerkleProofElements()).valueOr:
-      error "Failed to refresh merkle proof path", error = error
-      return
-    g.merkleProofCache = pathBytes
-    fetchOk = true
-    # Best-effort metric refresh - if there's a failure, it will update with the next root change
-    discard await g.updateMemberCount()
-
-  g.proofPathRefreshInFlightFut = doRefresh()
-  await g.proofPathRefreshInFlightFut.join()
-
-  if not fetchOk:
+  let pathBytes = refreshFut.read()
+  if pathBytes.len == 0:
     return err("merkle proof path refresh failed")
-  return ok()
+  return ok(pathBytes)
 
 method invalidateMerkleProofCache*(g: OnchainGroupManager) {.gcsafe, raises: [].} =
-  ## Empties the cache so the next `ensureFreshMerkleProofPath` refetches.
+  ## Empties the cache so the next `ensureFreshMerkleProofPath` refetches. The
+  ## generation bump makes a refetch already in flight fetch again rather than
+  ## serve a path that may predate this invalidate.
   g.merkleProofCache = @[]
+  g.merkleProofCacheGeneration.inc()
 
 method scheduleMerkleProofRefresh*(g: OnchainGroupManager) {.gcsafe, raises: [].} =
-  ## Empties the cache and spawns a detached refetch; a failed refetch is
+  ## Invalidates the cache and spawns a detached refetch; a failed refetch is
   ## logged and repaired by the next `ensureFreshMerkleProofPath`.
-  g.merkleProofCache = @[]
+  g.invalidateMerkleProofCache()
 
   proc refresh() {.async.} =
-    (await g.ensureFreshMerkleProofPath()).isOkOr:
-      warn "merkle proof refresh failed", error = error
+    let res = await g.ensureFreshMerkleProofPath()
+    if res.isErr():
+      warn "merkle proof refresh failed", error = res.error
 
   asyncSpawn refresh()
 
@@ -495,12 +514,9 @@ method generateProof*(
     return err("user message limit is not set")
 
   debug "Generating RLN proof"
-  ?(await g.ensureFreshMerkleProofPath())
+  let merkleProofPath = ?(await g.ensureFreshMerkleProofPath())
 
-  if g.merkleProofCache.len == 0:
-    return err("merkle proof cache is empty")
-
-  if (g.merkleProofCache.len mod 32) != 0:
+  if (merkleProofPath.len mod 32) != 0:
     return err("Invalid merkle proof cache length")
 
   let identity_secret = seqToField(g.idCredentials.get().idSecretHash)
@@ -509,8 +525,8 @@ method generateProof*(
   var path_elements = newSeq[byte](0)
 
   let identity_path_index = uint64ToIndex(g.membershipIndex.get(), 20)
-  for i in 0 ..< g.merkleProofCache.len div 32:
-    let chunk = g.merkleProofCache[i * 32 .. (i + 1) * 32 - 1]
+  for i in 0 ..< merkleProofPath.len div 32:
+    let chunk = merkleProofPath[i * 32 .. (i + 1) * 32 - 1]
     path_elements.add(chunk.reversed())
 
   let xCfr = hashToFieldLe(data).valueOr:

tests/waku_rln_relay/test_rln_group_manager_onchain.nim

@@ -432,6 +432,42 @@ suite "Onchain group manager":
       # replaced by a competing refresh.
       manager.rootsRefreshInFlightFut == inFlight
 
+  test "validateRoot: cancelling one caller does not cancel the shared roots refresh":
+    (waitFor manager.init()).isOkOr:
+      raiseAssert $error
+
+    let credentials = generateCredentials()
+    (waitFor manager.register(credentials, UserMessageLimit(20))).isOkOr:
+      assert false, "register failed: " & error
+
+    manager.lastRootsRefreshMoment = default(Moment)
+    manager.rootsRefreshInFlightFut = nil
+
+    var badRoot: MerkleNode
+    badRoot[0] = 0x66
+
+    let f1 = manager.validateRoot(badRoot)
+    let inFlight = manager.rootsRefreshInFlightFut
+    let f2 = manager.validateRoot(badRoot)
+
+    check:
+      inFlight != nil
+      not inFlight.finished()
+
+    # Cancel the initiating caller; the shared refresh must keep running for
+    # the coalesced one.
+    waitFor f1.cancelAndWait()
+
+    check:
+      f1.cancelled()
+      not inFlight.cancelled()
+
+    discard waitFor f2
+
+    check:
+      inFlight.completed()
+      manager.rootsRefreshInFlightFut == inFlight
+
   test "generateProof: fast-paths without refresh when cache is populated":
     (waitFor manager.init()).isOkOr:
       raiseAssert $error
@@ -642,6 +678,71 @@ suite "Onchain group manager":
     waitFor inFlight.join()
     check manager.merkleProofCache == goodCache
 
+  test "ensureFreshMerkleProofPath: invalidate during the refetch does not fail callers":
+    (waitFor manager.init()).isOkOr:
+      raiseAssert $error
+
+    let credentials = generateCredentials()
+    (waitFor manager.register(credentials, UserMessageLimit(20))).isOkOr:
+      assert false, "register failed: " & error
+
+    manager.merkleProofCache = @[]
+    manager.proofPathRefreshInFlightFut = nil
+
+    # Both callers coalesce onto a refetch suspended on the eth call.
+    let f1 = manager.ensureFreshMerkleProofPath()
+    let inFlight = manager.proofPathRefreshInFlightFut
+    let f2 = manager.ensureFreshMerkleProofPath()
+
+    check:
+      inFlight != nil
+      not inFlight.finished()
+
+    # A publish rejection lands mid-fetch. The invalidate must neither fail
+    # the coalesced callers nor be swallowed by the fetch it raced: the path
+    # they receive is refetched after the invalidate.
+    manager.invalidateMerkleProofCache()
+    check manager.merkleProofCache.len == 0
+
+    let r1 = waitFor f1
+    let r2 = waitFor f2
+
+    check:
+      r1.isOk()
+      r1.get().len > 0
+      r2.isOk()
+      r2.get() == r1.get()
+      # The post-invalidate refetch repopulated the cache.
+      manager.merkleProofCache == r1.get()
+
+  test "generateProof: succeeds when the cache is invalidated mid-refetch":
+    (waitFor manager.init()).isOkOr:
+      raiseAssert $error
+
+    let credentials = generateCredentials()
+    (waitFor manager.register(credentials, UserMessageLimit(20))).isOkOr:
+      assert false, "register failed: " & error
+
+    manager.merkleProofCache = @[]
+    manager.proofPathRefreshInFlightFut = nil
+
+    # generateProof runs down to the refetch's suspended eth call.
+    let proofFut = manager.generateProof(
+      data = "hello".toBytes(), epoch = default(Epoch), messageId = MessageId(1)
+    )
+
+    check:
+      manager.proofPathRefreshInFlightFut != nil
+      not manager.proofPathRefreshInFlightFut.finished()
+
+    # A concurrent publish rejection empties the cache while proof-gen waits
+    # on the refetch.
+    manager.invalidateMerkleProofCache()
+
+    let proofRes = waitFor proofFut
+    check:
+      proofRes.isOk()
+
   test "verifyProof: should verify valid proof":
     let credentials = generateCredentials()
     (waitFor manager.init()).isOkOr: