Add WithRequestGate handler option

handler.go

@@ -28,6 +28,7 @@ import (
 type Handler struct {
 	spec             Spec
 	implementation   StreamingHandlerFunc
+	requestGates     []func(ctx context.Context, spec Spec, peer Peer, header http.Header) (context.Context, error)
 	protocolHandlers map[string][]protocolHandler // Method to protocol handlers
 	allowMethod      string                       // Allow header
 	acceptPost       string                       // Accept-Post header
@@ -105,6 +106,7 @@ func NewUnaryHandler[Req, Res any](
 	return &Handler{
 		spec:             config.newSpec(),
 		implementation:   implementation,
+		requestGates:     config.RequestGates,
 		protocolHandlers: mappedMethodHandlers(protocolHandlers),
 		allowMethod:      sortedAllowMethodValue(protocolHandlers),
 		acceptPost:       sortedAcceptPostValue(protocolHandlers),
@@ -334,6 +336,18 @@ func (h *Handler) ServeHTTP(responseWriter http.ResponseWriter, request *http.Re
 		_ = connCloser.Close(timeoutErr)
 		return
 	}
+	// Gates run before the implementation, so before the interceptor chain and
+	// before any message is received.
+	for _, gate := range h.requestGates {
+		gateCtx, err := gate(ctx, h.spec, connCloser.Peer(), connCloser.RequestHeader())
+		if err != nil {
+			_ = connCloser.Close(err)
+			return
+		}
+		if gateCtx != nil {
+			ctx = gateCtx //nolint:fatcontext // Chaining is intended: each gate sees the previous gate's context.
+		}
+	}
 	_ = connCloser.Close(h.implementation(ctx, connCloser))
 }
 
@@ -352,6 +366,7 @@ type handlerConfig struct {
 	ReadMaxBytes                 int
 	SendMaxBytes                 int
 	StreamType                   StreamType
+	RequestGates                 []func(ctx context.Context, spec Spec, peer Peer, header http.Header) (context.Context, error)
 }
 
 func newHandlerConfig(procedure string, streamType StreamType, options []HandlerOption) *handlerConfig {
@@ -420,6 +435,7 @@ func newStreamHandler(
 	return &Handler{
 		spec:             config.newSpec(),
 		implementation:   implementation,
+		requestGates:     config.RequestGates,
 		protocolHandlers: mappedMethodHandlers(protocolHandlers),
 		allowMethod:      sortedAllowMethodValue(protocolHandlers),
 		acceptPost:       sortedAcceptPostValue(protocolHandlers),

handler_ext_test.go

@@ -568,6 +568,127 @@ func TestDynamicHandler(t *testing.T) {
 	})
 }
 
+func TestHandlerWithRequestGate(t *testing.T) {
+	t.Parallel()
+	errRejected := connect.NewError(connect.CodeUnauthenticated, errors.New("no credentials"))
+	const tokenHeader = "Test-Token"
+
+	t.Run("rejects_before_decode", func(t *testing.T) {
+		t.Parallel()
+		var handlerCalls, gateCalls int
+		var gateSpec connect.Spec
+		var gatePeer connect.Peer
+		var gateToken string
+		mux := http.NewServeMux()
+		mux.Handle(pingv1connect.NewPingServiceHandler(
+			&pluggablePingServer{ping: func(context.Context, *connect.Request[pingv1.PingRequest]) (*connect.Response[pingv1.PingResponse], error) {
+				handlerCalls++
+				return connect.NewResponse(&pingv1.PingResponse{}), nil
+			}},
+			connect.WithRequestGate(func(ctx context.Context, spec connect.Spec, peer connect.Peer, header http.Header) (context.Context, error) {
+				gateCalls++
+				gateSpec, gatePeer, gateToken = spec, peer, header.Get(tokenHeader)
+				return nil, errRejected
+			}),
+		))
+		server := memhttptest.NewServer(t, mux)
+
+		// A body that cannot be unmarshaled. Without a gate this request gets
+		// 400 invalid_argument from the codec, so a 401 proves that the gate ran
+		// before the receive.
+		request, err := http.NewRequestWithContext(
+			t.Context(),
+			http.MethodPost,
+			server.URL()+pingv1connect.PingServicePingProcedure,
+			bytes.NewReader([]byte("invalid message")),
+		)
+		assert.Nil(t, err)
+		request.Header.Set("Content-Type", "application/proto")
+		request.Header.Set(tokenHeader, "sesame")
+		response, err := server.Client().Do(request)
+		assert.Nil(t, err)
+		defer response.Body.Close()
+		body, err := io.ReadAll(response.Body)
+		assert.Nil(t, err)
+		assert.Equal(t, response.StatusCode, http.StatusUnauthorized)
+		assert.True(t, strings.Contains(string(body), "no credentials"), assert.Sprintf("body: %s", body))
+		assert.Equal(t, gateCalls, 1)
+		assert.Equal(t, handlerCalls, 0)
+		assert.Equal(t, gateSpec.Procedure, pingv1connect.PingServicePingProcedure)
+		assert.Equal(t, gateSpec.StreamType, connect.StreamTypeUnary)
+		assert.Equal(t, gateToken, "sesame")
+		assert.NotZero(t, gatePeer.Addr)
+	})
+
+	t.Run("passes_context_to_handler", func(t *testing.T) {
+		t.Parallel()
+		type gateKey struct{}
+		var seen any
+		mux := http.NewServeMux()
+		mux.Handle(pingv1connect.NewPingServiceHandler(
+			&pluggablePingServer{ping: func(ctx context.Context, request *connect.Request[pingv1.PingRequest]) (*connect.Response[pingv1.PingResponse], error) {
+				seen = ctx.Value(gateKey{})
+				return connect.NewResponse(&pingv1.PingResponse{Number: request.Msg.GetNumber()}), nil
+			}},
+			connect.WithRequestGate(func(ctx context.Context, _ connect.Spec, _ connect.Peer, _ http.Header) (context.Context, error) {
+				return context.WithValue(ctx, gateKey{}, "gated"), nil
+			}),
+		))
+		server := memhttptest.NewServer(t, mux)
+		client := pingv1connect.NewPingServiceClient(server.Client(), server.URL())
+		response, err := client.Ping(t.Context(), connect.NewRequest(&pingv1.PingRequest{Number: 42}))
+		assert.Nil(t, err)
+		assert.Equal(t, response.Msg.GetNumber(), int64(42))
+		assert.Equal(t, seen, "gated")
+	})
+
+	t.Run("runs_in_order", func(t *testing.T) {
+		t.Parallel()
+		var order []string
+		gate := func(name string, err error) connect.HandlerOption {
+			return connect.WithRequestGate(func(ctx context.Context, _ connect.Spec, _ connect.Peer, _ http.Header) (context.Context, error) {
+				order = append(order, name)
+				return ctx, err
+			})
+		}
+		mux := http.NewServeMux()
+		mux.Handle(pingv1connect.NewPingServiceHandler(
+			successPingServer{},
+			gate("first", nil),
+			gate("second", errRejected),
+			gate("third", nil),
+		))
+		server := memhttptest.NewServer(t, mux)
+		client := pingv1connect.NewPingServiceClient(server.Client(), server.URL())
+		_, err := client.Ping(t.Context(), connect.NewRequest(&pingv1.PingRequest{}))
+		assert.Equal(t, connect.CodeOf(err), connect.CodeUnauthenticated)
+		// The third gate does not run: the second short-circuits the RPC.
+		assert.Equal(t, order, []string{"first", "second"})
+	})
+
+	t.Run("gates_streaming_procedures", func(t *testing.T) {
+		t.Parallel()
+		var gateStreamType connect.StreamType
+		mux := http.NewServeMux()
+		mux.Handle(pingv1connect.NewPingServiceHandler(
+			successPingServer{},
+			connect.WithRequestGate(func(_ context.Context, spec connect.Spec, _ connect.Peer, _ http.Header) (context.Context, error) {
+				gateStreamType = spec.StreamType
+				return nil, errRejected
+			}),
+		))
+		server := memhttptest.NewServer(t, mux)
+		client := pingv1connect.NewPingServiceClient(server.Client(), server.URL())
+		stream := client.CumSum(t.Context())
+		assert.Nil(t, stream.Send(&pingv1.CumSumRequest{Number: 1}))
+		_, err := stream.Receive()
+		assert.Equal(t, connect.CodeOf(err), connect.CodeUnauthenticated)
+		assert.Equal(t, gateStreamType, connect.StreamTypeBidi)
+		assert.Nil(t, stream.CloseRequest())
+		assert.Nil(t, stream.CloseResponse())
+	})
+}
+
 type successPingServer struct {
 	pingv1connect.UnimplementedPingServiceHandler
 }

option.go

@@ -154,6 +154,26 @@ func WithRecover(handle func(context.Context, Spec, http.Header, any) error) Han
 	return WithInterceptors(&recoverHandlerInterceptor{handle: handle})
 }
 
+// WithRequestGate registers a function that runs once the request headers are
+// available and before any message is received. Returning a non-nil error
+// rejects the RPC: the error is sent to the client, and neither the interceptor
+// chain nor the handler runs, so the request message is never decompressed or
+// unmarshaled. Returning a non-nil context replaces the one passed to
+// interceptors and the handler.
+//
+// Gates suit authentication, which is costly to do in an [Interceptor]:
+// interceptors run after the request message has been received, so rejecting an
+// unauthenticated unary call there pays for the decode first. Gates run for
+// every stream type. Authenticating in net/http middleware, with
+// connectrpc.com/authn or otherwise, runs earlier still.
+//
+// Applying this option more than once registers multiple gates, which run in
+// the order they were applied. Gates must be safe to call concurrently, and
+// panics in a gate are not recovered by [WithRecover].
+func WithRequestGate(gate func(ctx context.Context, spec Spec, peer Peer, header http.Header) (context.Context, error)) HandlerOption {
+	return &requestGateOption{gate: gate}
+}
+
 // WithRequireConnectProtocolHeader configures the Handler to require requests
 // using the Connect RPC protocol to include the Connect-Protocol-Version
 // header. This ensures that HTTP proxies and net/http middleware can easily
@@ -498,6 +518,14 @@ func (o *handlerOptionsOption) applyToHandler(config *handlerConfig) {
 	}
 }
 
+type requestGateOption struct {
+	gate func(ctx context.Context, spec Spec, peer Peer, header http.Header) (context.Context, error)
+}
+
+func (o *requestGateOption) applyToHandler(config *handlerConfig) {
+	config.RequestGates = append(config.RequestGates, o.gate)
+}
+
 type requireConnectProtocolHeaderOption struct{}
 
 func (o *requireConnectProtocolHeaderOption) applyToHandler(config *handlerConfig) {