Tweak JWT, expect account vs system in field

Signed-off-by: Neil Twigg <[email protected]>

Author: neilalexander

go.mod

@@ -2,7 +2,7 @@ module github.com/nats-io/nats-server/v2
 
 go 1.21.0
 
-replace github.com/nats-io/jwt/v2 => github.com/nats-io/jwt/v2 v2.5.9-0.20240730132529-79732145f9be
+replace github.com/nats-io/jwt/v2 => github.com/nats-io/jwt/v2 v2.5.9-0.20240801130136-270cc45c44ee
 
 require (
 	github.com/google/go-tpm v0.9.0

go.sum

@@ -6,8 +6,8 @@ github.com/klauspost/compress v1.17.9 h1:6KIumPrER1LHsvBVuDa0r5xaG0Es51mhhB9BQB2
 github.com/klauspost/compress v1.17.9/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=
 github.com/minio/highwayhash v1.0.3 h1:kbnuUMoHYyVl7szWjSxJnxw11k2U709jqFPPmIUyD6Q=
 github.com/minio/highwayhash v1.0.3/go.mod h1:GGYsuwP/fPD6Y9hMiXuapVvlIUEhFhMTh0rxU3ik1LQ=
-github.com/nats-io/jwt/v2 v2.5.9-0.20240730132529-79732145f9be h1:0EpVNpiLRtjqTgsz8TFtVeNzOwpjC/Xu7yXDiqJSxyI=
-github.com/nats-io/jwt/v2 v2.5.9-0.20240730132529-79732145f9be/go.mod h1:ZdWS1nZa6WMZfFwwgpEaqBV8EPGVgOTDHN/wTbz0Y5A=
+github.com/nats-io/jwt/v2 v2.5.9-0.20240801130136-270cc45c44ee h1:kiTo11kMbk4UMNjdgXAFFSq8+p1RLI3XPcAukFaiw+g=
+github.com/nats-io/jwt/v2 v2.5.9-0.20240801130136-270cc45c44ee/go.mod h1:ZdWS1nZa6WMZfFwwgpEaqBV8EPGVgOTDHN/wTbz0Y5A=
 github.com/nats-io/nats.go v1.36.0 h1:suEUPuWzTSse/XhESwqLxXGuj8vGRuPRoG7MoRN/qyU=
 github.com/nats-io/nats.go v1.36.0/go.mod h1:Ubdu4Nh9exXdSz0RVWRFBbRfrbSxOYd26oF0wkWclB8=
 github.com/nats-io/nkeys v0.4.7 h1:RwNJbbIdYCoClSDNY7QVKZlyb/wfT6ugvFCiKy6vDvI=

server/accounts.go

@@ -3683,7 +3683,16 @@ func (s *Server) updateAccountClaimsWithRefresh(a *Account, ac *jwt.AccountClaim
 	if a.js != nil {
 		// Check whether the account NRG status changed. If it has then we need to notify the
 		// Raft groups running on the system so that they can move their subs if needed.
-		if wasAccountNRG := a.js.accountNRG.Swap(ac.AccountNRG); wasAccountNRG != ac.AccountNRG {
+		wantAccountNRG := a.js.accountNRG.Load()
+		switch strings.ToLower(ac.NRGAccount) {
+		case "account":
+			wantAccountNRG = true
+		case "system":
+			wantAccountNRG = false
+		default:
+			s.Errorf("Account claim for %q has invalid value %q for account NRG status", a.Name, ac.NRGAccount)
+		}
+		if wasAccountNRG := a.js.accountNRG.Swap(wantAccountNRG); wasAccountNRG != wantAccountNRG {
 			s.updateNRGAccountStatus()
 		}
 	}

server/jetstream_jwt_test.go

@@ -1571,8 +1571,8 @@ func TestJetStreamJWTClusterAccountNRG(t *testing.T) {
 
 	// We'll try flipping the state a few times and then do some sanity
 	// checks to check that it took effect.
-	for _, state := range []bool{true, false, true} {
-		accClaim.AccountNRG = state
+	for _, state := range []string{"account", "system", "account"} {
+		accClaim.NRGAccount = state
 		accJwt = encodeClaim(t, accClaim, aExpPub)
 
 		for _, s := range c.servers {
@@ -1586,7 +1586,7 @@ func TestJetStreamJWTClusterAccountNRG(t *testing.T) {
 			// Check that everything looks like it should.
 			require_True(t, acc != nil)
 			require_True(t, acc.js != nil)
-			require_Equal(t, acc.js.accountNRG.Load(), state)
+			require_Equal(t, acc.js.accountNRG.Load(), state == "account")
 
 			// Now get a list of all of the Raft nodes that should
 			// have been updated by now.
@@ -1608,7 +1608,7 @@ func TestJetStreamJWTClusterAccountNRG(t *testing.T) {
 				inAcc := rg.inAcc
 				rg.Unlock()
 
-				require_Equal(t, inAcc, state)
+				require_Equal(t, inAcc, state == "account")
 			}
 		}
 	}