Add requestGate handler option

packages/connect/src/protocol-connect/handler-factory.spec.ts

@@ -36,13 +36,16 @@ import {
   sinkAll,
   transformSplitEnvelope,
 } from "../protocol/index.js";
-import { Code, ConnectError } from "../index.js";
+import { Code, ConnectError, createContextKey } from "../index.js";
 import { errorFromJsonBytes } from "./error-json.js";
 import { endStreamFromJson } from "./end-stream.js";
 import { createTransport } from "./transport.js";
 import { requestHeader } from "./request-header.js";
 import { readAll } from "../protocol/async-iterable-helper.spec.js";
-import { contentTypeStreamProto } from "./content-type.js";
+import {
+  contentTypeStreamProto,
+  contentTypeUnaryProto,
+} from "./content-type.js";
 import { createServiceDesc } from "../descriptor-helper.spec.js";
 import {
   ApiSchema,
@@ -689,4 +692,97 @@ describe("createHandlerFactory()", () => {
       );
     });
   });
+
+  describe("requestGate", () => {
+    const denyAll = () => {
+      throw new ConnectError("no credentials", Code.Unauthenticated);
+    };
+    // A request body that records whether the handler read from it.
+    function trackedBody(state: { read: boolean }): AsyncIterable<Uint8Array> {
+      return {
+        [Symbol.asyncIterator]: () => ({
+          next: () => {
+            state.read = true;
+            return Promise.resolve({ done: true, value: undefined });
+          },
+        }),
+      };
+    }
+
+    for (const method of [
+      testService.method.unary,
+      testService.method.serverStreaming,
+    ]) {
+      it(`should not read the request body of a ${method.methodKind} RPC`, async () => {
+        const { handler } = setupTestHandler(
+          method,
+          { requestGate: denyAll },
+          () => assert.fail("implementation should not be called"),
+        );
+        const state = { read: false };
+        const res = await handler({
+          httpVersion: "2.0",
+          method: "POST",
+          url: `https://example.com/${method.parent.typeName}/${method.name}`,
+          header: new Headers({
+            "Content-Type":
+              method.methodKind === "unary"
+                ? contentTypeUnaryProto
+                : contentTypeStreamProto,
+          }),
+          body: trackedBody(state),
+          signal: new AbortController().signal,
+        });
+        assert.notStrictEqual(res.status, 415); // wrong content-type for this RPC
+        assert.strictEqual(state.read, false);
+      });
+    }
+
+    it("should reject a streaming RPC with the error from the gate", async () => {
+      const { transport, method } = setupTestHandler(
+        testService.method.serverStreaming,
+        { requestGate: denyAll },
+        () => assert.fail("implementation should not be called"),
+      );
+      await assert.rejects(
+        async () => {
+          const r = await transport.stream(
+            method,
+            undefined,
+            undefined,
+            undefined,
+            createAsyncIterable([create(Int32ValueSchema)]),
+          );
+          await pipeTo(r.message, sinkAll());
+        },
+        (e) => {
+          assert.ok(e instanceof ConnectError);
+          assert.strictEqual(e.code, Code.Unauthenticated);
+          return true;
+        },
+      );
+    });
+
+    it("should pass context values from the gate to the implementation", async () => {
+      const kUser = createContextKey<string>("anonymous");
+      const { transport, method } = setupTestHandler(
+        testService.method.unary,
+        {
+          requestGate: (ctx) => {
+            assert.strictEqual(ctx.requestHeader.get("authorization"), "token");
+            ctx.values.set(kUser, "alice");
+          },
+        },
+        (_req, ctx) => ({ value: ctx.values.get(kUser) }),
+      );
+      const res = await transport.unary(
+        method,
+        undefined,
+        undefined,
+        { authorization: "token" },
+        create(Int32ValueSchema, { value: 1 }),
+      );
+      assert.strictEqual(res.message.value, "alice");
+    });
+  });
 });

packages/connect/src/protocol-connect/handler-factory.ts

@@ -98,6 +98,7 @@ import { contentTypeMatcher } from "../protocol/content-type-matcher.js";
 import { createMethodUrl } from "../protocol/create-method-url.js";
 import type { EnvelopedMessage } from "../protocol/envelope.js";
 import {
+  applyRequestGate,
   invokeUnaryImplementation,
   transformInvokeImplementation,
 } from "../protocol/invoke-implementation.js";
@@ -216,6 +217,11 @@ function createUnaryHandler<I extends DescMessage, O extends DescMessage>(
     let status = uResponseOk.status;
     let body: Uint8Array;
     try {
+      // We run the request gate before receiving the body, so that a request it
+      // rejects is never read, decompressed, or parsed.
+      if (opt.requestGate !== undefined) {
+        await opt.requestGate(context);
+      }
       if (opt.requireConnectProtocolHeader) {
         if (isGet) {
           requireProtocolVersionParam(queryParams);
@@ -440,11 +446,15 @@ function createStreamHandler<I extends DescMessage, O extends DescMessage>(
         // raises an error, but we want to be lenient
       ),
     );
-    const it = transformInvokeImplementation<I, O>(
-      spec,
-      context,
-      opt.interceptors,
-    )(inputIt)[Symbol.asyncIterator]();
+    // We run the request gate before receiving the body, so that a request it
+    // rejects is never read, decompressed, or parsed.
+    const it = await applyRequestGate(context, opt.requestGate, () =>
+      transformInvokeImplementation<I, O>(
+        spec,
+        context,
+        opt.interceptors,
+      )(inputIt),
+    );
     const outputIt = pipe(
       // We wrap the iterator in an async iterator to ensure that the
       // abort signal is aborted when the iterator is done.

packages/connect/src/protocol-grpc-web/handler-factory.spec.ts

@@ -27,6 +27,7 @@ import {
   pipeTo,
   sinkAll,
 } from "../protocol/index.js";
+import { createContextKey } from "../index.js";
 import { createHandlerFactory } from "./handler-factory.js";
 import { createTransport } from "./transport.js";
 import { requestHeader } from "./request-header.js";
@@ -301,4 +302,92 @@ describe("createHandlerFactory()", () => {
       assert.strictEqual(handlerContextSignal?.reason, "test-reason");
     });
   });
+
+  describe("requestGate", () => {
+    const denyAll = () => {
+      throw new ConnectError("no credentials", Code.Unauthenticated);
+    };
+    // A request body that records whether the handler read from it.
+    function trackedBody(state: { read: boolean }): AsyncIterable<Uint8Array> {
+      return {
+        [Symbol.asyncIterator]: () => ({
+          next: () => {
+            state.read = true;
+            return Promise.resolve({ done: true, value: undefined });
+          },
+        }),
+      };
+    }
+
+    for (const method of [
+      testService.method.unary,
+      testService.method.serverStreaming,
+    ]) {
+      it(`should not read the request body of a ${method.methodKind} RPC`, async () => {
+        const { handler } = setupTestHandler(
+          method,
+          { requestGate: denyAll },
+          () => assert.fail("implementation should not be called"),
+        );
+        const state = { read: false };
+        const res = await handler({
+          httpVersion: "2.0",
+          method: "POST",
+          url: `https://example.com/${method.parent.typeName}/${method.name}`,
+          header: new Headers({ "Content-Type": contentTypeProto }),
+          body: trackedBody(state),
+          signal: new AbortController().signal,
+        });
+        assert.notStrictEqual(res.status, 415); // wrong content-type for this RPC
+        assert.strictEqual(state.read, false);
+      });
+    }
+
+    it("should reject a streaming RPC with the error from the gate", async () => {
+      const { transport, method } = setupTestHandler(
+        testService.method.serverStreaming,
+        { requestGate: denyAll },
+        () => assert.fail("implementation should not be called"),
+      );
+      await assert.rejects(
+        async () => {
+          const r = await transport.stream(
+            method,
+            undefined,
+            undefined,
+            undefined,
+            createAsyncIterable([create(Int32ValueSchema)]),
+          );
+          await pipeTo(r.message, sinkAll());
+        },
+        (e) => {
+          assert.ok(e instanceof ConnectError);
+          assert.strictEqual(e.code, Code.Unauthenticated);
+          return true;
+        },
+      );
+    });
+
+    it("should pass context values from the gate to the implementation", async () => {
+      const kUser = createContextKey<string>("anonymous");
+      const { transport, method } = setupTestHandler(
+        testService.method.unary,
+        {
+          requestGate: (ctx) => {
+            assert.strictEqual(ctx.requestHeader.get("authorization"), "token");
+            ctx.values.set(kUser, "alice");
+          },
+        },
+        (_req, ctx) => ({ value: ctx.values.get(kUser) }),
+      );
+      const res = await transport.unary(
+        method,
+        undefined,
+        undefined,
+        { authorization: "token" },
+        create(Int32ValueSchema, { value: 1 }),
+      );
+      assert.strictEqual(res.message.value, "alice");
+    });
+  });
 });

packages/connect/src/protocol-grpc-web/handler-factory.ts

@@ -52,7 +52,10 @@ import { compressionNegotiate } from "../protocol/compression.js";
 import { contentTypeMatcher } from "../protocol/content-type-matcher.js";
 import { createMethodUrl } from "../protocol/create-method-url.js";
 import type { EnvelopedMessage } from "../protocol/envelope.js";
-import { transformInvokeImplementation } from "../protocol/invoke-implementation.js";
+import {
+  applyRequestGate,
+  transformInvokeImplementation,
+} from "../protocol/invoke-implementation.js";
 import type { ProtocolHandlerFactory } from "../protocol/protocol-handler-factory.js";
 import { createMethodSerializationLookup } from "../protocol/serialization.js";
 import type { Serialization } from "../protocol/serialization.js";
@@ -172,11 +175,15 @@ function createHandler<I extends DescMessage, O extends DescMessage>(
         // raises an error, but we want to be lenient
       ),
     );
-    const it = transformInvokeImplementation<I, O>(
-      spec,
-      context,
-      opt.interceptors,
-    )(inputIt)[Symbol.asyncIterator]();
+    // We run the request gate before receiving the body, so that a request it
+    // rejects is never read, decompressed, or parsed.
+    const it = await applyRequestGate(context, opt.requestGate, () =>
+      transformInvokeImplementation<I, O>(
+        spec,
+        context,
+        opt.interceptors,
+      )(inputIt),
+    );
     const outputIt = pipe(
       // We wrap the iterator in an async iterator to ensure that the
       // abort signal is aborted when the iterator is done.

packages/connect/src/protocol-grpc/handler-factory.spec.ts

@@ -28,7 +28,7 @@ import {
 import { createHandlerFactory } from "./handler-factory.js";
 import { createTransport } from "./transport.js";
 import { requestHeader } from "./request-header.js";
-import { Code, ConnectError } from "../index.js";
+import { Code, ConnectError, createContextKey } from "../index.js";
 import { contentTypeProto } from "./content-type.js";
 import { createServiceDesc } from "../descriptor-helper.spec.js";
 import { Int32ValueSchema, StringValueSchema } from "@bufbuild/protobuf/wkt";
@@ -300,4 +300,92 @@ describe("createHandlerFactory()", () => {
       assert.strictEqual(handlerContextSignal?.reason, "test-reason");
     });
   });
+
+  describe("requestGate", () => {
+    const denyAll = () => {
+      throw new ConnectError("no credentials", Code.Unauthenticated);
+    };
+    // A request body that records whether the handler read from it.
+    function trackedBody(state: { read: boolean }): AsyncIterable<Uint8Array> {
+      return {
+        [Symbol.asyncIterator]: () => ({
+          next: () => {
+            state.read = true;
+            return Promise.resolve({ done: true, value: undefined });
+          },
+        }),
+      };
+    }
+
+    for (const method of [
+      testService.method.unary,
+      testService.method.serverStreaming,
+    ]) {
+      it(`should not read the request body of a ${method.methodKind} RPC`, async () => {
+        const { handler } = setupTestHandler(
+          method,
+          { requestGate: denyAll },
+          () => assert.fail("implementation should not be called"),
+        );
+        const state = { read: false };
+        const res = await handler({
+          httpVersion: "2.0",
+          method: "POST",
+          url: `https://example.com/${method.parent.typeName}/${method.name}`,
+          header: new Headers({ "Content-Type": contentTypeProto }),
+          body: trackedBody(state),
+          signal: new AbortController().signal,
+        });
+        assert.notStrictEqual(res.status, 415); // wrong content-type for this RPC
+        assert.strictEqual(state.read, false);
+      });
+    }
+
+    it("should reject a streaming RPC with the error from the gate", async () => {
+      const { transport, method } = setupTestHandler(
+        testService.method.serverStreaming,
+        { requestGate: denyAll },
+        () => assert.fail("implementation should not be called"),
+      );
+      await assert.rejects(
+        async () => {
+          const r = await transport.stream(
+            method,
+            undefined,
+            undefined,
+            undefined,
+            createAsyncIterable([create(Int32ValueSchema)]),
+          );
+          await pipeTo(r.message, sinkAll());
+        },
+        (e) => {
+          assert.ok(e instanceof ConnectError);
+          assert.strictEqual(e.code, Code.Unauthenticated);
+          return true;
+        },
+      );
+    });
+
+    it("should pass context values from the gate to the implementation", async () => {
+      const kUser = createContextKey<string>("anonymous");
+      const { transport, method } = setupTestHandler(
+        testService.method.unary,
+        {
+          requestGate: (ctx) => {
+            assert.strictEqual(ctx.requestHeader.get("authorization"), "token");
+            ctx.values.set(kUser, "alice");
+          },
+        },
+        (_req, ctx) => ({ value: ctx.values.get(kUser) }),
+      );
+      const res = await transport.unary(
+        method,
+        undefined,
+        undefined,
+        { authorization: "token" },
+        create(Int32ValueSchema, { value: 1 }),
+      );
+      assert.strictEqual(res.message.value, "alice");
+    });
+  });
 });