fix: remove mock signature from chat signature in production code

[hanakannzashi]

Fixes #308

[gemini-code-assist]

Summary of Changes

Hello @hanakannzashi, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request refactors the handling of chat completion request and response signature hashes, specifically for the MockProvider. The changes streamline the signature generation process by centralizing hash computation within the MockProvider itself, eliminating the need for the InferenceProviderPool to store these hashes and simplifying the AttestationService's logic. This ensures that mock provider signatures are generated consistently and deterministically, aligning their behavior more closely with real inference providers.

Highlights

• Refactored Signature Hash Management: The responsibility for computing and registering request and response hashes for chat completion signatures has been moved from the InferenceProviderPool and API route handler directly into the MockProvider. This centralizes hash generation logic.
• Simplified Attestation Service: The AttestationService no longer contains special conditional logic for handling MockProvider signatures. It now consistently requests signatures from the inference provider, regardless of whether it's a mock or real provider.
• Introduced SHA2 Hashing: The sha2 crate has been integrated into the MockProvider to deterministically compute SHA256 hashes for both request and response bodies, ensuring consistent and verifiable mock signatures.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request refactors the signature generation logic for the mock inference provider. The responsibility for calculating response hashes and preparing signature data is moved from the API route handlers and attestation service into the MockProvider itself. This is a great simplification, making the mock provider more self-contained and removing special-case logic from the core services.

My review focuses on the new logic within the MockProvider. I've identified a couple of places where unwrap() is used for JSON serialization, which could lead to panics during tests. I've suggested replacing these with proper error handling to improve robustness and provide clearer test failure messages, consistent with how errors are handled elsewhere in the codebase.

[claude]

Code Review

⚠️ CRITICAL ISSUE - Hash Mismatch in Disconnect Scenario

Location: crates/inference_providers/src/mock.rs:684-704

Problem: When disconnect_after_chunks is set, the code computes a signature hash that does not match what is actually streamed to the client:

// Line 686-687: chunks are truncated
if let Some(disconnect_at) = response_template.disconnect_after_chunks {
    chunks.truncate(disconnect_at);
}

// Line 695-701: hash is computed over TRUNCATED chunks + "[DONE]"
for chunk in &chunks {
    let json = serde_json::to_value(chunk).unwrap();
    let raw_bytes = Self::sse_data_static(&json);
    accumulated.extend_from_slice(&raw_bytes);
}
accumulated.extend_from_slice(b"data: [DONE]\n\n");  // ❌ WRONG!

Why This is Critical:

1. Security Issue: The signature hash should match what the client receives byte-for-byte
2. Disconnect simulation intentionally ends the stream early WITHOUT the [DONE] marker (see comment on line 685: "The stream will end abruptly without the final usage chunk")
3. Signature verification will fail: Any client trying to verify the signature will see a mismatch between the actual bytes received and the hash

Fix:

// Register signature hashes for this chat_id.
let chat_id_opt = chunks.first().map(|c| c.id.clone());
if let Some(chat_id) = chat_id_opt {
    let mut accumulated: Vec<u8> = Vec::new();
    for chunk in &chunks {
        let json = serde_json::to_value(chunk).unwrap();
        let raw_bytes = Self::sse_data_static(&json);
        accumulated.extend_from_slice(&raw_bytes);
    }
    
    // Only add [DONE] if disconnect simulation is NOT active
    if response_template.disconnect_after_chunks.is_none() {
        accumulated.extend_from_slice(b"data: [DONE]\n\n");
    }
    
    let response_hash = compute_sha256_hex(&accumulated);
    self.register_signature_hashes(chat_id, request_hash, response_hash)
        .await;
}

---

Other Observations (Non-blocking):

✅ Good refactoring: Moving hash computation into MockProvider improves separation of concerns and eliminates the gateway-side hash tracking complexity.

✅ Cleanup: Removed 136 lines of code by eliminating signature_hashes from InferenceProviderPool and special-case logic from AttestationService.

✅ Deterministic signatures: Using SHA256 for hash computation is appropriate for mock testing.

⚠️ Minor: The .unwrap() calls on lines 697, 709, 792 are acceptable since these are internal mock structures that should always serialize successfully. However, consider if these should panic or return errors more gracefully in a test environment.

---

Recommendation: Fix the disconnect scenario hash computation before merging.

[think-in-universe]

The reported security vulnerability is from ruint v0.17.0 (#304 / RUSTSEC-2025-0137), not related to the change in this PR.

We're waiting for a patch from ruint to fix this issue: recmo/uint#550

[Copilot]

Pull request overview

This PR refactors mock signature generation by moving the responsibility from the gateway layer to the MockProvider itself, simplifying the architecture and removing MockProvider-specific logic from the AttestationService.

Key changes:

• Signature hash registration moved from gateway API route to MockProvider implementation
• MockProvider now computes response hashes internally during chat_completion and chat_completion_stream
• Removed gateway-side signature synthesis logic from AttestationService

Reviewed changes

Copilot reviewed 5 out of 6 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
crates/services/src/inference_provider_pool/mod.rs	Removed signature_hashes field and associated methods (register_signature_hashes_for_chat, get_signature_hashes_for_chat) that were used for gateway-side hash tracking
crates/services/src/attestation/mod.rs	Removed MockProvider-specific signature synthesis logic; now always fetches signatures directly from the provider
crates/inference_providers/src/mock.rs	Added compute_sha256_hex helper and signature hash registration in both streaming and non-streaming chat completion methods; changed request_hash parameter from unused to used
crates/api/src/routes/completions.rs	Removed gateway-side response hash computation and signature registration code for streaming completions
crates/inference_providers/Cargo.toml	Added sha2 dependency for hash computation in MockProvider
Cargo.lock	Updated with sha2 dependency

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[think-in-universe]

Not relevant to the change in the PR, but I think there're several causes of the issue: @nickpismenkov @PierreLeGuen

1. We're mixing code for testing purpose with production code. Can we add test feature flag for code that is used for testing only?
2. Now all the model providers are mocked. There's no end-to-end tests against the real vllm-proxy instances, which makes it fail to catch some issues like signature signing address issue and leak the bugs into staging or production.

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.