Merge #154842

154842: rpc: respect interceptors ServerOption in NewDRPCServer r=tbg a=stevendanna

We were not adding the passed interceptor to the list of unary or stream interceptors. Now we do. This in turn addresses a race condition that we intended to solve in e005120.

Part of this PR was authored by Claude Code. After a hint about DRPC vs gRPC, it was able to diagnose the problem given race detector output and generate the fix included here.

Its initial attempts at creating a test either didn't actually assert anything relevant or was very large. After some prompting, it created a test similar to what is included here, but with many hallucinated functions. I then just fixed it by hand. 

Epic: none
Release note: None

Co-authored-by: Steven Danna <[email protected]>

Author: craig[bot]

pkg/rpc/drpc.go

@@ -236,6 +236,26 @@ func NewDRPCServer(_ context.Context, rpcCtx *Context, opts ...ServerOption) (DR
 		streamInterceptors = append(streamInterceptors, a.AuthDRPCStream())
 	}
 
+	if o.interceptor != nil {
+		unaryInterceptors = append(unaryInterceptors, func(
+			ctx context.Context, req interface{}, fullMethod string, handler drpcmux.UnaryHandler,
+		) (interface{}, error) {
+			if err := o.interceptor(fullMethod); err != nil {
+				return nil, err
+			}
+			return handler(ctx, req)
+		})
+
+		streamInterceptors = append(streamInterceptors, func(
+			stream drpc.Stream, fullMethod string, handler drpcmux.StreamHandler,
+		) (interface{}, error) {
+			if err := o.interceptor(fullMethod); err != nil {
+				return nil, err
+			}
+			return handler(stream)
+		})
+	}
+
 	if tracer := rpcCtx.Stopper.Tracer(); tracer != nil {
 		unaryInterceptors = append(unaryInterceptors, drpcinterceptor.ServerInterceptor(tracer))
 		streamInterceptors = append(streamInterceptors, drpcinterceptor.StreamServerInterceptor(tracer))

pkg/rpc/drpc_test.go

@@ -7,11 +7,19 @@ package rpc
 
 import (
 	"context"
+	"crypto/tls"
+	"sync/atomic"
 	"testing"
+	"time"
 
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/rpc/rpcbase"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/netutil"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
+	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
+	"github.com/cockroachdb/cockroach/pkg/util/uuid"
 	"github.com/cockroachdb/errors"
 	"github.com/stretchr/testify/require"
 	"storj.io/drpc"
@@ -129,3 +137,47 @@ func TestGatewayRequestDRPCRecoveryInterceptor(t *testing.T) {
 		require.ErrorIs(t, err, expectedErr)
 	})
 }
+
+// TestDRPCServerWithInterceptor verifies that configured interceptors are
+// invoked.
+func TestDRPCServerWithInterceptor(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	stopper := stop.NewStopper()
+	defer stopper.Stop(ctx)
+
+	loopbackL := netutil.NewLoopbackListener(ctx, stopper)
+
+	serverCtx := newTestContext(uuid.MakeV4(), timeutil.NewManualTime(timeutil.Unix(0, 1)), 0, stopper)
+	serverCtx.NodeID.Set(ctx, 1)
+	serverCtx.SetLoopbackDRPCDialer(loopbackL.Connect)
+	serverCtx.AdvertiseAddr = "127.0.0.1:1"
+	serverCtx.RPCHeartbeatInterval = 1 * time.Hour
+
+	var shouldBlock atomic.Bool
+	blockedErr := errors.New("RPC blocked by interceptor")
+	drpcServer, err := NewDRPCServer(ctx, serverCtx, WithInterceptor(func(rpcName string) error {
+		if shouldBlock.Load() && rpcName == "/cockroach.rpc.Heartbeat/Ping" {
+			return blockedErr
+		}
+		return nil
+	}))
+	require.NoError(t, err)
+	require.NoError(t, DRPCRegisterHeartbeat(drpcServer, serverCtx.NewHeartbeatService()))
+
+	tlsConfig, err := serverCtx.GetServerTLSConfig()
+	require.NoError(t, err)
+	tlsListener := tls.NewListener(loopbackL, tlsConfig)
+	require.NoError(t, stopper.RunAsyncTask(ctx, "drpc-server", func(ctx context.Context) {
+		netutil.FatalIfUnexpected(drpcServer.Serve(ctx, tlsListener))
+	}))
+
+	conn, err := serverCtx.DRPCDialNode("127.0.0.1:1", 1, roachpb.Locality{}, rpcbase.DefaultClass).Connect(ctx)
+	require.NoError(t, err)
+	client := NewDRPCHeartbeatClient(conn)
+	shouldBlock.Store(true)
+	_, err = client.Ping(ctx, &PingRequest{})
+	require.Contains(t, err.Error(), "RPC blocked by interceptor")
+}