rpc-servers: Allow chainHead methods to be called from a single connection

[Bullrich]

This PR adds middleware for the JSON rpc servers to ensure that the chainHead methods are called from a single connection.

The middleware intercepts the subscription IDs generated by the chainHead_follow.
Any further chainHead function call (ie chainHead_storage) must provide a subscription ID that was generated by the middleware's connection.
A middleware per connection is created in this context.

Closes paritytech#3207

Copied from paritytech#3343

Summary by CodeRabbit

• New Features
  • Introduced a new RPC middleware for managing chainHead method calls, ensuring they are handled from a single connection for better performance and reliability.
  • Added structures to manage subscriptions and connections specifically for chainHead methods.

[coderabbitai]

Walkthrough

The modifications introduce a mechanism in Substrate's RPC layer to ensure chainHead operations are confined to a single connection. This is achieved by implementing a new middleware that manages subscription IDs and validates method calls, ensuring operations like chainHead_follow and chainHead_storage are tightly coupled to the connection that initiated them. This change enhances security and integrity of the RPC server by preventing misuse of subscription IDs across different connections.

Changes

File Path	Change Summary
.../rpc-servers/src/lib.rs	Added Arc, Mutex, ConnectionData struct, and ChainHeadLayer integration into server startup.
.../rpc-servers/src/middleware/chain_head.rs	New module with ConnectionData, ChainHeadLayer, ChainHeadMiddleware, and related RPC logic.

Assessment against linked issues

Objective	Addressed	Explanation
#3207: Ensure chainHead operations are made from a single connection	✅

Poem

In the digital meadow, under the byte-sized sun,
A rabbit hopped and toiled, until the code was done.
🐇 "With locks and layers, now secure," it said,
"ChainHead's whispers, through a single thread."
🌟 In the server's heart, connections tightly knit,
A leap for blockchain, by a rabbit's wit.
"Fear not the chaos of scattered calls," it sung,
"For in unity, our strength is truly sprung." 🚀

Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media?

Share

• X
• Mastodon
• Reddit
• LinkedIn

---

Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>.
  • Generate unit-tests for this file.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai generate unit tests for this file.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai generate interesting stats about this repository from git and render them as a table.
  • @coderabbitai show all the console.log statements in this repository.
  • @coderabbitai read src/utils.ts and generate unit tests.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (invoked as PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger a review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai help to get help.

Additionally, you can add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.

CodeRabbit Configration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• The JSON schema for the configuration file is available here.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/coderabbit-overrides.v2.json

CodeRabbit Discord Community

Join our Discord Community to get help, request features, and share feedback.

[coderabbitai]

Review Status

Actionable comments generated: 0

Configuration used: CodeRabbit UI

Commits

Files that changed from the base of the PR and between 9bf1a5e and 8e2b268.

Files ignored due to path filters (2)

• Cargo.lock is excluded by: !**/*.lock
• substrate/client/rpc-servers/Cargo.toml is excluded by: !**/*.toml

Files selected for processing (3)

• substrate/client/rpc-servers/src/lib.rs (4 hunks)
• substrate/client/rpc-servers/src/middleware/chain_head.rs (1 hunks)
• substrate/client/rpc-servers/src/middleware/mod.rs (1 hunks)

Additional comments: 11

substrate/client/rpc-servers/src/middleware/mod.rs (1)

• 21-21: The addition of the chain_head module is correctly implemented and aligns with the PR objectives to introduce middleware for chainHead operations. This change is straightforward and follows Rust's module declaration syntax and naming conventions.

substrate/client/rpc-servers/src/middleware/chain_head.rs (7)

• 37-42: The ConnectionData struct is well-defined, serving as a container for active chainHead_follow subscriptions associated with a connection. This struct is crucial for tracking subscription IDs and ensuring they are exclusive to the connection that generated them, aligning with the PR's objectives.
• 44-54: The ChainHeadLayer struct and its constructor method new are correctly implemented. This struct plays a key role in enabling the chainHead RPC methods to be called from a single connection by associating them with ConnectionData. The use of Arc<Mutex<ConnectionData>> ensures thread-safe access to the connection data.
• 65-76: The ChainHeadMiddleware struct and its constructor method new are correctly implemented. This middleware intercepts RPC requests to enforce the single-connection policy for chainHead operations. The structure and logic are consistent with the PR's objectives.
• 85-130: The logic within the call method of ChainHeadMiddleware correctly intercepts and validates subscription IDs for chainHead operations. It ensures that subscription IDs generated by chainHead_follow are used exclusively within the same connection context. This implementation effectively addresses the PR's security and integrity concerns.
• 134-155: The get_subscription_id function correctly extracts the subscription ID from both positional and named parameters. This function is essential for the middleware's operation, ensuring that the subscription ID is correctly retrieved from RPC requests for validation.
• 157-177: The get_method_result function is well-implemented, extracting the result field from the JSON-RPC response. This function is crucial for registering new subscription IDs returned by the chainHead_follow method.
• 179-237: The ResponseFuture enum and its implementation of the Future trait are correctly designed to handle different outcomes of the middleware's operation. This includes immediate error responses, forwarding calls to other layers, and registering new subscription IDs. The use of pin_project for handling the futures is appropriate and follows Rust's asynchronous programming patterns.

substrate/client/rpc-servers/src/lib.rs (3)

• 58-58: The import of ChainHeadLayer and ConnectionData from the chain_head module is correctly implemented. This import is necessary for the subsequent integration of the middleware into the server startup logic.
• 153-154: The creation of the chain_head_data instance using Arc<Mutex<ConnectionData>> is correctly implemented. This instance is essential for managing subscription IDs and their association with specific connections, aligning with the PR's objectives.
• 169-170: The integration of the ChainHeadLayer into the server configuration through the layer method is correctly implemented. This step is crucial for enforcing the single-connection policy for chainHead operations by utilizing the middleware.

[Bullrich]

@coderabbitai generate unit tests for this file.

[coderabbitai]

@Bullrich, based on the objectives and summary of the PR, here's a general approach to generating unit tests for the chain_head.rs file:

1. ConnectionData Tests: Create tests to ensure that the ConnectionData struct correctly manages subscription IDs, associating them with unique connections. This could involve testing the addition, retrieval, and deletion of subscription IDs from the ConnectionData instance.

2. ChainHeadLayer and ChainHeadMiddleware Tests: Develop tests to verify that these components intercept RPC requests as intended. This includes ensuring that they correctly identify chainHead method calls and apply the single-connection policy by validating subscription IDs against the stored connection data.

3. Subscription ID Validation Tests: Implement tests to simulate scenarios where valid and invalid subscription IDs are used in chainHead operations. This should cover cases where subscription IDs are attempted to be used across multiple connections, ensuring such attempts are blocked as per the new policy.

4. Mocking and Integration Tests: Utilize mocking frameworks to simulate the RPC server and connections. This will allow you to test the middleware's integration with the server and its behavior in a controlled environment, ensuring it operates as expected under various scenarios.

This outline should serve as a starting point for developing comprehensive unit tests for the new middleware functionality. It's important to adapt and extend these tests based on the specific logic and requirements of your implementation.