diff --git a/src/mix_rln_spam_protection.nim b/src/mix_rln_spam_protection.nim index f7dec01..00d2f6c 100644 --- a/src/mix_rln_spam_protection.nim +++ b/src/mix_rln_spam_protection.nim @@ -26,8 +26,11 @@ ## await plugin.init() ## ## # Set up logos-messaging integration -## plugin.setPublishCallback(proc(topic: string, data: seq[byte]) {.async.} = +## plugin.setPublishCallback(proc( +## topic: string, data: seq[byte] +## ): Future[Result[void, string]] {.async.} = ## await logosMessaging.publish(topic, data) +## ok() ## ) ## ## # Start the plugin @@ -71,26 +74,10 @@ import ./mix_rln_spam_protection/[ - types, - constants, - protobuf, - codec, - rln_interface, - group_manager, - nullifier_log, - spam_protection, - coordination, - credentials + types, constants, protobuf, codec, rln_interface, group_manager, nullifier_log, + spam_protection, coordination, credentials, ] export - types, - constants, - protobuf, - codec, - spam_protection, - coordination, - credentials, - group_manager, - nullifier_log, - rln_interface + types, constants, protobuf, codec, spam_protection, coordination, credentials, + group_manager, nullifier_log, rln_interface diff --git a/src/mix_rln_spam_protection/coordination.nim b/src/mix_rln_spam_protection/coordination.nim index 77a2e44..d2392be 100644 --- a/src/mix_rln_spam_protection/coordination.nim +++ b/src/mix_rln_spam_protection/coordination.nim @@ -25,7 +25,8 @@ logScope: type # Subscription handler types MembershipSubscriptionHandler* = proc(update: MembershipUpdate) {.gcsafe, raises: [].} - MetadataSubscriptionHandler* = proc(metadata: ProofMetadataBroadcast) {.gcsafe, raises: [].} + MetadataSubscriptionHandler* = + proc(metadata: ProofMetadataBroadcast) {.gcsafe, raises: [].} # Coordination layer wrapper CoordinationLayer* = ref object @@ -45,7 +46,7 @@ proc newCoordinationLayer*(spamProtection: MixRlnSpamProtection): CoordinationLa spamProtection: spamProtection, publishCallback: none(PublishCallback), onMembershipUpdate: none(MembershipSubscriptionHandler), - onProofMetadata: none(MetadataSubscriptionHandler) + onProofMetadata: none(MetadataSubscriptionHandler), ) proc setPublishCallback*(cl: CoordinationLayer, callback: PublishCallback) = @@ -54,18 +55,20 @@ proc setPublishCallback*(cl: CoordinationLayer, callback: PublishCallback) = cl.publishCallback = some(callback) cl.spamProtection.setPublishCallback(callback) -proc setMembershipUpdateHandler*(cl: CoordinationLayer, handler: MembershipSubscriptionHandler) = +proc setMembershipUpdateHandler*( + cl: CoordinationLayer, handler: MembershipSubscriptionHandler +) = ## Set additional handler for membership updates (for custom processing). cl.onMembershipUpdate = some(handler) -proc setProofMetadataHandler*(cl: CoordinationLayer, handler: MetadataSubscriptionHandler) = +proc setProofMetadataHandler*( + cl: CoordinationLayer, handler: MetadataSubscriptionHandler +) = ## Set additional handler for proof metadata (for custom processing). cl.onProofMetadata = some(handler) proc handleIncomingMessage*( - cl: CoordinationLayer, - contentTopic: string, - data: seq[byte] + cl: CoordinationLayer, contentTopic: string, data: seq[byte] ): Future[RlnResult[void]] {.async.} = ## Route incoming messages from logos-messaging to appropriate handlers. ## @@ -85,7 +88,6 @@ proc handleIncomingMessage*( let update = MembershipUpdate.decode(data).valueOr: return err("Failed to decode update for handler: " & $error) cl.onMembershipUpdate.get()(update) - elif contentTopic == metadataTopic: # Handle proof metadata let metadataResult = cl.spamProtection.handleProofMetadata(data) @@ -97,7 +99,6 @@ proc handleIncomingMessage*( let metadata = ProofMetadataBroadcast.decode(data).valueOr: return err("Failed to decode metadata for handler: " & $error) cl.onProofMetadata.get()(metadata) - else: return err("Unknown content topic: " & contentTopic) @@ -112,12 +113,14 @@ proc getDefaultContentTopics*(): seq[string] = @[MembershipContentTopic, ProofMetadataContentTopic] # Helper function for building a logos-messaging subscription filter -proc buildSubscriptionFilter*(cl: CoordinationLayer): seq[tuple[contentTopic: string, handler: string]] = +proc buildSubscriptionFilter*( + cl: CoordinationLayer +): seq[tuple[contentTopic: string, handler: string]] = ## Build a subscription filter for logos-messaging. ## Returns tuples of (contentTopic, handlerName) for documentation. @[ (cl.spamProtection.getMembershipContentTopic(), "handleMembershipUpdate"), - (cl.spamProtection.getProofMetadataContentTopic(), "handleProofMetadata") + (cl.spamProtection.getProofMetadataContentTopic(), "handleProofMetadata"), ] # Integration example documentation @@ -140,8 +143,11 @@ await spamProtection.init() let coordination = newCoordinationLayer(spamProtection) # Wire up publish callback to logos-messaging -coordination.setPublishCallback(proc(topic: string, data: seq[byte]) {.async.} = +coordination.setPublishCallback(proc( + topic: string, data: seq[byte] +): Future[Result[void, string]] {.async.} = await logosMessaging.publish(topic, data) + ok() ) # Subscribe to RLN topics (uses configured content topics) @@ -167,20 +173,25 @@ let mixProto = MixProtocol.new( # Utility for creating a simple publish callback that logs (for testing) proc createLoggingPublishCallback*(): PublishCallback = ## Create a publish callback that just logs messages (for testing). - proc callback(contentTopic: string, data: seq[byte]): Future[void] {.async, gcsafe.} = + proc callback( + contentTopic: string, data: seq[byte] + ): Future[Result[void, string]] {.async, gcsafe.} = debug "Would publish to coordination layer", - topic = contentTopic, - dataLen = data.len + topic = contentTopic, dataLen = data.len + ok() + return callback # Utility for printing membership update in human-readable form proc formatMembershipUpdate*(update: MembershipUpdate): string = ## Format a membership update for logging/display. - let actionStr = case update.action + let actionStr = + case update.action of MembershipAction.Add: "ADD" of MembershipAction.Remove: "REMOVE" - result = actionStr & " member at index " & $update.index & - " (idCommitment: " & update.idCommitment[0..7].toHex() & "...)" + result = + actionStr & " member at index " & $update.index & " (idCommitment: " & + update.idCommitment[0 .. 7].toHex() & "...)" # Note: toHex is imported from types module via spam_protection diff --git a/src/mix_rln_spam_protection/group_manager.nim b/src/mix_rln_spam_protection/group_manager.nim index 9a31922..97ac998 100644 --- a/src/mix_rln_spam_protection/group_manager.nim +++ b/src/mix_rln_spam_protection/group_manager.nim @@ -68,12 +68,14 @@ type ## Group manager that propagates membership via logos-messaging content topics. ## Membership additions and deletions are broadcast to all nodes. publishCallback: Option[PublishCallback] - membershipByIdCommitment: Table[IDCommitment, MembershipIndex] ## idCommitment -> index (for spam recovery) + membershipByIdCommitment: Table[IDCommitment, MembershipIndex] + ## idCommitment -> index (for spam recovery) membershipByIndex: Table[MembershipIndex, IDCommitment] ## index -> idCommitment # We track per-user rate limits because the Merkle tree stores rateCommitment = Poseidon(idCommitment, userMessageLimit), # not the raw userMessageLimit. When serializing tree snapshots, we need the original userMessageLimit to correctly # recompute rateCommitment on load. Without this, we'd have to assume all members use the same limit. - rateLimitByIdCommitment: Table[IDCommitment, uint64] ## idCommitment -> userMessageLimit + rateLimitByIdCommitment: Table[IDCommitment, uint64] + ## idCommitment -> userMessageLimit nextIndex: MembershipIndex membershipContentTopic*: string ## Content topic for membership updates @@ -161,9 +163,7 @@ proc refreshProofCache*(gm: OffchainGroupManager): RlnResult[void] = return ok() let cache = gm.rlnInstance.generatePartialProofCache( - gm.credentials.get(), - gm.membershipIndex.get(), - gm.userMessageLimit, + gm.credentials.get(), gm.membershipIndex.get(), gm.userMessageLimit ).valueOr: gm.partialProofCache = none(PartialProofCache) return err("Failed to refresh partial proof cache: " & error) @@ -234,8 +234,7 @@ method generateProof*( let index = gm.membershipIndex.get() trace "Generating proof with credentials", - membershipIndex = index, - commitment = creds.idCommitment.toHex() + membershipIndex = index, commitment = creds.idCommitment.toHex() # Flush tree to ensure internal state is synced if not flush(gm.rlnInstance.ctx): @@ -252,8 +251,7 @@ method generateProof*( return err("Failed to compute expected rate commitment: " & error) trace "Tree state verification before proof generation", - membershipIndex = index, - commitmentsMatch = treeCommitment == expectedRateCommitment + membershipIndex = index, commitmentsMatch = treeCommitment == expectedRateCommitment if treeCommitment != expectedRateCommitment: error "CRITICAL: Tree commitment at index does not match our credentials!", @@ -270,9 +268,7 @@ method generateProof*( return err("Failed to get current Merkle root: " & error) trace "Generating RLN proof", - signalLen = signal.len, - messageId = messageId, - membershipIndex = index + signalLen = signal.len, messageId = messageId, membershipIndex = index # Happy path: cached partial-proof should be available and valid most of the time. # On cache miss/staleness, rebuild once and retry. Fall back to full proof generation. @@ -309,7 +305,8 @@ method generateProof*( ) usedPartialCache = proofResult.isOk if proofResult.isErr: - warn "Refreshed partial proof cache could not be used", error = proofResult.error + warn "Refreshed partial proof cache could not be used", + error = proofResult.error if not usedPartialCache: proofResult = gm.rlnInstance.generateRlnProofWithWitness( @@ -429,11 +426,7 @@ proc restoreMemberFromKeystore*( if not gm.isInitialized: return err("Group manager not initialized") - let memberLimit = - if userMessageLimit > 0: - userMessageLimit - else: - gm.userMessageLimit + let memberLimit = if userMessageLimit > 0: userMessageLimit else: gm.userMessageLimit # Compute rate commitment = Poseidon(idCommitment, userMessageLimit) # This is the actual leaf value stored in the RLN Merkle tree @@ -512,7 +505,8 @@ method register*( index: index, ) let data = update.toBytes() - await gm.publishCallback.get()(gm.membershipContentTopic, data) + (await gm.publishCallback.get()(gm.membershipContentTopic, data)).isOkOr: + warn "Failed to broadcast membership update", error = error # Call callback if gm.onRegister.isSome: @@ -538,7 +532,8 @@ proc registerWithLimit*( return err("Failed to compute rate commitment: " & error) let index = gm.nextIndex - trace "Registering member with custom limit", index = index, userMessageLimit = userMessageLimit + trace "Registering member with custom limit", + index = index, userMessageLimit = userMessageLimit gm.nextIndex += 1 # Insert rateCommitment into RLN tree @@ -564,13 +559,15 @@ proc registerWithLimit*( index: index, ) let data = update.toBytes() - await gm.publishCallback.get()(gm.membershipContentTopic, data) + (await gm.publishCallback.get()(gm.membershipContentTopic, data)).isOkOr: + warn "Failed to broadcast membership update", error = error # Call callback if gm.onRegister.isSome: await gm.onRegister.get()(commitment, index) - debug "Member registered with custom limit", index = index, userMessageLimit = userMessageLimit + debug "Member registered with custom limit", + index = index, userMessageLimit = userMessageLimit ok(index) method register*( @@ -613,9 +610,8 @@ method withdraw*( return err("Failed to delete member: " & deleteResult.error) # Get the member's rate limit before deleting (for broadcast) - let memberRateLimit = gm.rateLimitByIdCommitment.getOrDefault( - idCommitment, gm.userMessageLimit - ) + let memberRateLimit = + gm.rateLimitByIdCommitment.getOrDefault(idCommitment, gm.userMessageLimit) # Update local tracking gm.membershipByIdCommitment.del(idCommitment) @@ -635,7 +631,8 @@ method withdraw*( index: index, ) let data = update.toBytes() - await gm.publishCallback.get()(gm.membershipContentTopic, data) + (await gm.publishCallback.get()(gm.membershipContentTopic, data)).isOkOr: + warn "Failed to broadcast membership update", error = error # Call callback if gm.onWithdraw.isSome: @@ -789,8 +786,8 @@ proc getMemberRateLimit*( # └─────────────────────────────────────────────────────────────┘ const - SnapshotHeaderSize = 16 # member_count (8) + next_index (8) - SnapshotMemberSize = 48 # commitment (32) + index (8) + userMessageLimit (8) + SnapshotHeaderSize = 16 # member_count (8) + next_index (8) + SnapshotMemberSize = 48 # commitment (32) + index (8) + userMessageLimit (8) # Note: writeUint64LE and readUint64LE are imported from bytes_utils via types @@ -821,9 +818,8 @@ proc serializeTreeSnapshot*(gm: OffchainGroupManager): seq[byte] = offset += Uint64ByteSize # Write userMessageLimit - let rateLimit = gm.rateLimitByIdCommitment.getOrDefault( - commitment, gm.userMessageLimit - ) + let rateLimit = + gm.rateLimitByIdCommitment.getOrDefault(commitment, gm.userMessageLimit) result.writeUint64LE(offset, rateLimit) offset += Uint64ByteSize @@ -849,7 +845,10 @@ proc loadTreeSnapshot*(gm: OffchainGroupManager, data: seq[byte]): RlnResult[voi if data.len != expectedSize: error "Snapshot size mismatch", actual = data.len, expected = expectedSize, memberCount = memberCount - return err("Invalid snapshot: size mismatch (expected " & $expectedSize & ", got " & $data.len & ")") + return err( + "Invalid snapshot: size mismatch (expected " & $expectedSize & ", got " & $data.len & + ")" + ) trace "Loading tree snapshot", memberCount = memberCount, nextIndex = nextIndex @@ -883,7 +882,8 @@ proc loadTreeSnapshot*(gm: OffchainGroupManager, data: seq[byte]): RlnResult[voi # Insert into RLN Merkle tree let insertResult = gm.rlnInstance.insertMemberAt(index, rateCommitment) if insertResult.isErr: - error "Failed to insert member from snapshot", index = index, error = insertResult.error + error "Failed to insert member from snapshot", + index = index, error = insertResult.error return err("Failed to insert member: " & insertResult.error) # Update tracking tables - track by idCommitment for spam recovery @@ -897,7 +897,8 @@ proc loadTreeSnapshot*(gm: OffchainGroupManager, data: seq[byte]): RlnResult[voi gm.resetRootTrackerOrLog() gm.refreshProofCacheOrLog() - debug "Tree snapshot loaded successfully", memberCount = memberCount, nextIndex = nextIndex + debug "Tree snapshot loaded successfully", + memberCount = memberCount, nextIndex = nextIndex ok() proc saveTreeToFile*(gm: OffchainGroupManager, path: string): RlnResult[void] = @@ -929,9 +930,7 @@ proc loadTreeFromFile*(gm: OffchainGroupManager, path: string): RlnResult[void] if not flush(gm.rlnInstance.ctx): return err("Failed to flush tree after loading") - debug "Tree loaded and flushed", - path = path, - memberCount = gm.membershipByIndex.len + debug "Tree loaded and flushed", path = path, memberCount = gm.membershipByIndex.len return ok() except IOError as e: diff --git a/src/mix_rln_spam_protection/spam_protection.nim b/src/mix_rln_spam_protection/spam_protection.nim index cad13c9..f57327c 100644 --- a/src/mix_rln_spam_protection/spam_protection.nim +++ b/src/mix_rln_spam_protection/spam_protection.nim @@ -7,7 +7,7 @@ ## This module provides the MixRlnSpamProtection type that can be used with ## the mix protocol for per-hop proof generation and verification. -import std/[math, options, deques, times] +import std/[math, options, deques, times, sequtils] import chronos import results import chronicles @@ -65,6 +65,10 @@ type publishCallback: Option[PublishCallback] spamHandler: Option[SpamHandler] epochTimerLoop: Future[void] + pendingBroadcasts: seq[Future[void]] + ## In-flight best-effort proof-metadata broadcasts. verifyProof is + ## synchronous so it cannot await these; they are tracked here (pruned + ## when finished) and cancelled in stop() instead of being orphaned. proc defaultRlnIdentifier*(): RlnIdentifier = ## Get the default RLN identifier. @@ -257,11 +261,26 @@ proc stop*(sp: MixRlnSpamProtection) {.async.} = if not sp.epochTimerLoop.isNil: await sp.epochTimerLoop.cancelAndWait() + # Cancel any in-flight proof-metadata broadcasts. + for fut in sp.pendingBroadcasts: + if not fut.isNil and not fut.finished: + await fut.cancelAndWait() + sp.pendingBroadcasts.setLen(0) + await sp.groupManager.stop() await sp.nullifierLog.stop() info "MixRlnSpamProtection stopped" +proc broadcastProofMetadata( + sp: MixRlnSpamProtection, data: seq[byte] +) {.async, gcsafe.} = + ## Best-effort broadcast of proof metadata to the coordination layer. + ## verifyProof is synchronous (it overrides a sync base method), so this runs + ## as a tracked background future rather than a fire-and-forget asyncSpawn. + (await sp.publishCallback.get()(sp.config.proofMetadataContentTopic, data)).isOkOr: + warn "Failed to broadcast proof metadata", error = error + {.push raises: [], gcsafe.} proc isReady*(sp: MixRlnSpamProtection): bool = @@ -617,7 +636,11 @@ method verifyProof*( epoch: proof.epoch, ) let data = broadcast.toBytes() - asyncSpawn sp.publishCallback.get()(sp.config.proofMetadataContentTopic, data) + # verifyProof is synchronous, so track the background broadcast future + # (pruning finished ones) rather than firing an untracked asyncSpawn; + # stop() cancels any that are still in flight. + sp.pendingBroadcasts.keepItIf(not it.finished) + sp.pendingBroadcasts.add(sp.broadcastProofMetadata(data)) debug "Proof verified successfully", epoch = epochToUint64(proof.epoch), diff --git a/src/mix_rln_spam_protection/types.nim b/src/mix_rln_spam_protection/types.nim index ddfd093..d67c404 100644 --- a/src/mix_rln_spam_protection/types.nim +++ b/src/mix_rln_spam_protection/types.nim @@ -106,9 +106,11 @@ type epoch*: Epoch # Callback types for coordination layer integration - PublishCallback* = - proc(contentTopic: string, data: seq[byte]): Future[void] {.gcsafe, raises: [].} + PublishCallback* = proc( + contentTopic: string, data: seq[byte] + ): Future[Result[void, string]] {.gcsafe, raises: [].} ## Callback for publishing messages to logos-messaging. + ## Returns an error string if publishing failed so callers can handle it. # RLN Witness input for explicit Merkle proof-based proof generation Field* = array[32, byte] ## 32-byte field element representation (256 bits). @@ -150,8 +152,7 @@ type # as little-endian uint64. The remaining 24 bytes are zero-padded. proc calcEpoch*( - timestamp: float64, - epochDurationSeconds: float64 = EpochDurationSeconds, + timestamp: float64, epochDurationSeconds: float64 = EpochDurationSeconds ): Epoch = ## Calculate the epoch for a given Unix timestamp. let epochNum = @@ -165,16 +166,11 @@ proc calcEpoch*( for i in 0 ..< Uint64ByteSize: result[i] = epochBytes[i] -proc calcEpoch*( - t: Time, - epochDurationSeconds: float64 = EpochDurationSeconds, -): Epoch = +proc calcEpoch*(t: Time, epochDurationSeconds: float64 = EpochDurationSeconds): Epoch = ## Calculate the epoch for a given Time. calcEpoch(t.toUnixFloat(), epochDurationSeconds) -proc currentEpoch*( - epochDurationSeconds: float64 = EpochDurationSeconds, -): Epoch = +proc currentEpoch*(epochDurationSeconds: float64 = EpochDurationSeconds): Epoch = ## Get the current epoch based on system time. calcEpoch(getTime(), epochDurationSeconds)