drpc, server: implement filtering rules for server initialization state

Previously, in DRPC, the `batch` RPC was able to access
`node.Descriptor` which was yet to be initialized. This led to
a data race, described in #153948.
This wasn't the case for gRPC, because of an interceptor that
could filter RPCs during different server states. Particularly,
it allows only bootstrap, heartbeat, health and gossip methods
during initialization to prevent calls to potentially uninitialized
services.
This patch adds this interceptor support into drpc as well.
In-depth discussion on this: #153948 (comment)

Epic: None
Fixes: #153948
Release note: None

Author: Nukitt

pkg/server/drpc_server.go

@@ -31,7 +31,17 @@ type drpcServer struct {
 // newDRPCServer creates and configures a new drpcServer instance. It enables
 // DRPC if the experimental setting is on, otherwise returns a dummy server.
 func newDRPCServer(ctx context.Context, rpcCtx *rpc.Context) (*drpcServer, error) {
-	d, err := rpc.NewDRPCServer(ctx, rpcCtx)
+	drpcServer := &drpcServer{}
+	drpcServer.setMode(modeInitializing)
+
+	d, err := rpc.NewDRPCServer(
+		ctx,
+		rpcCtx,
+		rpc.WithInterceptor(
+			func(path string) error {
+				return drpcServer.intercept(path)
+			}),
+	)
 	if err != nil {
 		return nil, err
 	}
@@ -41,12 +51,8 @@ func newDRPCServer(ctx context.Context, rpcCtx *rpc.Context) (*drpcServer, error
 		return nil, err
 	}
 
-	drpcServer := &drpcServer{
-		DRPCServer: d,
-		tlsCfg:     tlsCfg,
-	}
-
-	drpcServer.setMode(modeInitializing)
+	drpcServer.DRPCServer = d
+	drpcServer.tlsCfg = tlsCfg
 
 	if err := rpc.DRPCRegisterHeartbeat(drpcServer, rpcCtx.NewHeartbeatService()); err != nil {
 		return nil, err

pkg/server/grpc_server.go

@@ -68,24 +68,6 @@ func (s *grpcServer) health(ctx context.Context) error {
 	}
 }
 
-var rpcsAllowedWhileBootstrapping = map[string]struct{}{
-	"/cockroach.rpc.Heartbeat/Ping":             {},
-	"/cockroach.gossip.Gossip/Gossip":           {},
-	"/cockroach.server.serverpb.Init/Bootstrap": {},
-	"/cockroach.server.serverpb.Admin/Health":   {},
-}
-
-// intercept implements filtering rules for each server state.
-func (s *grpcServer) intercept(fullName string) error {
-	if s.operational() {
-		return nil
-	}
-	if _, allowed := rpcsAllowedWhileBootstrapping[fullName]; !allowed {
-		return NewWaitingForInitError(fullName)
-	}
-	return nil
-}
-
 // NewWaitingForInitError creates an error indicating that the server cannot run
 // the specified method until the node has been initialized.
 func NewWaitingForInitError(methodName string) error {

pkg/server/serve_mode.go

@@ -36,6 +36,24 @@ const (
 	modeDraining
 )
 
+var rpcsAllowedWhileBootstrapping = map[string]struct{}{
+	"/cockroach.rpc.Heartbeat/Ping":             {},
+	"/cockroach.gossip.Gossip/Gossip":           {},
+	"/cockroach.server.serverpb.Init/Bootstrap": {},
+	"/cockroach.server.serverpb.Admin/Health":   {},
+}
+
+// intercept implements filtering rules for each server state.
+func (s *serveModeHandler) intercept(fullName string) error {
+	if s.operational() {
+		return nil
+	}
+	if _, allowed := rpcsAllowedWhileBootstrapping[fullName]; !allowed {
+		return NewWaitingForInitError(fullName)
+	}
+	return nil
+}
+
 func (s *serveMode) set(mode serveMode) {
 	atomic.StoreInt32((*int32)(s), int32(mode))
 }