[FIXED] 'cluster_traffic: owner' is forgotten after restart (#7191)

`cluster_traffic: owner` can be used to have replication traffic go over
the account "owning" the stream, versus this traffic going over the
system account (`cluster_traffic: system`/default).

When pushing an updated JWT to all servers, all these servers would
correctly update their `cluster_traffic` setting.

However, if a server was restarted it would "forget" `cluster_traffic:
owner` was set, and revert back to `cluster_traffic: system`. This would
put this single server to be unable to communicate with the remainder of
the cluster. Restarting the rest of the cluster would have them also
revert back to `cluster_traffic: system`, which would allow them to
communicate again. But, without respecting the `cluster_traffic: owner`
setting on the account.

This PR fixes that by ensuring `cluster_traffic` can be updated at the
same time as that JetStream is enabled for that particular account upon
startup.

Due to this issue, if a server with this fix is deployed, it will not be
able to communicate with the other servers that had reverted back to
`cluster_traffic: system`. A clean upgrade path for this would be:
- Temporarily update the account to have `cluster_traffic: system`. If
all servers were restarted they were already on this setting. Any
servers that were not yet restarted will now agree on this setting.
- Upgrade all servers to the new server version with this fix.
- Update the account to have `cluster_traffic: owner` again. It should
now be remembered even after a server restart.

Signed-off-by: Maurice van Veen <[email protected]>

Author: neilalexander

server/accounts.go

@@ -3784,6 +3784,11 @@ func (s *Server) updateAccountClaimsWithRefresh(a *Account, ac *jwt.AccountClaim
 			// Absent reload of js server cfg, this is going to be broken until js is disabled
 			a.incomplete = true
 			a.mu.Unlock()
+		} else {
+			a.mu.Lock()
+			// Refresh reference, we've just enabled JetStream, so it would have been nil before.
+			ajs = a.js
+			a.mu.Unlock()
 		}
 	} else if a.jsLimits != nil {
 		// We do not have JS enabled for this server, but the account has it enabled so setup

server/jetstream_jwt_test.go

@@ -1821,6 +1821,91 @@ func TestJetStreamJWTClusterAccountNRG(t *testing.T) {
 	}
 }
 
+func TestJetStreamJWTClusterAccountNRGPersistsAfterRestart(t *testing.T) {
+	_, syspub := createKey(t)
+	sysJwt := encodeClaim(t, jwt.NewAccountClaims(syspub), syspub)
+
+	aExpKp, aExpPub := createKey(t)
+	accClaim := jwt.NewAccountClaims(aExpPub)
+	accClaim.Name = "acc"
+	accClaim.ClusterTraffic = jwt.ClusterTrafficOwner
+	accClaim.Limits.JetStreamTieredLimits["R1"] = jwt.JetStreamLimits{DiskStorage: 1100, Consumer: 10, Streams: 1}
+	accClaim.Limits.JetStreamTieredLimits["R3"] = jwt.JetStreamLimits{DiskStorage: 1100, Consumer: 1, Streams: 1}
+	accJwt := encodeClaim(t, accClaim, aExpPub)
+	accCreds := newUser(t, aExpKp)
+
+	tmlp := `
+		listen: 127.0.0.1:-1
+		server_name: %s
+		jetstream: {max_mem_store: 256MB, max_file_store: 2GB, store_dir: '%s'}
+		leaf {
+			listen: 127.0.0.1:-1
+		}
+		cluster {
+			name: %s
+			listen: 127.0.0.1:%d
+			routes = [%s]
+		}
+	` + fmt.Sprintf(`
+		operator: %s
+		system_account: %s
+		resolver = MEMORY
+		resolver_preload = {
+			%s : %s
+			%s : %s
+		}
+	`, ojwt, syspub, syspub, sysJwt, aExpPub, accJwt)
+
+	c := createJetStreamClusterWithTemplate(t, tmlp, "cluster", 3)
+	defer c.shutdown()
+
+	nc, _ := jsClientConnect(t, c.randomServer(), nats.UserCredentials(accCreds))
+	_, err := jsStreamCreate(t, nc, &StreamConfig{
+		Name:     "TEST",
+		Replicas: 3,
+		Storage:  FileStorage,
+	})
+	require_NoError(t, err)
+
+	// The account had cluster traffic set to "owner" already. Restarting servers should remember this setting.
+	for _, s := range c.servers {
+		acc, err := s.lookupAccount(aExpPub)
+		require_NoError(t, err)
+
+		// Check that everything looks like it should.
+		require_True(t, acc != nil)
+		require_True(t, acc.js != nil)
+		require_Equal(t, acc.nrgAccount, aExpPub)
+
+		// Now get a list of all the Raft nodes that should have the correct cluster traffic set.
+		s.rnMu.Lock()
+		raftNodes := make([]*raft, 0, len(s.raftNodes))
+		for _, n := range s.raftNodes {
+			rg := n.(*raft)
+			if rg.accName != acc.Name {
+				continue
+			}
+			raftNodes = append(raftNodes, rg)
+		}
+		s.rnMu.Unlock()
+
+		// Get the Raftz state also.
+		rz := s.Raftz(&RaftzOptions{AccountFilter: aExpPub})
+		require_NotNil(t, rz)
+		rza := (*rz)[aExpPub]
+		require_NotNil(t, rza)
+
+		for _, rg := range raftNodes {
+			rg.Lock()
+			rgAcc := rg.acc
+			rg.Unlock()
+			require_Equal(t, rgAcc.Name, aExpPub)
+			require_Equal(t, rza[rg.group].SystemAcc, false)
+			require_Equal(t, rza[rg.group].TrafficAcc, aExpPub)
+		}
+	}
+}
+
 func TestJetStreamJWTUpdateWithPreExistingStream(t *testing.T) {
 	updateJwt := func(url string, creds string, pubKey string, jwt string) {
 		t.Helper()