Repository: nats-io/nats-server (go)

Issue:
Corrupted messages when websocket consumer is active
### Observed behavior

I've observed both reordering of the bytes within the message, as well as extra bytes being added (e.g. `PONG`  added at a random point within the message). This is not consistent for the same message - if a message is published to jetstream and then repeatedly requested, corruption may only occur some of the time and the corrupted message may be different each time (even though no new publishes have happened).

This only happens when a websocket client is connected and receiving messages, but messages sent to other clients (e.g. nats-cli) during this time may also be corrupted. Other clients return to normal ~30s after the websocket client is closed.

Only occurs for fairly specific message sizes - I can reliably reproduce with ~790000 bytes.

### Expected behavior

Messages always arrive intact.

### Server and client version

nats-server 2.10.5, still occurs on `main`

Bisected to 6b65452bc70973d9bcf59857afa94ae46b04bdf4

Clients:
nats-cli 0.1.1
nats.ws 1.19.1

### Host environment

Arch Linux x86_64

Have reproduced on multiple machines with different specs.

### Steps to reproduce

<details>
<summary>
Server config
</summary>

```
listen: localhost:4222
max_payload: 8MB

http: localhost:8222

jetstream {
    store_dir: store
}

websocket {
    listen: localhost:8443
    no_tls: true
}
```
</details>

<details>
<summary>
Stream config
</summary>

```
{
    "name": "test",
    "subjects": [
      "test"
    ],
    "retention": "limits",
    "max_consumers": -1,
    "max_msgs_per_subject": -1,
    "max_msgs": -1,
    "max_bytes": -1,
    "max_age": 604800000000000,
    "max_msg_size": -1,
    "storage": "file",
    "discard": "old",
    "num_replicas": 1,
    "duplicate_window": 120000000000,
    "sealed": false,
    "deny_delete": false,
    "deny_purge": false,
    "allow_rollup_hdrs": false,
    "allow_direct": false,
    "mirror_direct": false
}
```
</details>

<details>
<summary>
Websocket client:
</summary>

```
<!DOCTYPE html>
<script type="module">
import { connect, DeliverPolicy } from "https://unpkg.com/nats.ws@1.19.1/esm/nats.js";
const nats = await connect({
    servers: "ws://localhost:8443",
    waitOnFirstConnect: true,
    // no limit
    maxReconnectAttempts: -1,
});

for (;;) {
    const consumer = await nats
        .jetstream()
        .consumers.get("test", {
            filterSubjects: "test",
            deliver_policy: DeliverPolicy.Last,
        });
    for await (const msg of await consumer.consume()) {
        console.log("received");
    }
}
</script>
```
</details>

Steps
```
# generate 79000 random bytes
cat /dev/urandom | base64 | head -c 790000 > rand.dat

# publish to jetstream
nats publish test < rand.dat

# keep trying to read the message until it is corrupted
while nats subscribe --raw --count 1 --last-per-subject test | diff -qw - rand.dat ; do :; done

# open websocket client, observe that above loop stops as message is corrupted
```

Pull request description:
Fix data race in WebSockets (fixes #4808)
In certain conditions, a WebSocket with a maximum frame size (i.e. web browsers) and specific message sizes could result in multiple overlapping slices being returned to the buffer pool. This would result in a data race that would result in data corruption on the wire.

This condition shouldn't affect WebSocket connections that don't have a maximum frame size.

Signed-off-by: Neil Twigg <neil@nats.io>

Patch:
diff --git a/server/client.go b/server/client.go
index 72346e69..e4078416 100644
--- a/server/client.go
+++ b/server/client.go
@@ -336,6 +336,9 @@ var nbPoolLarge = &sync.Pool{
 	},
 }
 
+// nbPoolGet returns a frame that is a best-effort match for the given size.
+// Once a pooled frame is no longer needed, it should be recycled by passing
+// it to nbPoolPut.
 func nbPoolGet(sz int) []byte {
 	switch {
 	case sz <= nbPoolSizeSmall:
@@ -347,6 +350,10 @@ func nbPoolGet(sz int) []byte {
 	}
 }
 
+// nbPoolPut recycles a frame that was retrieved from nbPoolGet. It is not
+// safe to return multiple slices referring to chunks of the same underlying
+// array as this may create overlaps when the buffers are returned to their
+// original size, resulting in race conditions.
 func nbPoolPut(b []byte) {
 	switch cap(b) {
 	case nbPoolSizeSmall:
@@ -1490,8 +1497,8 @@ func closedStateForErr(err error) ClosedState {
 	return ReadError
 }
 
-// collapsePtoNB will place primary onto nb buffer as needed in prep for WriteTo.
-// This will return a copy on purpose.
+// collapsePtoNB will either returned framed WebSocket buffers or it will
+// return a reference to c.out.nb.
 func (c *client) collapsePtoNB() (net.Buffers, int64) {
 	if c.isWebsocket() {
 		return c.wsCollapsePtoNB()
diff --git a/server/websocket.go b/server/websocket.go
index 0f45f91e..e026674d 100644
--- a/server/websocket.go
+++ b/server/websocket.go
@@ -1269,13 +1269,12 @@ func (s *Server) createWSClient(conn net.Conn, ws *websocket) *client {
 }
 
 func (c *client) wsCollapsePtoNB() (net.Buffers, int64) {
-	var nb net.Buffers
+	nb := c.out.nb
 	var mfs int
 	var usz int
 	if c.ws.browser {
 		mfs = wsFrameSizeForBrowsers
 	}
-	nb = c.out.nb
 	mask := c.ws.maskwrite
 	// Start with possible already framed buffers (that we could have
 	// got from partials or control messages such as ws pings or pongs).
@@ -1378,8 +1377,10 @@ func (c *client) wsCollapsePtoNB() (net.Buffers, int64) {
 			for i := 0; i < len(nb); i++ {
 				b := nb[i]
 				if total+len(b) <= mfs {
-					bufs = append(bufs, b)
+					buf := nbPoolGet(len(b))
+					bufs = append(bufs, append(buf, b...))
 					total += len(b)
+					nbPoolPut(nb[i])
 					continue
 				}
 				for len(b) > 0 {
@@ -1394,9 +1395,11 @@ func (c *client) wsCollapsePtoNB() (net.Buffers, int64) {
 					if endStart {
 						fhIdx = startFrame()
 					}
-					bufs = append(bufs, b[:total])
+					buf := nbPoolGet(total)
+					bufs = append(bufs, append(buf, b[:total]...))
 					b = b[total:]
 				}
+				nbPoolPut(nb[i]) // No longer needed as copied into smaller frames.
 			}
 			if total > 0 {
 				endFrame(fhIdx, total)


Test patch:
diff --git a/server/norace_test.go b/server/norace_test.go
index 06a8d2b3..a773751c 100644
--- a/server/norace_test.go
+++ b/server/norace_test.go
@@ -9141,3 +9141,7 @@ func TestNoRaceFileStoreLargeMsgsAndFirstMatching(t *testing.T) {
 	fs.LoadNextMsg("*.baz.*", true, fseq, nil)
 	require_True(t, time.Since(start) < 200*time.Microsecond)
 }
+
+func TestNoRaceWSNoCorruptionWithFrameSizeLimit(t *testing.T) {
+	testWSNoCorruptionWithFrameSizeLimit(t, 50000)
+}
diff --git a/server/websocket_test.go b/server/websocket_test.go
index 3fdb1fdb..5853731d 100644
--- a/server/websocket_test.go
+++ b/server/websocket_test.go
@@ -30,10 +30,12 @@ import (
 	"reflect"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
 	"github.com/nats-io/jwt/v2"
+	"github.com/nats-io/nats.go"
 	"github.com/nats-io/nkeys"
 
 	"github.com/klauspost/compress/flate"
@@ -4066,6 +4068,123 @@ func TestWSWithPartialWrite(t *testing.T) {
 	}
 }
 
+func testWSNoCorruptionWithFrameSizeLimit(t *testing.T, total int) {
+	tmpl := `
+               listen: "127.0.0.1:-1"
+               cluster {
+                       name: "local"
+                       port: -1
+                       %s
+               }
+               websocket {
+                       listen: "127.0.0.1:-1"
+                       no_tls: true
+               }
+       `
+	conf1 := createConfFile(t, []byte(fmt.Sprintf(tmpl, _EMPTY_)))
+	s1, o1 := RunServerWithConfig(conf1)
+	defer s1.Shutdown()
+
+	routes := fmt.Sprintf("routes: [\"nats://127.0.0.1:%d\"]", o1.Cluster.Port)
+	conf2 := createConfFile(t, []byte(fmt.Sprintf(tmpl, routes)))
+	s2, o2 := RunServerWithConfig(conf2)
+	defer s2.Shutdown()
+
+	conf3 := createConfFile(t, []byte(fmt.Sprintf(tmpl, routes)))
+	s3, o3 := RunServerWithConfig(conf3)
+	defer s3.Shutdown()
+
+	checkClusterFormed(t, s1, s2, s3)
+
+	nc3 := natsConnect(t, fmt.Sprintf("ws://127.0.0.1:%d", o3.Websocket.Port))
+	defer nc3.Close()
+
+	nc2 := natsConnect(t, fmt.Sprintf("ws://127.0.0.1:%d", o2.Websocket.Port))
+	defer nc2.Close()
+
+	payload := make([]byte, 100000)
+	for i := 0; i < len(payload); i++ {
+		payload[i] = 'A' + byte(i%26)
+	}
+	errCh := make(chan error, 1)
+	doneCh := make(chan struct{}, 1)
+	count := int32(0)
+
+	createSub := func(nc *nats.Conn) {
+		sub := natsSub(t, nc, "foo", func(m *nats.Msg) {
+			if !bytes.Equal(m.Data, payload) {
+				stop := len(m.Data)
+				if l := len(payload); l < stop {
+					stop = l
+				}
+				start := 0
+				for i := 0; i < stop; i++ {
+					if m.Data[i] != payload[i] {
+						start = i
+						break
+					}
+				}
+				if stop-start > 20 {
+					stop = start + 20
+				}
+				select {
+				case errCh <- fmt.Errorf("Invalid message: [%d bytes same]%s[...]", start, m.Data[start:stop]):
+				default:
+				}
+				return
+			}
+			if n := atomic.AddInt32(&count, 1); int(n) == 2*total {
+				doneCh <- struct{}{}
+			}
+		})
+		sub.SetPendingLimits(-1, -1)
+	}
+	createSub(nc2)
+	createSub(nc3)
+
+	checkSubInterest(t, s1, globalAccountName, "foo", time.Second)
+
+	nc1 := natsConnect(t, fmt.Sprintf("ws://127.0.0.1:%d", o1.Websocket.Port))
+	defer nc1.Close()
+	natsFlush(t, nc1)
+
+	// Change websocket connections to force a max frame size.
+	for _, s := range []*Server{s1, s2, s3} {
+		s.mu.RLock()
+		for _, c := range s.clients {
+			c.mu.Lock()
+			if c.ws != nil {
+				c.ws.browser = true
+			}
+			c.mu.Unlock()
+		}
+		s.mu.RUnlock()
+	}
+
+	for i := 0; i < total; i++ {
+		natsPub(t, nc1, "foo", payload)
+		if i%100 == 0 {
+			select {
+			case err := <-errCh:
+				t.Fatalf("Error: %v", err)
+			default:
+			}
+		}
+	}
+	select {
+	case err := <-errCh:
+		t.Fatalf("Error: %v", err)
+	case <-doneCh:
+		return
+	case <-time.After(10 * time.Second):
+		t.Fatalf("Test timed out")
+	}
+}
+
+func TestWSNoCorruptionWithFrameSizeLimit(t *testing.T) {
+	testWSNoCorruptionWithFrameSizeLimit(t, 1000)
+}
+
 // ==================================================================
 // = Benchmark tests
 // ==================================================================