Merge #34197 #36952

34197: server,rpc: validate node IDs in RPC heartbeats r=tbg a=knz

Fixes #34158.

Prior to this patch, it was possible for a RPC client to dial a node
ID and get a connection to another node instead. This is because the
mapping of node ID -> address may be stale, and a different node could
take the address of the intended node from "under" the dialer.

(See #34155 for a scenario.)

This happened to be "safe" in many cases where it matters because:

- RPC requests for distSQL are OK with being served on a different
  node than intended (with potential performance drop);
- RPC requests to the KV layer are OK with being served on a different
  node than intended (they would route underneath);
- RPC requests to the storage layer are rejected by the
  remote node because the store ID in the request would not match.

However this safety is largely accidental, and we should not work with
the assumption that any RPC request is safe to be mis-routed. (In
fact, we have not audited all the RPC endpoints and cannot establish
this safety exists throughout.)

This patch works to prevent these mis-routings by adding a check of
the intended node ID during RPC heartbeats (including the initial
heartbeat), when the intended node ID is known. A new API
`GRPCDialNode()` is introduced to establish such connections.

Release note (bug fix): CockroachDB now performs fewer attempts to
communicate with the wrong node, when a node is restarted with another
node's address.

36952: storage: deflake TestNodeLivenessStatusMap r=tbg a=knz

Fixes #35675.

Prior to this patch, this test would fail `stressrace` after a few
dozen iterations.

With this patch, `stressrace` succeeds thousands of iterations.

I have checked that the test logic is preserved: if I change one of
the expected statuses in `testData`, the test still fail properly.

Release note: None

Co-authored-by: Raphael 'kena' Poss <knz@cockroachlabs.com>

Author: craig[bot]

pkg/cli/debug_test.go

@@ -252,7 +252,10 @@ func TestRemoveDeadReplicas(t *testing.T) {
 	tc := testcluster.StartTestCluster(t, 3, clusterArgs)
 	defer tc.Stopper().Stop(ctx)
 
-	grpcConn, err := tc.Server(0).RPCContext().GRPCDial(tc.Server(0).ServingAddr()).Connect(ctx)
+	grpcConn, err := tc.Server(0).RPCContext().GRPCDialNode(
+		tc.Server(0).ServingAddr(),
+		tc.Server(0).NodeID(),
+	).Connect(ctx)
 	if err != nil {
 		t.Fatal(err)
 	}

pkg/cli/start.go

@@ -1132,7 +1132,9 @@ func getClientGRPCConn(ctx context.Context) (*grpc.ClientConn, *hlc.Clock, func(
 		stopper.Stop(ctx)
 		return nil, nil, nil, err
 	}
-	conn, err := rpcContext.GRPCDial(addr).Connect(ctx)
+	// We use GRPCUnvalidatedDial() here because it does not matter
+	// to which node we're talking to.
+	conn, err := rpcContext.GRPCUnvalidatedDial(addr).Connect(ctx)
 	if err != nil {
 		stopper.Stop(ctx)
 		return nil, nil, nil, err

pkg/gossip/client.go

@@ -109,7 +109,7 @@ func (c *client) startLocked(
 			// asynchronous from the caller's perspective, so the only effect of
 			// `WithBlock` here is blocking shutdown - at the time of this writing,
 			// that ends ups up making `kv` tests take twice as long.
-			conn, err := rpcCtx.GRPCDial(c.addr.String()).Connect(ctx)
+			conn, err := rpcCtx.GRPCUnvalidatedDial(c.addr.String()).Connect(ctx)
 			if err != nil {
 				return err
 			}

pkg/gossip/client_test.go

@@ -66,6 +66,7 @@ func startGossipAtAddr(
 	registry *metric.Registry,
 ) *Gossip {
 	rpcContext := newInsecureRPCContext(stopper)
+	rpcContext.NodeID.Set(context.TODO(), nodeID)
 	server := rpc.NewServer(rpcContext)
 	g := NewTest(nodeID, rpcContext, server, stopper, registry)
 	ln, err := netutil.ListenAndServeGRPC(stopper, server, addr)

pkg/gossip/gossip_test.go

@@ -61,8 +61,7 @@ func TestGossipInfoStore(t *testing.T) {
 }
 
 // TestGossipMoveNode verifies that if a node is moved to a new address, it
-// gets properly updated in gossip (including that any other node that was
-// previously at that address gets removed from the cluster).
+// gets properly updated in gossip.
 func TestGossipMoveNode(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	stopper := stop.NewStopper()
@@ -462,7 +461,7 @@ func TestGossipNoForwardSelf(t *testing.T) {
 		c := newClient(log.AmbientContext{Tracer: tracing.NewTracer()}, local.GetNodeAddr(), makeMetrics())
 
 		testutils.SucceedsSoon(t, func() error {
-			conn, err := peer.rpcContext.GRPCDial(c.addr.String()).Connect(ctx)
+			conn, err := peer.rpcContext.GRPCUnvalidatedDial(c.addr.String()).Connect(ctx)
 			if err != nil {
 				return err
 			}

pkg/kv/send_test.go

@@ -66,6 +66,12 @@ func TestSendToOneClient(t *testing.T) {
 		stopper,
 		&cluster.MakeTestingClusterSettings().Version,
 	)
+
+	// This test uses the testing function sendBatch() which does not
+	// support setting the node ID on GRPCDialNode(). Disable Node ID
+	// checks to avoid log.Fatal.
+	rpcContext.TestingAllowNamedRPCToAnonymousServer = true
+
 	s := rpc.NewServer(rpcContext)
 	roachpb.RegisterInternalServer(s, Node(0))
 	ln, err := netutil.ListenAndServeGRPC(rpcContext.Stopper, s, util.TestAddr)
@@ -136,6 +142,10 @@ func TestComplexScenarios(t *testing.T) {
 		stopper,
 		&cluster.MakeTestingClusterSettings().Version,
 	)
+
+	// We're going to serve multiple node IDs with that one
+	// context. Disable node ID checks.
+	nodeContext.TestingAllowNamedRPCToAnonymousServer = true
 	nodeDialer := nodedialer.New(nodeContext, nil)
 
 	// TODO(bdarnell): the retryable flag is no longer used for RPC errors.

pkg/rpc/context.go

@@ -276,10 +276,12 @@ func NewServerWithInterceptor(
 
 	s := grpc.NewServer(opts...)
 	RegisterHeartbeatServer(s, &HeartbeatService{
-		clock:              ctx.LocalClock,
-		remoteClockMonitor: ctx.RemoteClocks,
-		clusterID:          &ctx.ClusterID,
-		version:            ctx.version,
+		clock:                                 ctx.LocalClock,
+		remoteClockMonitor:                    ctx.RemoteClocks,
+		clusterID:                             &ctx.ClusterID,
+		nodeID:                                &ctx.NodeID,
+		version:                               ctx.version,
+		testingAllowNamedRPCToAnonymousServer: ctx.TestingAllowNamedRPCToAnonymousServer,
 	})
 	return s
 }
@@ -298,14 +300,20 @@ type Connection struct {
 	initialHeartbeatDone chan struct{} // closed after first heartbeat
 	stopper              *stop.Stopper
 
+	// remoteNodeID implies checking the remote node ID. 0 when unknown,
+	// non-zero to check with remote node. This is constant throughout
+	// the lifetime of a Connection object.
+	remoteNodeID roachpb.NodeID
+
 	initOnce      sync.Once
 	validatedOnce sync.Once
 }
 
-func newConnection(stopper *stop.Stopper) *Connection {
+func newConnectionToNodeID(stopper *stop.Stopper, remoteNodeID roachpb.NodeID) *Connection {
 	c := &Connection{
 		initialHeartbeatDone: make(chan struct{}),
 		stopper:              stopper,
+		remoteNodeID:         remoteNodeID,
 	}
 	c.heartbeatResult.Store(heartbeatResult{err: ErrNotHeartbeated})
 	return c
@@ -372,11 +380,23 @@ type Context struct {
 	stats StatsHandler
 
 	ClusterID base.ClusterIDContainer
+	NodeID    base.NodeIDContainer
 	version   *cluster.ExposedClusterVersion
 
 	// For unittesting.
 	BreakerFactory  func() *circuit.Breaker
 	testingDialOpts []grpc.DialOption
+
+	// For testing. See the comment on the same field in HeartbeatService.
+	TestingAllowNamedRPCToAnonymousServer bool
+}
+
+// connKey is used as key in the Context.conns map.  Different remote
+// node IDs get different *Connection objects, to ensure that we don't
+// mis-route RPC requests.
+type connKey struct {
+	targetAddr string
+	nodeID     roachpb.NodeID
 }
 
 // NewContext creates an rpc Context with the supplied values.
@@ -422,7 +442,7 @@ func NewContext(
 					conn.dialErr = &roachpb.NodeUnavailableError{}
 				}
 			})
-			ctx.removeConn(k.(string), conn)
+			ctx.removeConn(k.(connKey), conn)
 			return true
 		})
 	})
@@ -439,8 +459,10 @@ func (ctx *Context) GetStatsMap() *syncmap.Map {
 
 // GetLocalInternalClientForAddr returns the context's internal batch client
 // for target, if it exists.
-func (ctx *Context) GetLocalInternalClientForAddr(target string) roachpb.InternalClient {
-	if target == ctx.AdvertiseAddr {
+func (ctx *Context) GetLocalInternalClientForAddr(
+	target string, nodeID roachpb.NodeID,
+) roachpb.InternalClient {
+	if target == ctx.AdvertiseAddr && nodeID == ctx.NodeID.Get() {
 		return ctx.localInternalClient
 	}
 	return nil
@@ -544,15 +566,15 @@ func (ctx *Context) SetLocalInternalServer(internalServer roachpb.InternalServer
 	ctx.localInternalClient = internalClientAdapter{internalServer}
 }
 
-func (ctx *Context) removeConn(key string, conn *Connection) {
+func (ctx *Context) removeConn(key connKey, conn *Connection) {
 	ctx.conns.Delete(key)
 	if log.V(1) {
-		log.Infof(ctx.masterCtx, "closing %s", key)
+		log.Infof(ctx.masterCtx, "closing %+v", key)
 	}
 	if grpcConn := conn.grpcConn; grpcConn != nil {
 		if err := grpcConn.Close(); err != nil && !grpcutil.IsClosedConnection(err) {
 			if log.V(1) {
-				log.Errorf(ctx.masterCtx, "failed to close client connection: %s", err)
+				log.Errorf(ctx.masterCtx, "failed to close client connection: %v", err)
 			}
 		}
 	}
@@ -675,11 +697,43 @@ func (ctx *Context) GRPCDialRaw(target string) (*grpc.ClientConn, <-chan struct{
 	return conn, dialer.redialChan, err
 }
 
-// GRPCDial calls grpc.Dial with options appropriate for the context.
-func (ctx *Context) GRPCDial(target string) *Connection {
-	value, ok := ctx.conns.Load(target)
+// GRPCUnvalidatedDial uses GRPCDialNode and disables validation of the
+// node ID between client and server. This function should only be
+// used with the gossip client and CLI commands which can talk to any
+// node.
+func (ctx *Context) GRPCUnvalidatedDial(target string) *Connection {
+	return ctx.grpcDialNodeInternal(target, 0)
+}
+
+// GRPCDialNode calls grpc.Dial with options appropriate for the context.
+//
+// The remoteNodeID becomes a constraint on the expected node ID of
+// the remote node; this is checked during heartbeats. The caller is
+// responsible for ensuring the remote node ID is known prior to using
+// this function.
+func (ctx *Context) GRPCDialNode(target string, remoteNodeID roachpb.NodeID) *Connection {
+	if remoteNodeID == 0 && !ctx.TestingAllowNamedRPCToAnonymousServer {
+		log.Fatalf(context.TODO(), "invalid node ID 0 in GRPCDialNode()")
+	}
+	return ctx.grpcDialNodeInternal(target, remoteNodeID)
+}
+
+func (ctx *Context) grpcDialNodeInternal(target string, remoteNodeID roachpb.NodeID) *Connection {
+	thisConnKey := connKey{target, remoteNodeID}
+	value, ok := ctx.conns.Load(thisConnKey)
 	if !ok {
-		value, _ = ctx.conns.LoadOrStore(target, newConnection(ctx.Stopper))
+		value, _ = ctx.conns.LoadOrStore(thisConnKey, newConnectionToNodeID(ctx.Stopper, remoteNodeID))
+		if remoteNodeID != 0 {
+			// If the first connection established at a target address is
+			// for a specific node ID, then we want to reuse that connection
+			// also for other dials (eg for gossip) which don't require a
+			// specific node ID. (We do this as an optimization to reduce
+			// the number of TCP connections alive between nodes. This is
+			// not strictly required for correctness.) This LoadOrStore will
+			// ensure we're registering the connection we just created for
+			// future use by these other dials.
+			_, _ = ctx.conns.LoadOrStore(connKey{target, 0}, value)
+		}
 	}
 
 	conn := value.(*Connection)
@@ -694,11 +748,11 @@ func (ctx *Context) GRPCDial(target string) *Connection {
 						if err != nil && !grpcutil.IsClosedConnection(err) {
 							log.Errorf(masterCtx, "removing connection to %s due to error: %s", target, err)
 						}
-						ctx.removeConn(target, conn)
+						ctx.removeConn(thisConnKey, conn)
 					})
 				}); err != nil {
 				conn.dialErr = err
-				ctx.removeConn(target, conn)
+				ctx.removeConn(thisConnKey, conn)
 			}
 		}
 	})
@@ -720,7 +774,7 @@ func (ctx *Context) NewBreaker(name string) *circuit.Breaker {
 // the first heartbeat.
 var ErrNotHeartbeated = errors.New("not yet heartbeated")
 
-// ConnHealth returns nil if we have an open connection to the given
+// TestingConnHealth returns nil if we have an open connection to the given
 // target that succeeded on its most recent heartbeat. Otherwise, it
 // kicks off a connection attempt (unless one is already in progress
 // or we are in a backoff state) and returns an error (typically
@@ -729,27 +783,26 @@ var ErrNotHeartbeated = errors.New("not yet heartbeated")
 // error will be returned. This method should therefore be used to
 // prioritize among a list of candidate nodes, but not to filter out
 // "unhealthy" nodes.
-func (ctx *Context) ConnHealth(target string) error {
-	if ctx.GetLocalInternalClientForAddr(target) != nil {
+//
+// This is used in tests only; in clusters use (*Dialer).ConnHealth()
+// instead which automates the address resolution.
+//
+// TODO(knz): remove this altogether. Use the dialer in all cases.
+func (ctx *Context) TestingConnHealth(target string, nodeID roachpb.NodeID) error {
+	if ctx.GetLocalInternalClientForAddr(target, nodeID) != nil {
 		// The local server is always considered healthy.
 		return nil
 	}
-	conn := ctx.GRPCDial(target)
+	conn := ctx.GRPCDialNode(target, nodeID)
 	return conn.Health()
 }
 
 func (ctx *Context) runHeartbeat(
 	conn *Connection, target string, redialChan <-chan struct{},
 ) error {
 	maxOffset := ctx.LocalClock.MaxOffset()
-	clusterID := ctx.ClusterID.Get()
+	maxOffsetNanos := maxOffset.Nanoseconds()
 
-	request := PingRequest{
-		Addr:           ctx.Addr,
-		MaxOffsetNanos: maxOffset.Nanoseconds(),
-		ClusterID:      &clusterID,
-		ServerVersion:  ctx.version.ServerVersion,
-	}
 	heartbeatClient := NewHeartbeatClient(conn.grpcConn)
 
 	var heartbeatTimer timeutil.Timer
@@ -768,14 +821,24 @@ func (ctx *Context) runHeartbeat(
 			heartbeatTimer.Read = true
 		}
 
+		// We re-mint the PingRequest to pick up any asynchronous update to clusterID.
+		clusterID := ctx.ClusterID.Get()
+		request := &PingRequest{
+			Addr:           ctx.Addr,
+			MaxOffsetNanos: maxOffsetNanos,
+			ClusterID:      &clusterID,
+			NodeID:         conn.remoteNodeID,
+			ServerVersion:  ctx.version.ServerVersion,
+		}
+
 		var response *PingResponse
 		sendTime := ctx.LocalClock.PhysicalTime()
 		err := contextutil.RunWithTimeout(ctx.masterCtx, "rpc heartbeat", ctx.heartbeatTimeout,
 			func(goCtx context.Context) error {
 				// NB: We want the request to fail-fast (the default), otherwise we won't
 				// be notified of transport failures.
 				var err error
-				response, err = heartbeatClient.Ping(goCtx, &request)
+				response, err = heartbeatClient.Ping(goCtx, request)
 				return err
 			})