feat: Add MiddlewareContext parameter to InternalProvider.handle()

packages/eth-json-rpc-provider/CHANGELOG.md

@@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Add `providerFromMiddlewareV2` ([#7001](https://github.com/MetaMask/core/pull/7001))
   - This accepts the new middleware from `@metamask/json-rpc-engine/v2`.
+- Add `context` option to `InternalProvider.request()` ([#7056](https://github.com/MetaMask/core/pull/7056))
+  - Enables passing a `MiddlewareContext` to the JSON-RPC server.
 
 ### Changed
 

packages/eth-json-rpc-provider/src/internal-provider.test.ts

@@ -2,7 +2,10 @@ import { Web3Provider } from '@ethersproject/providers';
 import EthQuery from '@metamask/eth-query';
 import EthJsQuery from '@metamask/ethjs-query';
 import { asV2Middleware, JsonRpcEngine } from '@metamask/json-rpc-engine';
-import type { JsonRpcMiddleware } from '@metamask/json-rpc-engine/v2';
+import type {
+  JsonRpcMiddleware,
+  MiddlewareContext,
+} from '@metamask/json-rpc-engine/v2';
 import { JsonRpcEngineV2 } from '@metamask/json-rpc-engine/v2';
 import { providerErrors, rpcErrors } from '@metamask/rpc-errors';
 import { type JsonRpcRequest, type Json } from '@metamask/utils';
@@ -16,7 +19,9 @@ import {
 
 jest.mock('uuid');
 
-type ResultParam = Json | ((req?: JsonRpcRequest) => Json);
+type ResultParam =
+  | Json
+  | ((req?: JsonRpcRequest, context?: MiddlewareContext) => Json);
 
 const createLegacyEngine = (method: string, result: ResultParam) => {
   const engine = new JsonRpcEngine();
@@ -33,9 +38,11 @@ const createLegacyEngine = (method: string, result: ResultParam) => {
 const createV2Engine = (method: string, result: ResultParam) => {
   return JsonRpcEngineV2.create<JsonRpcMiddleware<JsonRpcRequest>>({
     middleware: [
-      ({ request, next }) => {
+      ({ request, next, context }) => {
         if (request.method === method) {
-          return typeof result === 'function' ? result(request) : result;
+          return typeof result === 'function'
+            ? result(request as JsonRpcRequest, context)
+            : result;
         }
         return next();
       },
@@ -245,6 +252,29 @@ describe.each([
       expect(response.result).toBe(42);
     });
 
+    it('forwards the context to the JSON-RPC handler', async () => {
+      const rpcHandler = createRpcHandler('test', (request, context) => {
+        // @ts-expect-error - Intentional type abuse.
+        // eslint-disable-next-line jest/no-conditional-in-test
+        return context?.assertGet('foo') ?? request.foo;
+      });
+      const provider = new InternalProvider({ engine: rpcHandler });
+
+      const request = {
+        id: 1,
+        jsonrpc: '2.0' as const,
+        method: 'test',
+      };
+
+      const result = await provider.request(request, {
+        context: {
+          foo: 'bar',
+        },
+      });
+
+      expect(result).toBe('bar');
+    });
+
     it('handles a successful EIP-1193 object request', async () => {
       let req: JsonRpcRequest | undefined;
       const rpcHandler = createRpcHandler('test', (request) => {

packages/eth-json-rpc-provider/src/internal-provider.ts

@@ -1,9 +1,10 @@
 import { asV2Middleware, type JsonRpcEngine } from '@metamask/json-rpc-engine';
-import type {
-  ContextConstraint,
-  MiddlewareContext,
+import {
+  type HandleOptions,
+  type ContextConstraint,
+  type MiddlewareContext,
+  JsonRpcEngineV2,
 } from '@metamask/json-rpc-engine/v2';
-import { JsonRpcEngineV2 } from '@metamask/json-rpc-engine/v2';
 import type { JsonRpcSuccess } from '@metamask/utils';
 import {
   type Json,
@@ -37,16 +38,18 @@ type Options<
  * This provider loosely follows conventions that pre-date EIP-1193.
  * It is not compliant with any Ethereum provider standard.
  */
-export class InternalProvider {
-  readonly #engine: JsonRpcEngineV2<JsonRpcRequest, MiddlewareContext>;
+export class InternalProvider<
+  Context extends ContextConstraint = MiddlewareContext,
+> {
+  readonly #engine: JsonRpcEngineV2<JsonRpcRequest, Context>;
 
   /**
    * Construct a InternalProvider from a JSON-RPC server or legacy engine.
    *
    * @param options - Options.
    * @param options.engine - The JSON-RPC engine used to process requests.
    */
-  constructor({ engine }: Options) {
+  constructor({ engine }: Options<JsonRpcRequest, Context>) {
     this.#engine =
       'push' in engine
         ? JsonRpcEngineV2.create({
@@ -59,14 +62,17 @@ export class InternalProvider {
    * Send a provider request asynchronously.
    *
    * @param eip1193Request - The request to send.
+   * @param options - The options for the request operation.
+   * @param options.context - The context to pass to the server.
    * @returns The JSON-RPC response.
    */
   async request<Params extends JsonRpcParams, Result extends Json>(
     eip1193Request: Eip1193Request<Params>,
+    options?: HandleOptions<Context>,
   ): Promise<Result> {
     const jsonRpcRequest =
       convertEip1193RequestToJsonRpcRequest(eip1193Request);
-    return (await this.#handle<Result>(jsonRpcRequest)).result;
+    return (await this.#handle<Result>(jsonRpcRequest, options)).result;
   }
 
   /**
@@ -118,12 +124,14 @@ export class InternalProvider {
 
   readonly #handle = async <Result extends Json>(
     jsonRpcRequest: JsonRpcRequest,
+    options?: HandleOptions<Context>,
   ): Promise<JsonRpcSuccess<Result>> => {
     const { id, jsonrpc } = jsonRpcRequest;
-    // This typecast is technicaly unsafe, but we need it to preserve the provider's
-    // public interface, which allows you to typecast results.
+    // The `result` typecast is unsafe, but we need it to preserve the provider's
+    // public interface, which allows you to unsafely typecast results.
     const result = (await this.#engine.handle(
       jsonRpcRequest,
+      options,
     )) as unknown as Result;
 
     return {

packages/json-rpc-engine/CHANGELOG.md

@@ -9,7 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
-- Add `JsonRpcEngineV2` ([#6176](https://github.com/MetaMask/core/pull/6176), [#6971](https://github.com/MetaMask/core/pull/6971), [#6975](https://github.com/MetaMask/core/pull/6975), [#6990](https://github.com/MetaMask/core/pull/6990), [#6991](https://github.com/MetaMask/core/pull/6991), [#7001](https://github.com/MetaMask/core/pull/7001), [#7032](https://github.com/MetaMask/core/pull/7032))
+- Add `JsonRpcEngineV2` ([#6176](https://github.com/MetaMask/core/pull/6176), [#6971](https://github.com/MetaMask/core/pull/6971), [#6975](https://github.com/MetaMask/core/pull/6975), [#6990](https://github.com/MetaMask/core/pull/6990), [#6991](https://github.com/MetaMask/core/pull/6991), [#7001](https://github.com/MetaMask/core/pull/7001), [#7032](https://github.com/MetaMask/core/pull/7032), [#7056](https://github.com/MetaMask/core/pull/7056))
   - This is a complete rewrite of `JsonRpcEngine`, intended to replace the original implementation. See the readme for details.
 
 ## [10.1.1]

packages/json-rpc-engine/src/v2/JsonRpcEngineV2.test.ts

@@ -371,6 +371,24 @@ describe('JsonRpcEngineV2', () => {
         expect(result).toBe('bar');
       });
 
+      it('accepts an initial context as a KeyValues object', async () => {
+        const initialContext = { foo: 'bar' } as const;
+        const middleware: JsonRpcMiddleware<
+          JsonRpcRequest,
+          string,
+          MiddlewareContext<Record<string, string>>
+        > = ({ context }) => context.assertGet('foo');
+        const engine = JsonRpcEngineV2.create({
+          middleware: [middleware],
+        });
+
+        const result = await engine.handle(makeRequest(), {
+          context: initialContext,
+        });
+
+        expect(result).toBe('bar');
+      });
+
       it('accepts middleware with different context types', async () => {
         const middleware1: JsonRpcMiddleware<
           JsonRpcCall,

packages/json-rpc-engine/src/v2/JsonRpcEngineV2.ts

@@ -7,7 +7,11 @@ import {
 } from '@metamask/utils';
 import deepFreeze from 'deep-freeze-strict';
 
-import type { ContextConstraint, MergeContexts } from './MiddlewareContext';
+import type {
+  ContextConstraint,
+  InferKeyValues,
+  MergeContexts,
+} from './MiddlewareContext';
 import { MiddlewareContext } from './MiddlewareContext';
 import {
   isNotification,
@@ -60,8 +64,11 @@ type RequestState<Request extends JsonRpcCall> = {
   result: Readonly<ResultConstraint<Request>> | undefined;
 };
 
-type HandleOptions<Context extends MiddlewareContext> = {
-  context?: Context;
+/**
+ * The options for the JSON-RPC request/notification handling operation.
+ */
+export type HandleOptions<Context extends ContextConstraint> = {
+  context?: Context | InferKeyValues<Context>;
 };
 
 type ConstructorOptions<
@@ -286,12 +293,14 @@ export class JsonRpcEngineV2<
    * operation. Permits returning an `undefined` result.
    *
    * @param originalRequest - The JSON-RPC request to handle.
-   * @param context - The context to pass to the middleware.
+   * @param rawContext - The context to pass to the middleware.
    * @returns The result from the middleware.
    */
   async #handle(
     originalRequest: Request,
-    context: Context = new MiddlewareContext() as Context,
+    rawContext:
+      | Context
+      | InferKeyValues<Context> = new MiddlewareContext() as Context,
   ): Promise<RequestState<Request>> {
     this.#assertIsNotDestroyed();
 
@@ -303,6 +312,9 @@ export class JsonRpcEngineV2<
     };
     const middlewareIterator = this.#makeMiddlewareIterator();
     const firstMiddleware = middlewareIterator.next().value;
+    const context = MiddlewareContext.isInstance(rawContext)
+      ? rawContext
+      : (new MiddlewareContext(rawContext) as Context);
 
     const makeNext = this.#makeNextFactory(middlewareIterator, state, context);
 

packages/json-rpc-engine/src/v2/JsonRpcServer.test.ts

@@ -3,6 +3,7 @@ import { rpcErrors } from '@metamask/rpc-errors';
 import type { JsonRpcMiddleware } from './JsonRpcEngineV2';
 import { JsonRpcEngineV2 } from './JsonRpcEngineV2';
 import { JsonRpcServer } from './JsonRpcServer';
+import type { MiddlewareContext } from './MiddlewareContext';
 import type { JsonRpcNotification, JsonRpcRequest } from './utils';
 import { isRequest, JsonRpcEngineError, stringify } from './utils';
 
@@ -108,6 +109,39 @@ describe('JsonRpcServer', () => {
     expect(response).toBeUndefined();
   });
 
+  it('forwards the context to the engine', async () => {
+    const middleware: JsonRpcMiddleware<
+      JsonRpcRequest,
+      string,
+      MiddlewareContext<{ foo: string }>
+    > = ({ context }) => {
+      return context.assertGet('foo');
+    };
+    const server = new JsonRpcServer({
+      middleware: [middleware],
+      onError: () => undefined,
+    });
+
+    const response = await server.handle(
+      {
+        jsonrpc,
+        id: 1,
+        method: 'hello',
+      },
+      {
+        context: {
+          foo: 'bar',
+        },
+      },
+    );
+
+    expect(response).toStrictEqual({
+      jsonrpc,
+      id: 1,
+      result: 'bar',
+    });
+  });
+
   it('returns an error response for a failed request', async () => {
     const server = new JsonRpcServer({
       engine: makeEngine(),

packages/json-rpc-engine/src/v2/JsonRpcServer.ts

@@ -9,6 +9,7 @@ import type {
 import { hasProperty, isObject } from '@metamask/utils';
 
 import type {
+  HandleOptions,
   JsonRpcMiddleware,
   MergedContextOf,
   MiddlewareConstraint,
@@ -102,9 +103,14 @@ export class JsonRpcServer<
    * engine. The request will fail if the engine can only handle notifications.
    *
    * @param request - The request to handle.
+   * @param options - The options for the handle operation.
+   * @param options.context - The context to pass to the middleware.
    * @returns The JSON-RPC response.
    */
-  async handle(request: JsonRpcRequest): Promise<JsonRpcResponse>;
+  async handle(
+    request: JsonRpcRequest,
+    options?: HandleOptions<MergedContextOf<Middleware>>,
+  ): Promise<JsonRpcResponse>;
 
   /**
    * Handle a JSON-RPC notification.
@@ -116,8 +122,13 @@ export class JsonRpcServer<
    * engine. The request will fail if the engine cannot handle notifications.
    *
    * @param notification - The notification to handle.
+   * @param options - The options for the handle operation.
+   * @param options.context - The context to pass to the middleware.
    */
-  async handle(notification: JsonRpcNotification): Promise<void>;
+  async handle(
+    notification: JsonRpcNotification,
+    options?: HandleOptions<MergedContextOf<Middleware>>,
+  ): Promise<void>;
 
   /**
    * Handle an alleged JSON-RPC request or notification. Permits any plain
@@ -133,12 +144,20 @@ export class JsonRpcServer<
    * response) is not of the type expected by the underlying engine.
    *
    * @param rawRequest - The raw request to handle.
+   * @param options - The options for the handle operation.
+   * @param options.context - The context to pass to the middleware.
    * @returns The JSON-RPC response, or `undefined` if the request is a
    * notification.
    */
-  async handle(rawRequest: unknown): Promise<JsonRpcResponse | void>;
+  async handle(
+    rawRequest: unknown,
+    options?: HandleOptions<MergedContextOf<Middleware>>,
+  ): Promise<JsonRpcResponse | void>;
 
-  async handle(rawRequest: unknown): Promise<JsonRpcResponse | void> {
+  async handle(
+    rawRequest: unknown,
+    options?: HandleOptions<MergedContextOf<Middleware>>,
+  ): Promise<JsonRpcResponse | void> {
     // If rawRequest is not a notification, the originalId will be attached
     // to the response. We attach our own, trusted id in #coerceRequest()
     // while the request is being handled.
@@ -148,7 +167,7 @@ export class JsonRpcServer<
       const request = this.#coerceRequest(rawRequest, isRequest);
       // @ts-expect-error - The request may not be of the type expected by the engine,
       // and we intentionally allow this to happen.
-      const result = await this.#engine.handle(request);
+      const result = await this.#engine.handle(request, options);
 
       if (result !== undefined) {
         return {

packages/json-rpc-engine/src/v2/MiddlewareContext.test.ts

@@ -11,6 +11,11 @@ describe('MiddlewareContext', () => {
     expect(context.get(symbol)).toBe('value');
   });
 
+  it('can be constructed with a KeyValues object', () => {
+    const context = new MiddlewareContext<{ test: string }>({ test: 'value' });
+    expect(context.get('test')).toBe('value');
+  });
+
   it('is frozen', () => {
     const context = new MiddlewareContext();
     expect(Object.isFrozen(context)).toBe(true);