Inference msgs optimization: optimize key verification#779
Inference msgs optimization: optimize key verification#779DimaOrekhovPS wants to merge 27 commits intoupgrade-v0.2.11from
Conversation
There was a problem hiding this comment.
Pull request overview
This PR optimizes signature verification for inference messages by implementing a dual-path approach that handles out-of-order message arrival (start-first vs finish-first). The optimization reduces redundant cryptographic signature verification: the first message (whether Start or Finish) performs full signature verification, while the second message performs component comparison only. Additionally, the PR changes the StartProcessed() detection logic to use MaxTokens instead of PromptHash to avoid false positives in the finish-first flow.
Changes:
- Modified
StartProcessed()to useMaxTokens != 0instead ofPromptHash != ""as the marker for whether StartInference has been processed - Implemented signature verification policy: first message verifies signatures, second message performs equality checks
- Added new error types for component mismatches:
ErrDevComponentMismatch,ErrTAComponentMismatch,ErrInferenceRoleMismatch - Added comparison functions to validate component consistency between Start and Finish messages in out-of-order scenarios
Reviewed changes
Copilot reviewed 11 out of 11 changed files in this pull request and generated 16 comments.
Show a summary per file
| File | Description |
|---|---|
| inference-chain/x/inference/types/inference.go | Changed StartProcessed() to use MaxTokens instead of PromptHash |
| inference-chain/x/inference/types/errors.go | Added three new error types for component mismatches |
| inference-chain/x/inference/keeper/msg_server_validation_test.go | Added missing Creator field to FinishInference test message |
| inference-chain/x/inference/keeper/msg_server_start_inference.go | Implemented dual-path signature verification: verify on first, compare on second |
| inference-chain/x/inference/keeper/msg_server_out_of_order_inference_test.go | Updated mock expectations and added OriginalPromptHash assertions |
| inference-chain/x/inference/keeper/msg_server_finish_inference_test.go | Added comprehensive tests for start-first/finish-first scenarios and mismatch detection |
| inference-chain/x/inference/keeper/msg_server_finish_inference.go | Implemented dual-path signature verification for FinishInference |
| inference-chain/x/inference/keeper/msg_server_finish_first_missing_signature_fields_test.go | New test documenting bug where hash fields aren't persisted in finish-first flow |
| inference-chain/x/inference/calculations/inference_state_test.go | Updated test to reflect new StartProcessed() logic using MaxTokens |
| inference-chain/x/inference/calculations/inference_state.go | Updated startProcessed() function to use MaxTokens instead of PromptHash |
| decentralized-api/internal/server/public/post_chat_handler.go | Added TODO comment about rejecting empty PromptHash |
inference-chain/x/inference/keeper/msg_server_out_of_order_inference_test.go
Show resolved
Hide resolved
inference-chain/x/inference/keeper/msg_server_start_inference.go
Outdated
Show resolved
Hide resolved
inference-chain/x/inference/keeper/msg_server_finish_first_missing_signature_fields_test.go
Outdated
Show resolved
Hide resolved
| // TODO: We need to include inferenceId in the TA signature to make sure executor can't substitute the modified prompt | ||
| // TODO: any error here should lead to punishing the TA |
There was a problem hiding this comment.
This TODO highlights a critical security concern: without the inferenceId in the TA signature, an executor could potentially substitute a modified prompt while keeping the same hash. This should be addressed before the signature verification optimization is considered complete.
Additionally, the second TODO mentions that any error in these comparisons should lead to punishing the TA. Consider implementing this punishment mechanism as it's important for maintaining system integrity.
| // TODO: punish executor if Dev fails | ||
| if err := compareFinishDevComponents(msg, &existingInference); err != nil { | ||
| k.LogError("FinishInference: dev component mismatch", types.Inferences, "error", err, "inferenceId", msg.InferenceId) | ||
| return failedFinish(ctx, err, msg), nil | ||
| } | ||
| // TODO: re-check TA signature if any of the checks fail and punish the TA if the signature is invalid |
There was a problem hiding this comment.
These TODOs indicate incomplete implementation of the punishment mechanism. The comparison checks are meant to detect malicious behavior, but without punishing the bad actors, the system remains vulnerable to:
- Executors submitting incorrect dev components (line 82)
- Transfer agents submitting invalid signatures that fail comparison (line 87)
Implementing these punishment mechanisms should be a priority to ensure the signature verification optimization doesn't create security gaps.
inference-chain/x/inference/keeper/msg_server_finish_inference.go
Outdated
Show resolved
Hide resolved
inference-chain/x/inference/keeper/msg_server_out_of_order_inference_test.go
Show resolved
Hide resolved
| // StartInference always assigns MaxTokens (explicit or default). | ||
| // Finish-first flow can populate PromptHash early, so use MaxTokens to detect | ||
| // whether StartInference has already been processed. | ||
| return inference.MaxTokens != 0 |
There was a problem hiding this comment.
Alternatively we can check inference.AssignedTo != ""
| return currentInference, &payments, nil | ||
| } | ||
|
|
||
| func startProcessed(inference *types.Inference) bool { |
There was a problem hiding this comment.
Removed duplicate implementations, see inference.go
slandymani
left a comment
slandymani
left a comment
There was a problem hiding this comment.
change pr description after removing failOnCompareMismatch
| } | ||
| if request.PromptHash != "" && computedPromptHash != request.PromptHash { | ||
| if request.PromptHash == "" { | ||
| logging.Error("Empty prompt hash", types.Inferences, "error", err) |
| inference.PromptHash, | ||
| ) | ||
| } | ||
| if inference.RequestTimestamp != msg.RequestTimestamp { |
There was a problem hiding this comment.
duplicated comparison of msg.RequestTimestamp and msg.TransferredBy, we already checked them in compareFinishDevComponents
| inference.PromptHash, | ||
| ) | ||
| } | ||
| if inference.RequestTimestamp != msg.RequestTimestamp { |
There was a problem hiding this comment.
duplicated comparison of msg.RequestTimestamp and inference.TransferredBy, we already checked them in compareStartDevComponents
|
|
||
| // Record the current price only if this is the first message (FinishInference not processed yet) | ||
| // This ensures consistent pricing regardless of message arrival order | ||
| if !existingInference.FinishedProcessed() { |
There was a problem hiding this comment.
should this be moved to the previous if statement so as not to duplicate the check?
|
|
||
| // Record the current price only if this is the first message (StartInference not processed yet) | ||
| // This ensures consistent pricing regardless of message arrival order | ||
| if !existingInference.StartProcessed() { |
There was a problem hiding this comment.
move price recording to previous if statement, merge model checks with compareFinishModelField
…inference-ignite into do/inference-optimization-keys-2 # Conflicts: # inference-chain/x/inference/keeper/msg_server_out_of_order_inference_test.go # inference-chain/x/inference/types/errors.go
…inference-ignite into do/inference-optimization-keys-2 # Conflicts: # inference-chain/x/inference/keeper/msg_server_finish_inference.go # inference-chain/x/inference/keeper/msg_server_start_inference.go
…t-science/inference-ignite into do/inference-optimization-keys-2
Optimize inference signature verification flow
Signature verification policy
Whichever message arrives first (start or finish) performs cryptographic signature verification. The second message only compares its signature-bound fields against what was already persisted — no re-verification.
original_prompt_hash,request_timestamp,transfer_agent,requested_by.Field persistence for cross-message comparison
The new flow requires the first message to persist all fields that the second message will compare against. Added
OriginalPromptHash,PromptHash, andRequestedByassignments toProcessFinishInference, andOriginalPromptHashtoProcessStartInferenceininference_state.go.Also updated
calculations.startProcessed()fromPromptHash != ""toMaxTokens != 0— now thatProcessFinishInferencepopulatesPromptHash, it's no longer a reliable start-processed marker.Configurable mismatch behavior
Added
failOnCompareMismatchflag onmsgServer(defaulttrue). Set tofalseviaNewMsgServerImplWithCompareMismatchModeto log mismatches without rejecting — useful for safe rollout.