nats-server websockets are vulnerable to pre-auth memory DoS

@@ -31,6 +31,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 	"unicode/utf8"
 
@@ -211,6 +212,7 @@ func (c *client) wsRead(r *wsReadInfo, ior io.Reader, buf []byte) ([][]byte, err
 		err    error
 		pos    int
 		max    = len(buf)
+		mpay   = int(atomic.LoadInt32(&c.mpay))
 	)
 	for pos != max {
 		if r.fs {
@@ -324,7 +326,7 @@ func (c *client) wsRead(r *wsReadInfo, ior io.Reader, buf []byte) ([][]byte, err
 				// When we have the final frame and we have read the full payload,
 				// we can decompress it.
 				if r.ff && r.rem == 0 {
-					b, err = r.decompress()
+					b, err = r.decompress(mpay)
 					if err != nil {
 						return bufs, err
 					}
@@ -398,7 +400,16 @@ func (r *wsReadInfo) ReadByte() (byte, error) {
 	return b, nil
 }
 
-func (r *wsReadInfo) decompress() ([]byte, error) {
+// decompress decompresses the collected buffers.
+// The size of the decompressed buffer will be limited to the `mpay` value.
+// If, while decompressing, the resulting uncompressed buffer exceeds this
+// limit, the decompression stops and an empty buffer and the ErrMaxPayload
+// error are returned.
+func (r *wsReadInfo) decompress(mpay int) ([]byte, error) {
+	// If not limit is specified, use the default maximum payload size.
+	if mpay <= 0 {
+		mpay = MAX_PAYLOAD_SIZE
+	}
 	r.coff = 0
 	// As per https://tools.ietf.org/html/rfc7692#section-7.2.2
 	// add 0x00, 0x00, 0xff, 0xff and then a final block so that flate reader
@@ -413,8 +424,15 @@ func (r *wsReadInfo) decompress() ([]byte, error) {
 	} else {
 		d.(flate.Resetter).Reset(r, nil)
 	}
-	// This will do the decompression.
-	b, err := io.ReadAll(d)
+	// Use a LimitedReader to limit the decompressed size.
+	// We use "limit+1" bytes for "N" so we can detect if the limit is exceeded.
+	lr := io.LimitedReader{R: d, N: int64(mpay + 1)}
+	b, err := io.ReadAll(&lr)
+	if err == nil && len(b) > mpay {
+		// Decompressed data exceeds the maximum payload size.
+		b, err = nil, ErrMaxPayload
+	}
+	lr.R = nil
 	decompressorPool.Put(d)
 	// Now reset the compressed buffers list.
 	r.cbufs = nil

@@ -4473,6 +4473,98 @@ func TestWSNoCorruptionWithFrameSizeLimit(t *testing.T) {
 	testWSNoCorruptionWithFrameSizeLimit(t, 1000)
 }
 
+func TestWSDecompressLimit(t *testing.T) {
+	for _, test := range []struct {
+		name    string
+		mpayCfg string
+	}{
+		{"not explicitly configured", _EMPTY_},
+		{"explicit high", "max_payload: 2097152"},
+		{"explicit low", "max_payload: 4096"},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			conf := createConfFile(t, fmt.Appendf(nil, `
+				listen: "127.0.0.1:-1"
+				websocket {
+					listen: "127.0.0.1:-1"
+					no_tls: true
+				}
+				%s
+			`, test.mpayCfg))
+			s, o := RunServerWithConfig(conf)
+			defer s.Shutdown()
+
+			l := &captureErrorLogger{errCh: make(chan string, 10)}
+			s.SetLogger(l, false, false)
+
+			// Create a client that will use compression.
+			wsc, br, _ := testNewWSClient(t, testWSClientOptions{
+				compress: true,
+				host:     o.Websocket.Host,
+				port:     o.Websocket.Port,
+				noTLS:    true,
+			})
+			// We will hand-craft a frame that would use a 10MB of uncompressed zeros
+			// that should compress really small.
+			buf := &bytes.Buffer{}
+			compressor, _ := flate.NewWriter(buf, 1)
+			chunk := make([]byte, 1024*1024)
+			// Compress the equivalent of 100MB of data. We do by chunks to limit
+			// memory usage here.
+			for range 100 {
+				compressor.Write(chunk)
+			}
+			compressor.Flush()
+			payload := buf.Bytes()
+			// The last 4 bytes are dropped
+			payload = payload[:len(payload)-4]
+			lenPayload := len(payload)
+			frame := make([]byte, 14+lenPayload)
+			frame[0] = byte(wsBinaryMessage)
+			frame[0] |= wsFinalBit
+			frame[0] |= wsRsv1Bit
+			pos := 1
+			switch {
+			case lenPayload <= 125:
+				frame[pos] = byte(lenPayload) | wsMaskBit
+				pos++
+			case lenPayload < 65536:
+				frame[pos] = 126 | wsMaskBit
+				binary.BigEndian.PutUint16(frame[2:], uint16(lenPayload))
+				pos += 3
+			default:
+				frame[1] = 127 | wsMaskBit
+				binary.BigEndian.PutUint64(frame[2:], uint64(lenPayload))
+				pos += 9
+			}
+			key := []byte{1, 2, 3, 4}
+			copy(frame[pos:], key)
+			pos += 4
+			copy(frame[pos:], payload)
+			testWSSimpleMask(key, frame[pos:])
+			pos += lenPayload
+			toSend := frame[:pos]
+			if _, err := wsc.Write(toSend); err != nil {
+				t.Fatalf("Error sending message: %v", err)
+			}
+
+			// We should have been disconnected.
+			rbuf := make([]byte, 1024)
+			_, err := br.Read(rbuf)
+			require_Error(t, err)
+
+			select {
+			case err := <-l.errCh:
+				if !strings.Contains(err, ErrMaxPayload.Error()) {
+					t.Fatalf("Expected %s error, got %s", ErrMaxPayload, err)
+				}
+			case <-time.After(time.Second):
+				t.Fatal("Did not get the expected error")
+			}
+		})
+	}
+}
+
 // ==================================================================
 // = Benchmark tests
 // ==================================================================