Repository: nats-io/nats-server
Pull request #6249 (2024-12-13 08:01:58)

Adds a number of helper structs for z responses
These are used by the cli and other tools to decode the responses from the z endpoints

Signed-off-by: R.I.Pienaar <rip@devco.net>


ripienaar: These types doesn't exist anywhere its constructed from a `any` wrapped into all in the same structure

So to use them everyone either have to double json decode (This is how MH designed it and insisted its used) or write these structs themselves, so this just helps out.
derekcollison: ok let's comment why they are here and not used in any of the server code base just for clairity.

Patch:
diff --git a/README.md b/README.md
index 3c850d93..c5d524f3 100644
--- a/README.md
+++ b/README.md
@@ -36,7 +36,7 @@ If you are interested in contributing to NATS, read about our...
 [Fossa-Url]: https://app.fossa.io/projects/git%2Bgithub.com%2Fnats-io%2Fnats-server?ref=badge_shield
 [Fossa-Image]: https://app.fossa.io/api/projects/git%2Bgithub.com%2Fnats-io%2Fnats-server.svg?type=shield
 [Build-Status-Url]: https://travis-ci.com/github/nats-io/nats-server
-[Build-Status-Image]: https://travis-ci.org/nats-io/nats-server.svg?branch=main
+[Build-Status-Image]: https://travis-ci.com/nats-io/nats-server.svg?branch=main
 [Release-Url]: https://github.com/nats-io/nats-server/releases/latest
 [Release-image]: https://img.shields.io/github/v/release/nats-io/nats-server
 [Coverage-Url]: https://coveralls.io/r/nats-io/nats-server?branch=main
diff --git a/go.mod b/go.mod
index 60453ff5..f55ab729 100644
--- a/go.mod
+++ b/go.mod
@@ -8,9 +8,9 @@ require (
 	github.com/google/go-tpm v0.9.0
 	github.com/klauspost/compress v1.17.11
 	github.com/minio/highwayhash v1.0.3
-	github.com/nats-io/jwt/v2 v2.7.3
+	github.com/nats-io/jwt/v2 v2.7.2
 	github.com/nats-io/nats.go v1.36.0
-	github.com/nats-io/nkeys v0.4.9
+	github.com/nats-io/nkeys v0.4.8
 	github.com/nats-io/nuid v1.0.1
 	go.uber.org/automaxprocs v1.6.0
 	golang.org/x/crypto v0.31.0
diff --git a/go.sum b/go.sum
index 28407433..cbffa134 100644
--- a/go.sum
+++ b/go.sum
@@ -6,12 +6,12 @@ github.com/klauspost/compress v1.17.11 h1:In6xLpyWOi1+C7tXUUWv2ot1QvBjxevKAaI6IX
 github.com/klauspost/compress v1.17.11/go.mod h1:pMDklpSncoRMuLFrf1W9Ss9KT+0rH90U12bZKk7uwG0=
 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.7.3 h1:6bNPK+FXgBeAqdj4cYQ0F8ViHRbi7woQLq4W29nUAzE=
-github.com/nats-io/jwt/v2 v2.7.3/go.mod h1:GvkcbHhKquj3pkioy5put1wvPxs78UlZ7D/pY+BgZk4=
+github.com/nats-io/jwt/v2 v2.7.2 h1:SCRjfDLJ2q8naXp8YlGJJS5/yj3wGSODFYVi4nnwVMw=
+github.com/nats-io/jwt/v2 v2.7.2/go.mod h1:kB6QUmqHG6Wdrzj0KP2L+OX4xiTPBeV+NHVstFaATXU=
 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.9 h1:qe9Faq2Gxwi6RZnZMXfmGMZkg3afLLOtrU+gDZJ35b0=
-github.com/nats-io/nkeys v0.4.9/go.mod h1:jcMqs+FLG+W5YO36OX6wFIFcmpdAns+w1Wm6D3I/evE=
+github.com/nats-io/nkeys v0.4.8 h1:+wee30071y3vCZAYRsnrmIPaOe47A/SkK/UBDPdIV70=
+github.com/nats-io/nkeys v0.4.8/go.mod h1:kqXRgRDPlGy7nGaEDMuYzmiJCIAAWDK0IMBtDmGD0nc=
 github.com/nats-io/nuid v1.0.1 h1:5iA8DT8V7q8WK2EScv2padNa/rTESc1KdnPw4TC2paw=
 github.com/nats-io/nuid v1.0.1/go.mod h1:19wcPz3Ph3q0Jbyiqsd0kePYG7A95tJPxeL+1OSON2c=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
diff --git a/server/client.go b/server/client.go
index b5fc75a8..3d05ea01 100644
--- a/server/client.go
+++ b/server/client.go
@@ -4288,8 +4288,9 @@ func getHeader(key string, hdr []byte) []byte {
 
 // For bytes.HasPrefix below.
 var (
-	jsRequestNextPreB = []byte(jsRequestNextPre)
-	jsDirectGetPreB   = []byte(jsDirectGetPre)
+	jsRequestNextPreB  = []byte(jsRequestNextPre)
+	jsDirectGetPreB    = []byte(jsDirectGetPre)
+	jsConsumerInfoPreB = []byte(JSApiConsumerInfoPre)
 )
 
 // processServiceImport is an internal callback when a subscription matches an imported service
@@ -4309,12 +4310,16 @@ func (c *client) processServiceImport(si *serviceImport, acc *Account, msg []byt
 		}
 	}
 
+	var checkJS, checkConsumerInfo bool
+
 	acc.mu.RLock()
-	var checkJS bool
 	shouldReturn := si.invalid || acc.sl == nil
 	if !shouldReturn && !isResponse && si.to == jsAllAPI {
 		if bytes.HasPrefix(c.pa.subject, jsDirectGetPreB) || bytes.HasPrefix(c.pa.subject, jsRequestNextPreB) {
 			checkJS = true
+		} else if len(c.pa.psi) == 0 && bytes.HasPrefix(c.pa.subject, jsConsumerInfoPreB) {
+			// Only check if we are clustered and expecting a reply.
+			checkConsumerInfo = len(c.pa.reply) > 0 && c.srv.JetStreamIsClustered()
 		}
 	}
 	siAcc := si.acc
@@ -4329,6 +4334,15 @@ func (c *client) processServiceImport(si *serviceImport, acc *Account, msg []byt
 		return
 	}
 
+	// Here we will do a fast check for consumer info only to check if it does not exists. This will spread the
+	// load to all servers with connected clients since service imports are processed at point of entry.
+	// Only call for clustered setups.
+	if checkConsumerInfo && si.se != nil && si.se.acc == c.srv.SystemAccount() {
+		if c.srv.jsConsumerProcessMissing(c, acc) {
+			return
+		}
+	}
+
 	mt, traceOnly := c.isMsgTraceEnabled()
 
 	var nrr []byte
diff --git a/server/events.go b/server/events.go
index 3021edb7..362ec2a5 100644
--- a/server/events.go
+++ b/server/events.go
@@ -2004,7 +2004,9 @@ type ServerAPIResponse struct {
 	compress compressionType
 }
 
-// Specialized response types for unmarshalling.
+// Specialized response types for unmarshalling. These structures are not
+// used in the server code and only there for users of the Z endpoints to
+// unmarshal the data without having to create these structs in their code
 
 // ServerAPIConnzResponse is the response type connz
 type ServerAPIConnzResponse struct {
@@ -2013,6 +2015,69 @@ type ServerAPIConnzResponse struct {
 	Error  *ApiError   `json:"error,omitempty"`
 }
 
+// ServerAPIRoutezResponse is the response type for routez
+type ServerAPIRoutezResponse struct {
+	Server *ServerInfo `json:"server"`
+	Data   *Routez     `json:"data,omitempty"`
+	Error  *ApiError   `json:"error,omitempty"`
+}
+
+// ServerAPIGatewayzResponse is the response type for gatewayz
+type ServerAPIGatewayzResponse struct {
+	Server *ServerInfo `json:"server"`
+	Data   *Gatewayz   `json:"data,omitempty"`
+	Error  *ApiError   `json:"error,omitempty"`
+}
+
+// ServerAPIJszResponse is the response type for jsz
+type ServerAPIJszResponse struct {
+	Server *ServerInfo `json:"server"`
+	Data   *JSInfo     `json:"data,omitempty"`
+	Error  *ApiError   `json:"error,omitempty"`
+}
+
+// ServerAPIHealthzResponse is the response type for healthz
+type ServerAPIHealthzResponse struct {
+	Server *ServerInfo   `json:"server"`
+	Data   *HealthStatus `json:"data,omitempty"`
+	Error  *ApiError     `json:"error,omitempty"`
+}
+
+// ServerAPIVarzResponse is the response type for varz
+type ServerAPIVarzResponse struct {
+	Server *ServerInfo `json:"server"`
+	Data   *Varz       `json:"data,omitempty"`
+	Error  *ApiError   `json:"error,omitempty"`
+}
+
+// ServerAPISubszResponse is the response type for subsz
+type ServerAPISubszResponse struct {
+	Server *ServerInfo `json:"server"`
+	Data   *Subsz      `json:"data,omitempty"`
+	Error  *ApiError   `json:"error,omitempty"`
+}
+
+// ServerAPILeafzResponse is the response type for leafz
+type ServerAPILeafzResponse struct {
+	Server *ServerInfo `json:"server"`
+	Data   *Leafz      `json:"data,omitempty"`
+	Error  *ApiError   `json:"error,omitempty"`
+}
+
+// ServerAPIAccountzResponse is the response type for accountz
+type ServerAPIAccountzResponse struct {
+	Server *ServerInfo `json:"server"`
+	Data   *Accountz   `json:"data,omitempty"`
+	Error  *ApiError   `json:"error,omitempty"`
+}
+
+// ServerAPIExpvarzResponse is the response type for expvarz
+type ServerAPIExpvarzResponse struct {
+	Server *ServerInfo    `json:"server"`
+	Data   *ExpvarzStatus `json:"data,omitempty"`
+	Error  *ApiError      `json:"error,omitempty"`
+}
+
 // statszReq is a request for us to respond with current statsz.
 func (s *Server) statszReq(sub *subscription, c *client, _ *Account, subject, reply string, hdr, msg []byte) {
 	if !s.EventsEnabled() {
diff --git a/server/jetstream_api.go b/server/jetstream_api.go
index 90a5838d..45a6ddc0 100644
--- a/server/jetstream_api.go
+++ b/server/jetstream_api.go
@@ -159,8 +159,9 @@ const (
 
 	// JSApiConsumerInfo is for obtaining general information about a consumer.
 	// Will return JSON response.
-	JSApiConsumerInfo  = "$JS.API.CONSUMER.INFO.*.*"
-	JSApiConsumerInfoT = "$JS.API.CONSUMER.INFO.%s.%s"
+	JSApiConsumerInfoPre = "$JS.API.CONSUMER.INFO."
+	JSApiConsumerInfo    = "$JS.API.CONSUMER.INFO.*.*"
+	JSApiConsumerInfoT   = "$JS.API.CONSUMER.INFO.%s.%s"
 
 	// JSApiConsumerDelete is the endpoint to delete consumers.
 	// Will return JSON response.
@@ -1033,6 +1034,15 @@ func (s *Server) sendAPIErrResponse(ci *ClientInfo, acc *Account, subject, reply
 	s.sendJetStreamAPIAuditAdvisory(ci, acc, subject, request, response)
 }
 
+// Use the account acc to send actual result from non-system account.
+func (s *Server) sendAPIErrResponseFromAccount(ci *ClientInfo, acc *Account, subject, reply, request, response string) {
+	acc.trackAPIErr()
+	if reply != _EMPTY_ {
+		s.sendInternalAccountMsg(acc, reply, response)
+	}
+	s.sendJetStreamAPIAuditAdvisory(ci, acc, subject, request, response)
+}
+
 const errRespDelay = 500 * time.Millisecond
 
 type delayedAPIResponse struct {
@@ -4591,6 +4601,55 @@ func (s *Server) jsConsumerListRequest(sub *subscription, c *client, _ *Account,
 	s.sendAPIResponse(ci, acc, subject, reply, string(msg), s.jsonResponse(resp))
 }
 
+// This will be a quick check on point of entry for a consumer that does
+// not exist. If that is the case we will return the response and return
+// true which will shortcut the service import to alleviate pressure on
+// the JS API queues.
+func (s *Server) jsConsumerProcessMissing(c *client, acc *Account) bool {
+	subject := bytesToString(c.pa.subject)
+	streamName, consumerName := streamNameFromSubject(subject), consumerNameFromSubject(subject)
+
+	// Check to make sure the consumer is assigned.
+	// All JS servers will have the meta information.
+	js, cc := s.getJetStreamCluster()
+	if js == nil || cc == nil {
+		return false
+	}
+	js.mu.RLock()
+	sa, ca := js.assignments(acc.Name, streamName, consumerName)
+	js.mu.RUnlock()
+
+	// If we have a consumer assignment return false here and let normally processing takeover.
+	if ca != nil {
+		return false
+	}
+
+	// We can't find the consumer, so mimic what would be the errors below.
+	var resp = JSApiConsumerInfoResponse{ApiResponse: ApiResponse{Type: JSApiConsumerInfoResponseType}}
+
+	// Need to make subject and reply real here for queued response processing.
+	subject = string(c.pa.subject)
+	reply := string(c.pa.reply)
+
+	ci := c.getClientInfo(true)
+
+	if hasJS, doErr := acc.checkJetStream(); !hasJS {
+		if doErr {
+			resp.Error = NewJSNotEnabledForAccountError()
+			s.sendAPIErrResponseFromAccount(ci, acc, subject, reply, _EMPTY_, s.jsonResponse(&resp))
+		}
+	} else if sa == nil {
+		resp.Error = NewJSStreamNotFoundError()
+		s.sendAPIErrResponseFromAccount(ci, acc, subject, reply, _EMPTY_, s.jsonResponse(&resp))
+	} else {
+		// If we are here the consumer is not present.
+		resp.Error = NewJSConsumerNotFoundError()
+		s.sendAPIErrResponseFromAccount(ci, acc, subject, reply, _EMPTY_, s.jsonResponse(&resp))
+	}
+
+	return true
+}
+
 // Request for information about an consumer.
 func (s *Server) jsConsumerInfoRequest(sub *subscription, c *client, _ *Account, subject, reply string, rmsg []byte) {
 	if c == nil || !s.JetStreamEnabled() {
diff --git a/server/jetstream_cluster.go b/server/jetstream_cluster.go
index fb01f22b..5bd3cf6c 100644
--- a/server/jetstream_cluster.go
+++ b/server/jetstream_cluster.go
@@ -4777,6 +4777,15 @@ func (js *jetStream) consumerAssignment(account, stream, consumer string) *consu
 	return nil
 }
 
+// Return both the stream and consumer assignments.
+// Lock should be held.
+func (js *jetStream) assignments(account, stream, consumer string) (*streamAssignment, *consumerAssignment) {
+	if sa := js.streamAssignment(account, stream); sa != nil {
+		return sa, sa.consumers[consumer]
+	}
+	return nil, nil
+}
+
 // consumerAssigned informs us if this server has this consumer assigned.
 func (jsa *jsAccount) consumerAssigned(stream, consumer string) bool {
 	jsa.mu.RLock()
diff --git a/server/websocket.go b/server/websocket.go
index 664acdee..7ebbaa85 100644
--- a/server/websocket.go
+++ b/server/websocket.go
@@ -67,6 +67,7 @@ const (
 	wsCloseStatusProtocolError      = 1002
 	wsCloseStatusUnsupportedData    = 1003
 	wsCloseStatusNoStatusReceived   = 1005
+	wsCloseStatusAbnormalClosure    = 1006
 	wsCloseStatusInvalidPayloadData = 1007
 	wsCloseStatusPolicyViolation    = 1008
 	wsCloseStatusMessageTooBig      = 1009
@@ -461,21 +462,9 @@ func (c *client) wsHandleControlFrame(r *wsReadInfo, frameType wsOpCode, nc io.R
 				}
 			}
 		}
-		// If the status indicates that nothing was received, then we don't
-		// send anything back.
-		// From https://datatracker.ietf.org/doc/html/rfc6455#section-7.4
-		// it says that code 1005 is a reserved value and MUST NOT be set as a
-		// status code in a Close control frame by an endpoint.  It is
-		// designated for use in applications expecting a status code to indicate
-		// that no status code was actually present.
-		var clm []byte
-		if status != wsCloseStatusNoStatusReceived {
-			clm = wsCreateCloseMessage(status, body)
-		}
+		clm := wsCreateCloseMessage(status, body)
 		c.wsEnqueueControlMessage(wsCloseMessage, clm)
-		if len(clm) > 0 {
-			nbPoolPut(clm) // wsEnqueueControlMessage has taken a copy.
-		}
+		nbPoolPut(clm) // wsEnqueueControlMessage has taken a copy.
 		// Return io.EOF so that readLoop will close the connection as ClientClosed
 		// after processing pending buffers.
 		return pos, io.EOF
@@ -662,11 +651,10 @@ func (c *client) wsEnqueueCloseMessage(reason ClosedState) {
 		status = wsCloseStatusProtocolError
 	case MaxPayloadExceeded:
 		status = wsCloseStatusMessageTooBig
-	case WriteError, ReadError, StaleConnection, ServerShutdown:
-		// We used to have WriteError, ReadError and StaleConnection result in
-		// code 1006, which the spec says that it must not be used to set the
-		// status in the close message. So using this one instead.
+	case ServerShutdown:
 		status = wsCloseStatusGoingAway
+	case WriteError, ReadError, StaleConnection:
+		status = wsCloseStatusAbnormalClosure
 	default:
 		status = wsCloseStatusInternalSrvError
 	}


Test patch:
diff --git a/server/jetstream_cluster_1_test.go b/server/jetstream_cluster_1_test.go
index 1676d508..c882d5a9 100644
--- a/server/jetstream_cluster_1_test.go
+++ b/server/jetstream_cluster_1_test.go
@@ -3412,7 +3412,6 @@ func TestJetStreamClusterExtendedAccountInfo(t *testing.T) {
 	// Go client will lag so use direct for now.
 	getAccountInfo := func() *nats.AccountInfo {
 		t.Helper()
-
 		info, err := js.AccountInfo()
 		if err != nil {
 			t.Fatalf("Unexpected error: %v", err)
@@ -3437,10 +3436,13 @@ func TestJetStreamClusterExtendedAccountInfo(t *testing.T) {
 	js.ConsumerInfo("TEST-2", "NO-CONSUMER")
 	js.ConsumerInfo("TEST-3", "NO-CONSUMER")
 
-	ai = getAccountInfo()
-	if ai.API.Errors != 4 {
-		t.Fatalf("Expected 4 API calls to be errors, got %d", ai.API.Errors)
-	}
+	checkFor(t, 2*time.Second, 250*time.Millisecond, func() error {
+		ai = getAccountInfo()
+		if ai.API.Errors != 4 {
+			return fmt.Errorf("Expected 4 API calls to be errors, got %d", ai.API.Errors)
+		}
+		return nil
+	})
 }
 
 func TestJetStreamClusterPeerRemovalAPI(t *testing.T) {
@@ -4319,7 +4321,8 @@ func TestJetStreamClusterNoQuorumStepdown(t *testing.T) {
 	if err := js.DeleteConsumer("NO-Q", "dlc"); !notAvailableErr(err) {
 		t.Fatalf("Expected an 'unavailable' error, got %v", err)
 	}
-	if _, err := js.ConsumerInfo("NO-Q", "dlc"); !notAvailableErr(err) {
+	// Since we did not create the consumer our bypass will respond from the local server.
+	if _, err := js.ConsumerInfo("NO-Q", "dlc"); err != nats.ErrConsumerNotFound {
 		t.Fatalf("Expected an 'unavailable' error, got %v", err)
 	}
 	// Listers
@@ -6900,48 +6903,6 @@ func TestJetStreamClusterCatchupMustStallWhenBehindOnApplies(t *testing.T) {
 	require_Len(t, len(msg.Data), 0)
 }
 
-func TestJetStreamClusterConsumerInfoAfterCreate(t *testing.T) {
-	c := createJetStreamClusterExplicit(t, "R3S", 3)
-	defer c.shutdown()
-
-	nl := c.randomNonLeader()
-	nc, js := jsClientConnect(t, nl)
-	defer nc.Close()
-
-	_, err := js.AddStream(&nats.StreamConfig{
-		Name:     "TEST",
-		Subjects: []string{"foo"},
-		Replicas: 3,
-	})
-	require_NoError(t, err)
-
-	// We pause applies for the server we're connected to.
-	// This is fine for the RAFT log and allowing the consumer to be created,
-	// but we will not be able to apply the consumer assignment for some time.
-	mjs := nl.getJetStream()
-	require_NotNil(t, js)
-	mg := mjs.getMetaGroup()
-	require_NotNil(t, mg)
-	err = mg.(*raft).PauseApply()
-	require_NoError(t, err)
-
-	// Add consumer.
-	_, err = js.AddConsumer("TEST", &nats.ConsumerConfig{Durable: "CONSUMER"})
-	require_NoError(t, err)
-
-	// Consumer info should not fail, this server should not short-circuit because
-	// it was not able to apply the consumer assignment.
-	_, err = js.ConsumerInfo("TEST", "CONSUMER")
-	require_NoError(t, err)
-
-	// Resume applies.
-	mg.(*raft).ResumeApply()
-
-	// Check consumer info still works.
-	_, err = js.ConsumerInfo("TEST", "CONSUMER")
-	require_NoError(t, err)
-}
-
 //
 // DO NOT ADD NEW TESTS IN THIS FILE (unless to balance test times)
 // Add at the end of jetstream_cluster_<n>_test.go, with <n> being the highest value.
diff --git a/server/norace_test.go b/server/norace_test.go
index 81209dc1..58594241 100644
--- a/server/norace_test.go
+++ b/server/norace_test.go
@@ -11267,3 +11267,60 @@ func TestNoRaceJetStreamClusterLargeMetaSnapshotTiming(t *testing.T) {
 	require_NoError(t, n.InstallSnapshot(snap))
 	t.Logf("Took %v to snap meta with size of %v\n", time.Since(start), friendlyBytes(len(snap)))
 }
+
+func TestNoRaceJetStreamClusterInfoOnMissingConsumers(t *testing.T) {
+	c := createJetStreamClusterExplicit(t, "R3F", 3)
+	defer c.shutdown()
+
+	nc, js := jsClientConnect(t, c.randomServer())
+	defer nc.Close()
+
+	// Create a stream just so the consumer info processing misses on the consumer only.
+	_, err := js.AddStream(&nats.StreamConfig{
+		Name:     "TEST",
+		Subjects: []string{"foo"},
+		Replicas: 3,
+	})
+	require_NoError(t, err)
+
+	done := make(chan bool)
+	pending := make(chan int, 1)
+
+	// Check to make sure we never have any pending on the API queue.
+	go func() {
+		ml := c.leader()
+		for {
+			select {
+			case <-done:
+				return
+			case <-time.After(100 * time.Millisecond):
+				qlen := ml.jsAPIRoutedReqs.len() + int(ml.jsAPIRoutedReqs.inProgress())
+				if qlen > 0 {
+					pending <- qlen
+					return
+				}
+			}
+		}
+	}()
+
+	wg := sync.WaitGroup{}
+	wg.Add(500)
+	for i := 0; i < 500; i++ {
+		go func() {
+			defer wg.Done()
+			s := c.randomServer()
+			nc, js := jsClientConnect(t, s)
+			defer nc.Close()
+			// Check for non-existent consumers.
+			for c := 0; c < 1000; c++ {
+				_, err := js.ConsumerInfo("TEST", fmt.Sprintf("C-%d", c))
+				require_Error(t, err)
+			}
+		}()
+	}
+	wg.Wait()
+	close(done)
+	if len(pending) > 0 {
+		t.Fatalf("Saw API pending of %d, expected always 0", <-pending)
+	}
+}
diff --git a/server/websocket_test.go b/server/websocket_test.go
index b1f0ac03..2c8001b8 100644
--- a/server/websocket_test.go
+++ b/server/websocket_test.go
@@ -646,24 +646,19 @@ func TestWSReadPongFrame(t *testing.T) {
 
 func TestWSReadCloseFrame(t *testing.T) {
 	for _, test := range []struct {
-		name     string
-		noStatus bool
-		payload  []byte
+		name    string
+		payload []byte
 	}{
-		{"status without payload", false, nil},
-		{"status with payload", false, []byte("optional payload")},
-		{"no status no payload", true, nil},
+		{"without payload", nil},
+		{"with payload", []byte("optional payload")},
 	} {
 		t.Run(test.name, func(t *testing.T) {
 			c, ri, tr := testWSSetupForRead()
-			var payload []byte
-			if !test.noStatus {
-				// a close message has a status in 2 bytes + optional payload
-				payload = make([]byte, 2+len(test.payload))
-				binary.BigEndian.PutUint16(payload[:2], wsCloseStatusNormalClosure)
-				if len(test.payload) > 0 {
-					copy(payload[2:], test.payload)
-				}
+			// a close message has a status in 2 bytes + optional payload
+			payload := make([]byte, 2+len(test.payload))
+			binary.BigEndian.PutUint16(payload[:2], wsCloseStatusNormalClosure)
+			if len(test.payload) > 0 {
+				copy(payload[2:], test.payload)
 			}
 			close := testWSCreateClientMsg(wsCloseMessage, 1, true, false, payload)
 			// Have a normal frame prior to close to make sure that wsRead returns
@@ -689,11 +684,7 @@ func TestWSReadCloseFrame(t *testing.T) {
 			if n := len(nb); n == 0 {
 				t.Fatalf("Expected buffers, got %v", n)
 			}
-			if test.noStatus {
-				if expected := 2; expected != len(nb[0]) {
-					t.Fatalf("Expected buffer to be %v bytes long, got %v", expected, len(nb[0]))
-				}
-			} else if expected := 2 + 2 + len(test.payload); expected != len(nb[0]) {
+			if expected := 2 + 2 + len(test.payload); expected != len(nb[0]) {
 				t.Fatalf("Expected buffer to be %v bytes long, got %v", expected, len(nb[0]))
 			}
 			b := nb[0][0]
@@ -703,14 +694,12 @@ func TestWSReadCloseFrame(t *testing.T) {
 			if b&byte(wsCloseMessage) == 0 {
 				t.Fatalf("Should have been a CLOSE, it wasn't: %v", b)
 			}
-			if !test.noStatus {
-				if status := binary.BigEndian.Uint16(nb[0][2:4]); status != wsCloseStatusNormalClosure {
-					t.Fatalf("Expected status to be %v, got %v", wsCloseStatusNormalClosure, status)
-				}
-				if len(test.payload) > 0 {
-					if !bytes.Equal(nb[0][4:], test.payload) {
-						t.Fatalf("Unexpected content: %s", nb[0][4:])
-					}
+			if status := binary.BigEndian.Uint16(nb[0][2:4]); status != wsCloseStatusNormalClosure {
+				t.Fatalf("Expected status to be %v, got %v", wsCloseStatusNormalClosure, status)
+			}
+			if len(test.payload) > 0 {
+				if !bytes.Equal(nb[0][4:], test.payload) {
+					t.Fatalf("Unexpected content: %s", nb[0][4:])
 				}
 			}
 		})
@@ -788,6 +777,7 @@ func TestWSCloseFrameWithPartialOrInvalid(t *testing.T) {
 
 	// Now test close with invalid status size (1 instead of 2 bytes)
 	c, ri, tr = testWSSetupForRead()
+	payload[0] = 100
 	binary.BigEndian.PutUint16(payload, wsCloseStatusNormalClosure)
 	closeMsg = testWSCreateClientMsg(wsCloseMessage, 1, true, false, payload[:1])
 
@@ -804,15 +794,14 @@ func TestWSCloseFrameWithPartialOrInvalid(t *testing.T) {
 	if n := len(bufs); n != 0 {
 		t.Fatalf("Unexpected buffer returned: %v", n)
 	}
-	// Since no status was received, the server will send a close frame without
-	// status code nor payload.
+	// A CLOSE should have been queued with the payload of the original close message.
 	c.mu.Lock()
 	nb, _ = c.collapsePtoNB()
 	c.mu.Unlock()
 	if n := len(nb); n == 0 {
 		t.Fatalf("Expected buffers, got %v", n)
 	}
-	if expected := 2; expected != len(nb[0]) {
+	if expected := 2 + 2; expected != len(nb[0]) {
 		t.Fatalf("Expected buffer to be %v bytes long, got %v", expected, len(nb[0]))
 	}
 	b = nb[0][0]
@@ -822,6 +811,13 @@ func TestWSCloseFrameWithPartialOrInvalid(t *testing.T) {
 	if b&byte(wsCloseMessage) == 0 {
 		t.Fatalf("Should have been a CLOSE, it wasn't: %v", b)
 	}
+	// Since satus was not valid, we should get wsCloseStatusNoStatusReceived
+	if status := binary.BigEndian.Uint16(nb[0][2:4]); status != wsCloseStatusNoStatusReceived {
+		t.Fatalf("Expected status to be %v, got %v", wsCloseStatusNoStatusReceived, status)
+	}
+	if len(nb[0][:]) != 4 {
+		t.Fatalf("Unexpected content: %s", nb[0][2:])
+	}
 }
 
 func TestWSReadGetErrors(t *testing.T) {
@@ -1018,9 +1014,9 @@ func TestWSEnqueueCloseMsg(t *testing.T) {
 		{BadClientProtocolVersion, wsCloseStatusProtocolError},
 		{MaxPayloadExceeded, wsCloseStatusMessageTooBig},
 		{ServerShutdown, wsCloseStatusGoingAway},
-		{WriteError, wsCloseStatusGoingAway},
-		{ReadError, wsCloseStatusGoingAway},
-		{StaleConnection, wsCloseStatusGoingAway},
+		{WriteError, wsCloseStatusAbnormalClosure},
+		{ReadError, wsCloseStatusAbnormalClosure},
+		{StaleConnection, wsCloseStatusAbnormalClosure},
 		{ClosedState(254), wsCloseStatusInternalSrvError},
 	} {
 		t.Run(test.reason.String(), func(t *testing.T) {