Files
mesh-controller/vendor/github.com/nats-io/nats-server/v2/server/websocket.go
T
jochen be92762969 Give every test a bus of its own, at the release the mesh runs (hq ADR 0237)
The live tests reached one shared bus and assert, read and remove the mesh's own objects by
their fixed names, so packages run in parallel deleted what each other read and the suite
passed only one package at a time; a red suite read as noise. internal/testbus starts a server
per test, linked in at the nats-server release go.mod pins, and a test holds that pin to the
catalogue's bus image and to the facts snapshot's bus when there is one, so the tests never run
a bus the mesh does not. The waiter test read a timing (the most connections held at one look)
and now reads the state it means (the fewest held across the wait). make check runs the packages
in parallel under the race detector, with a timeout.
2026-10-06 21:17:02 +02:00

1666 lines
47 KiB
Go

// Copyright 2020-2025 The NATS Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package server
import (
"bytes"
crand "crypto/rand"
"crypto/sha1"
"crypto/tls"
"encoding/base64"
"encoding/binary"
"errors"
"fmt"
"io"
"log"
mrand "math/rand"
"net"
"net/http"
"net/url"
"strconv"
"strings"
"sync"
"sync/atomic"
"time"
"unicode/utf8"
"github.com/klauspost/compress/flate"
)
type wsOpCode int
const (
// From https://tools.ietf.org/html/rfc6455#section-5.2
wsTextMessage = wsOpCode(1)
wsBinaryMessage = wsOpCode(2)
wsCloseMessage = wsOpCode(8)
wsPingMessage = wsOpCode(9)
wsPongMessage = wsOpCode(10)
wsFinalBit = 1 << 7
wsRsv1Bit = 1 << 6 // Used for compression, from https://tools.ietf.org/html/rfc7692#section-6
wsRsv2Bit = 1 << 5
wsRsv3Bit = 1 << 4
wsMaskBit = 1 << 7
wsContinuationFrame = 0
wsMaxFrameHeaderSize = 14 // Since LeafNode may need to behave as a client
wsMaxControlPayloadSize = 125
wsFrameSizeForBrowsers = 4096 // From experiment, webrowsers behave better with limited frame size
wsCompressThreshold = 64 // Don't compress for small buffer(s)
wsMaxMsgPayloadMultiple = 8
wsMaxMsgPayloadLimit = 64 * 1024 * 1024
wsCloseSatusSize = 2
// From https://tools.ietf.org/html/rfc6455#section-11.7
wsCloseStatusNormalClosure = 1000
wsCloseStatusGoingAway = 1001
wsCloseStatusProtocolError = 1002
wsCloseStatusUnsupportedData = 1003
wsCloseStatusNoStatusReceived = 1005
wsCloseStatusInvalidPayloadData = 1007
wsCloseStatusPolicyViolation = 1008
wsCloseStatusMessageTooBig = 1009
wsCloseStatusInternalSrvError = 1011
wsCloseStatusTLSHandshake = 1015
wsFirstFrame = true
wsContFrame = false
wsFinalFrame = true
wsUncompressedFrame = false
wsSchemePrefix = "ws"
wsSchemePrefixTLS = "wss"
wsNoMaskingHeader = "Nats-No-Masking"
wsNoMaskingValue = "true"
wsXForwardedForHeader = "X-Forwarded-For"
wsNoMaskingFullResponse = wsNoMaskingHeader + ": " + wsNoMaskingValue + CR_LF
wsPMCExtension = "permessage-deflate" // per-message compression
wsPMCSrvNoCtx = "server_no_context_takeover"
wsPMCCliNoCtx = "client_no_context_takeover"
wsPMCReqHeaderValue = wsPMCExtension + "; " + wsPMCSrvNoCtx + "; " + wsPMCCliNoCtx
wsPMCFullResponse = "Sec-WebSocket-Extensions: " + wsPMCExtension + "; " + wsPMCSrvNoCtx + "; " + wsPMCCliNoCtx + _CRLF_
wsSecProto = "Sec-Websocket-Protocol"
wsMQTTSecProtoVal = "mqtt"
wsMQTTSecProto = wsSecProto + ": " + wsMQTTSecProtoVal + CR_LF
)
var decompressorPool sync.Pool
var compressLastBlock = []byte{0x00, 0x00, 0xff, 0xff, 0x01, 0x00, 0x00, 0xff, 0xff}
// From https://tools.ietf.org/html/rfc6455#section-1.3
var wsGUID = []byte("258EAFA5-E914-47DA-95CA-C5AB0DC85B11")
// Test can enable this so that server does not support "no-masking" requests.
var wsTestRejectNoMasking = false
type websocket struct {
frames net.Buffers
fs int64
closeMsg []byte
compress bool
closeSent bool
browser bool
nocompfrag bool // No fragment for compressed frames
maskread bool
maskwrite bool
compressor *flate.Writer
cookieJwt string
cookieUsername string
cookiePassword string
cookieToken string
clientIP string
}
type srvWebsocket struct {
mu sync.RWMutex
server *http.Server
listener net.Listener
listenerErr error
allowedOrigins map[string][]*allowedOrigin // host will be the key
sameOrigin bool
connectURLs []string
connectURLsMap refCountedUrlSet
authOverride bool // indicate if there is auth override in websocket config
rawHeaders string // raw headers to be used in the upgrade response.
// These are immutable and can be accessed without lock.
// This is the case when generating the client INFO.
tls bool // True if TLS is required (TLSConfig is specified).
host string // Host/IP the webserver is listening on (shortcut to opts.Websocket.Host).
port int // Port the webserver is listening on. This is after an ephemeral port may have been selected (shortcut to opts.Websocket.Port).
}
type allowedOrigin struct {
scheme string
port string
}
type wsUpgradeResult struct {
conn net.Conn
ws *websocket
kind int
}
type wsReadInfo struct {
rem uint64
fs bool
ff bool
fc bool
mask bool // Incoming leafnode connections may not have masking.
mkpos byte
mkey [4]byte
cbufs [][]byte
coff int
csz uint64
}
func (r *wsReadInfo) init() {
r.fs, r.ff = true, true
}
func (r *wsReadInfo) resetCompressedState() {
r.fs = true
r.ff = true
r.fc = false
r.rem = 0
r.cbufs = nil
r.coff = 0
r.csz = 0
}
// Compressed WebSocket messages have to be accumulated before they can be
// decompressed and handed to the parser, so this transport limit needs to
// allow batching several max_payload-sized NATS operations while still
// capping resource usage on the buffered compressed path.
func wsMaxMessageSize(mpay int) uint64 {
if mpay <= 0 {
mpay = MAX_PAYLOAD_SIZE
}
limit := uint64(mpay) * wsMaxMsgPayloadMultiple
if limit > wsMaxMsgPayloadLimit {
limit = wsMaxMsgPayloadLimit
}
return limit
}
// Returns a slice containing `needed` bytes from the given buffer `buf`
// starting at position `pos`, and possibly read from the given reader `r`.
// When bytes are present in `buf`, the `pos` is incremented by the number
// of bytes found up to `needed` and the new position is returned. If not
// enough bytes are found, the bytes found in `buf` are copied to the returned
// slice and the remaning bytes are read from `r`.
func wsGet(r io.Reader, buf []byte, pos, needed uint64) ([]byte, uint64, error) {
avail := uint64(len(buf)) - pos
if avail >= needed {
return buf[pos : pos+needed], pos + needed, nil
}
b := make([]byte, needed)
start := uint64(copy(b, buf[pos:]))
for start != needed {
n, err := r.Read(b[start:cap(b)])
if err != nil {
return nil, 0, err
}
start += uint64(n)
}
return b, pos + avail, nil
}
// Returns true if this connection is from a Websocket client.
// Lock held on entry.
func (c *client) isWebsocket() bool {
return c.ws != nil
}
// Returns a slice of byte slices corresponding to payload of websocket frames.
// The byte slice `buf` is filled with bytes from the connection's read loop.
// This function will decode the frame headers and unmask the payload(s).
// It is possible that the returned slices point to the given `buf` slice, so
// `buf` should not be overwritten until the returned slices have been parsed.
//
// Client lock MUST NOT be held on entry.
func (c *client) wsRead(r *wsReadInfo, ior io.Reader, buf []byte) ([][]byte, error) {
var bufs [][]byte
err := c.wsReadLoop(r, ior, buf, func(b []byte, compressed, final bool) error {
if compressed {
return errors.New("compressed websocket frames require wsReadAndParse")
}
bufs = append(bufs, b)
return nil
})
return bufs, err
}
func (c *client) wsReadAndParse(r *wsReadInfo, ior io.Reader, buf []byte) error {
mpay := int(atomic.LoadInt32(&c.mpay))
if mpay <= 0 {
mpay = MAX_PAYLOAD_SIZE
}
return c.wsReadLoop(r, ior, buf, func(b []byte, compressed, final bool) error {
if compressed {
if err := c.wsDecompressAndParse(r, b, final, mpay); err != nil {
r.resetCompressedState()
return err
}
if final {
r.fc = false
}
return nil
}
return c.parse(b)
})
}
func (c *client) wsReadLoop(r *wsReadInfo, ior io.Reader, buf []byte, handle func([]byte, bool, bool) error) error {
var (
tmpBuf []byte
err error
pos uint64
max = uint64(len(buf))
)
for pos != max {
if r.fs {
b0 := buf[pos]
frameType := wsOpCode(b0 & 0xF)
final := b0&wsFinalBit != 0
compressed := b0&wsRsv1Bit != 0
if b0&(wsRsv2Bit|wsRsv3Bit) != 0 {
return c.wsHandleProtocolError("RSV2 and RSV3 must be clear")
}
if compressed && !c.ws.compress {
return c.wsHandleProtocolError("compressed frame received without negotiated permessage-deflate")
}
pos++
tmpBuf, pos, err = wsGet(ior, buf, pos, 1)
if err != nil {
return err
}
b1 := tmpBuf[0]
// Clients MUST set the mask bit. If not set, reject.
// However, LEAF by default will not have masking, unless they are forced to, by configuration.
if r.mask && b1&wsMaskBit == 0 {
return c.wsHandleProtocolError("mask bit missing")
}
// Store size in case it is < 125
r.rem = uint64(b1 & 0x7F)
switch frameType {
case wsPingMessage, wsPongMessage, wsCloseMessage:
if r.rem > wsMaxControlPayloadSize {
return c.wsHandleProtocolError(
fmt.Sprintf("control frame length bigger than maximum allowed of %v bytes",
wsMaxControlPayloadSize))
}
if !final {
return c.wsHandleProtocolError("control frame does not have final bit set")
}
if compressed {
return c.wsHandleProtocolError("control frame must not be compressed")
}
case wsTextMessage, wsBinaryMessage:
if !r.ff {
return c.wsHandleProtocolError("new message started before final frame for previous message was received")
}
r.ff = final
r.fc = compressed
case wsContinuationFrame:
// Compressed bit must be only set in the first frame
if r.ff || compressed {
return c.wsHandleProtocolError("invalid continuation frame")
}
r.ff = final
default:
return c.wsHandleProtocolError(fmt.Sprintf("unknown opcode %v", frameType))
}
switch r.rem {
case 126:
tmpBuf, pos, err = wsGet(ior, buf, pos, 2)
if err != nil {
return err
}
r.rem = uint64(binary.BigEndian.Uint16(tmpBuf))
case 127:
tmpBuf, pos, err = wsGet(ior, buf, pos, 8)
if err != nil {
return err
}
if r.rem = binary.BigEndian.Uint64(tmpBuf); r.rem&(uint64(1)<<63) != 0 {
return c.wsHandleProtocolError("invalid 64-bit payload length")
}
}
if r.mask {
// Read masking key
tmpBuf, pos, err = wsGet(ior, buf, pos, 4)
if err != nil {
return err
}
copy(r.mkey[:], tmpBuf)
r.mkpos = 0
}
// Handle control messages in place...
if wsIsControlFrame(frameType) {
pos, err = c.wsHandleControlFrame(r, frameType, ior, buf, pos)
if err != nil {
return err
}
continue
}
// Done with the frame header
r.fs = false
}
if pos < max {
n := r.rem
if pos+n > max {
n = max - pos
}
b := buf[pos : pos+n]
pos += n
r.rem -= n
// If needed, unmask the buffer
if r.mask {
r.unmask(b)
}
if err := handle(b, r.fc, r.ff && r.rem == 0); err != nil {
return err
}
if r.rem == 0 {
r.fs = true
}
}
}
return nil
}
func (r *wsReadInfo) Read(dst []byte) (int, error) {
if len(dst) == 0 {
return 0, nil
}
if len(r.cbufs) == 0 {
return 0, io.EOF
}
copied := 0
rem := len(dst)
for buf := r.cbufs[0]; buf != nil && rem > 0; {
n := len(buf[r.coff:])
if n > rem {
n = rem
}
copy(dst[copied:], buf[r.coff:r.coff+n])
copied += n
rem -= n
r.coff += n
buf = r.nextCBuf()
}
return copied, nil
}
func (r *wsReadInfo) nextCBuf() []byte {
// We still have remaining data in the first buffer
if r.coff != len(r.cbufs[0]) {
return r.cbufs[0]
}
// We read the full first buffer. Reset offset.
r.coff = 0
// We were at the last buffer, so we are done.
if len(r.cbufs) == 1 {
r.cbufs = nil
return nil
}
// Here we move to the next buffer.
r.cbufs = r.cbufs[1:]
return r.cbufs[0]
}
func (r *wsReadInfo) ReadByte() (byte, error) {
for len(r.cbufs) > 0 && len(r.cbufs[0]) == 0 {
r.nextCBuf()
}
if len(r.cbufs) == 0 {
return 0, io.EOF
}
b := r.cbufs[0][r.coff]
r.coff++
r.nextCBuf()
return b, nil
}
func (c *client) wsDecompressAndParse(r *wsReadInfo, b []byte, final bool, mpay int) error {
limit := wsMaxMessageSize(mpay)
if len(b) > 0 {
if r.csz+uint64(len(b)) > limit {
return ErrMaxPayload
}
r.cbufs = append(r.cbufs, append([]byte(nil), b...))
r.csz += uint64(len(b))
}
if !final {
return nil
}
if r.csz+uint64(len(compressLastBlock)) > limit {
return ErrMaxPayload
}
r.cbufs = append(r.cbufs, compressLastBlock)
r.csz += uint64(len(compressLastBlock))
r.coff = 0
d, _ := decompressorPool.Get().(io.ReadCloser)
if d == nil {
d = flate.NewReader(r)
} else {
d.(flate.Resetter).Reset(r, nil)
}
defer func() {
d.Close()
decompressorPool.Put(d)
r.cbufs = nil
r.coff = 0
r.csz = 0
}()
lr := io.LimitedReader{R: d, N: int64(mpay + 1)}
buf := make([]byte, 32*1024)
total := 0
for {
n, err := lr.Read(buf)
if n > 0 {
pn := n
if total+n > mpay {
pn = mpay - total
}
if pn > 0 {
if err := c.parse(buf[:pn]); err != nil {
return err
}
}
total += n
if total > mpay {
return ErrMaxPayload
}
}
if err == nil {
continue
}
if err == io.EOF {
return nil
}
return err
}
}
// Handles the PING, PONG and CLOSE websocket control frames.
//
// Client lock MUST NOT be held on entry.
func (c *client) wsHandleControlFrame(r *wsReadInfo, frameType wsOpCode, nc io.Reader, buf []byte, pos uint64) (uint64, error) {
var payload []byte
var err error
if r.rem > 0 {
payload, pos, err = wsGet(nc, buf, pos, r.rem)
if err != nil {
return pos, err
}
if r.mask {
r.unmask(payload)
}
r.rem = 0
}
switch frameType {
case wsCloseMessage:
status := wsCloseStatusNoStatusReceived
var body string
lp := len(payload)
if lp == 1 {
return pos, c.wsHandleProtocolError("close frame payload cannot be 1 byte")
}
// If there is a payload, the status is represented as a 2-byte
// unsigned integer (in network byte order). Then, there may be an
// optional body.
hasStatus, hasBody := lp >= wsCloseSatusSize, lp > wsCloseSatusSize
if hasStatus {
// Decode the status
status = int(binary.BigEndian.Uint16(payload[:wsCloseSatusSize]))
if !wsIsValidCloseStatus(status) {
return pos, c.wsHandleProtocolError(fmt.Sprintf("invalid close status code %v", status))
}
// Now if there is a body, capture it and make sure this is a valid UTF-8.
if hasBody {
body = string(payload[wsCloseSatusSize:])
if !utf8.ValidString(body) {
// https://tools.ietf.org/html/rfc6455#section-5.5.1
// If body is present, it must be a valid utf8
status = wsCloseStatusInvalidPayloadData
body = "invalid utf8 body in close frame"
}
}
}
// 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)
}
c.wsEnqueueControlMessage(wsCloseMessage, clm)
if len(clm) > 0 {
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
case wsPingMessage:
c.wsEnqueueControlMessage(wsPongMessage, payload)
case wsPongMessage:
// Nothing to do..
}
return pos, nil
}
// Unmask the given slice.
func (r *wsReadInfo) unmask(buf []byte) {
p := int(r.mkpos)
if len(buf) < 16 {
for i := 0; i < len(buf); i++ {
buf[i] ^= r.mkey[p&3]
p++
}
r.mkpos = byte(p & 3)
return
}
var k [8]byte
for i := 0; i < 8; i++ {
k[i] = r.mkey[(p+i)&3]
}
km := binary.BigEndian.Uint64(k[:])
n := (len(buf) / 8) * 8
for i := 0; i < n; i += 8 {
tmp := binary.BigEndian.Uint64(buf[i : i+8])
tmp ^= km
binary.BigEndian.PutUint64(buf[i:], tmp)
}
buf = buf[n:]
for i := 0; i < len(buf); i++ {
buf[i] ^= r.mkey[p&3]
p++
}
r.mkpos = byte(p & 3)
}
// Returns true if the op code corresponds to a control frame.
func wsIsControlFrame(frameType wsOpCode) bool {
return frameType >= wsCloseMessage
}
// Create the frame header.
// Encodes the frame type and optional compression flag, and the size of the payload.
func wsCreateFrameHeader(useMasking, compressed bool, frameType wsOpCode, l int) ([]byte, []byte) {
fh := nbPoolGet(wsMaxFrameHeaderSize)[:wsMaxFrameHeaderSize]
n, key := wsFillFrameHeader(fh, useMasking, wsFirstFrame, wsFinalFrame, compressed, frameType, l)
return fh[:n], key
}
func wsFillFrameHeader(fh []byte, useMasking, first, final, compressed bool, frameType wsOpCode, l int) (int, []byte) {
var n int
var b byte
if first {
b = byte(frameType)
}
if final {
b |= wsFinalBit
}
if compressed {
b |= wsRsv1Bit
}
b1 := byte(0)
if useMasking {
b1 |= wsMaskBit
}
switch {
case l <= 125:
n = 2
fh[0] = b
fh[1] = b1 | byte(l)
case l < 65536:
n = 4
fh[0] = b
fh[1] = b1 | 126
binary.BigEndian.PutUint16(fh[2:], uint16(l))
default:
n = 10
fh[0] = b
fh[1] = b1 | 127
binary.BigEndian.PutUint64(fh[2:], uint64(l))
}
var key []byte
if useMasking {
var keyBuf [4]byte
if _, err := io.ReadFull(crand.Reader, keyBuf[:4]); err != nil {
kv := mrand.Int31()
binary.LittleEndian.PutUint32(keyBuf[:4], uint32(kv))
}
copy(fh[n:], keyBuf[:4])
key = fh[n : n+4]
n += 4
}
return n, key
}
// Invokes wsEnqueueControlMessageLocked under client lock.
//
// Client lock MUST NOT be held on entry
func (c *client) wsEnqueueControlMessage(controlMsg wsOpCode, payload []byte) {
c.mu.Lock()
c.wsEnqueueControlMessageLocked(controlMsg, payload)
c.mu.Unlock()
}
// Mask the buffer with the given key
func wsMaskBuf(key, buf []byte) {
for i := 0; i < len(buf); i++ {
buf[i] ^= key[i&3]
}
}
// Mask the buffers, as if they were contiguous, with the given key
func wsMaskBufs(key []byte, bufs [][]byte) {
pos := 0
for i := 0; i < len(bufs); i++ {
buf := bufs[i]
for j := 0; j < len(buf); j++ {
buf[j] ^= key[pos&3]
pos++
}
}
}
// Enqueues a websocket control message.
// If the control message is a wsCloseMessage, then marks this client
// has having sent the close message (since only one should be sent).
// This will prevent the generic closeConnection() to enqueue one.
//
// Client lock held on entry.
func (c *client) wsEnqueueControlMessageLocked(controlMsg wsOpCode, payload []byte) {
// Control messages are never compressed and their size will be
// less than wsMaxControlPayloadSize, which means the frame header
// will be only 2 or 6 bytes.
useMasking := c.ws.maskwrite
sz := 2
if useMasking {
sz += 4
}
cm := nbPoolGet(sz + len(payload))
cm = cm[:cap(cm)]
n, key := wsFillFrameHeader(cm, useMasking, wsFirstFrame, wsFinalFrame, wsUncompressedFrame, controlMsg, len(payload))
cm = cm[:n]
// Note that payload is optional.
if len(payload) > 0 {
cm = append(cm, payload...)
if useMasking {
wsMaskBuf(key, cm[n:])
}
}
c.out.pb += int64(len(cm))
if controlMsg == wsCloseMessage {
// We can't add the close message to the frames buffers
// now. It will be done on a flushOutbound() when there
// are no more pending buffers to send.
c.ws.closeSent = true
c.ws.closeMsg = cm
} else {
c.ws.frames = append(c.ws.frames, cm)
c.ws.fs += int64(len(cm))
}
c.flushSignal()
}
// Enqueues a websocket close message with a status mapped from the given `reason`.
//
// Client lock held on entry
func (c *client) wsEnqueueCloseMessage(reason ClosedState) {
var status int
switch reason {
case ClientClosed:
status = wsCloseStatusNormalClosure
case AuthenticationTimeout, AuthenticationViolation, SlowConsumerPendingBytes, SlowConsumerWriteDeadline,
MaxAccountConnectionsExceeded, MaxConnectionsExceeded, MaxControlLineExceeded, MaxSubscriptionsExceeded,
MissingAccount, AuthenticationExpired, Revocation:
status = wsCloseStatusPolicyViolation
case TLSHandshakeError:
status = wsCloseStatusTLSHandshake
case ParseError, ProtocolViolation, BadClientProtocolVersion:
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.
status = wsCloseStatusGoingAway
default:
status = wsCloseStatusInternalSrvError
}
body := wsCreateCloseMessage(status, reason.String())
c.wsEnqueueControlMessageLocked(wsCloseMessage, body)
nbPoolPut(body) // wsEnqueueControlMessageLocked has taken a copy.
}
// Create and then enqueue a close message with a protocol error and the
// given message. This is invoked when parsing websocket frames.
//
// Lock MUST NOT be held on entry.
func (c *client) wsHandleProtocolError(message string) error {
buf := wsCreateCloseMessage(wsCloseStatusProtocolError, message)
c.wsEnqueueControlMessage(wsCloseMessage, buf)
nbPoolPut(buf) // wsEnqueueControlMessage has taken a copy.
return errors.New(message)
}
func wsIsValidCloseStatus(code int) bool {
switch code {
case wsCloseStatusNoStatusReceived, 1004, 1006, wsCloseStatusTLSHandshake:
return false
}
if code < 1000 || code >= 5000 {
return false
}
// 1016-2999 are currently reserved.
if code >= 1016 && code <= 2999 {
return false
}
return true
}
// Create a close message with the given `status` and `body`.
// If the `body` is more than the maximum allows control frame payload size,
// it is truncated and "..." is added at the end (as a hint that message
// is not complete).
func wsCreateCloseMessage(status int, body string) []byte {
// Since a control message payload is limited in size, we
// will limit the text and add trailing "..." if truncated.
// The body of a Close Message must be preceded with 2 bytes,
// so take that into account for limiting the body length.
if len(body) > wsMaxControlPayloadSize-2 {
body = body[:wsMaxControlPayloadSize-5]
body += "..."
}
buf := nbPoolGet(2 + len(body))[:2+len(body)]
// We need to have a 2 byte unsigned int that represents the error status code
// https://tools.ietf.org/html/rfc6455#section-5.5.1
binary.BigEndian.PutUint16(buf[:2], uint16(status))
copy(buf[2:], []byte(body))
return buf
}
// Process websocket client handshake. On success, returns the raw net.Conn that
// will be used to create a *client object.
// Invoked from the HTTP server listening on websocket port.
func (s *Server) wsUpgrade(w http.ResponseWriter, r *http.Request) (*wsUpgradeResult, error) {
kind := CLIENT
if r.URL != nil {
ep := r.URL.EscapedPath()
if strings.HasSuffix(ep, leafNodeWSPath) {
kind = LEAF
} else if strings.HasSuffix(ep, mqttWSPath) {
kind = MQTT
}
}
opts := s.getOpts()
// From https://tools.ietf.org/html/rfc6455#section-4.2.1
// Point 1.
if r.Method != "GET" {
return nil, wsReturnHTTPError(w, r, http.StatusMethodNotAllowed, "request method must be GET")
}
// Point 2.
if r.Host == _EMPTY_ {
return nil, wsReturnHTTPError(w, r, http.StatusBadRequest, "'Host' missing in request")
}
// Point 3.
if !wsHeaderContains(r.Header, "Upgrade", "websocket") {
return nil, wsReturnHTTPError(w, r, http.StatusBadRequest, "invalid value for header 'Upgrade'")
}
// Point 4.
if !wsHeaderContains(r.Header, "Connection", "Upgrade") {
return nil, wsReturnHTTPError(w, r, http.StatusBadRequest, "invalid value for header 'Connection'")
}
// Point 5.
key := r.Header.Get("Sec-Websocket-Key")
if key == _EMPTY_ {
return nil, wsReturnHTTPError(w, r, http.StatusBadRequest, "key missing")
}
decoded, err := base64.StdEncoding.DecodeString(key)
if err != nil || len(decoded) != 16 {
return nil, wsReturnHTTPError(w, r, http.StatusBadRequest, "invalid websocket key")
}
// Point 6.
if !wsHeaderContains(r.Header, "Sec-Websocket-Version", "13") {
return nil, wsReturnHTTPError(w, r, http.StatusBadRequest, "invalid version")
}
// Others are optional
// Point 7.
if err := s.websocket.checkOrigin(r); err != nil {
return nil, wsReturnHTTPError(w, r, http.StatusForbidden, fmt.Sprintf("origin not allowed: %v", err))
}
// Point 8.
// We don't have protocols, so ignore.
// Point 9.
// Extensions, only support for compression at the moment
compress := opts.Websocket.Compression
if compress {
// Simply check if permessage-deflate extension is present.
compress, _ = wsPMCExtensionSupport(r.Header, true)
}
// We will do masking if asked (unless we reject for tests)
noMasking := r.Header.Get(wsNoMaskingHeader) == wsNoMaskingValue && !wsTestRejectNoMasking
h, ok := w.(http.Hijacker)
if !ok {
return nil, wsReturnHTTPError(w, r, http.StatusBadRequest, "websocket upgrade not supported")
}
conn, brw, err := h.Hijack()
if err != nil {
if conn != nil {
conn.Close()
}
return nil, wsReturnHTTPError(w, r, http.StatusInternalServerError, err.Error())
}
if brw.Reader.Buffered() > 0 {
conn.Close()
return nil, wsReturnHTTPError(w, r, http.StatusBadRequest, "client sent data before handshake is complete")
}
var buf [1024]byte
p := buf[:0]
// From https://tools.ietf.org/html/rfc6455#section-4.2.2
p = append(p, "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: "...)
p = append(p, wsAcceptKey(key)...)
p = append(p, _CRLF_...)
if compress {
p = append(p, wsPMCFullResponse...)
}
if noMasking {
p = append(p, wsNoMaskingFullResponse...)
}
if kind == MQTT {
p = append(p, wsMQTTSecProto...)
}
if s.websocket.rawHeaders != _EMPTY_ {
p = append(p, s.websocket.rawHeaders...)
}
p = append(p, _CRLF_...)
if _, err = conn.Write(p); err != nil {
conn.Close()
return nil, err
}
// If there was a deadline set for the handshake, clear it now.
if opts.Websocket.HandshakeTimeout > 0 {
conn.SetDeadline(time.Time{})
}
// Server always expect "clients" to send masked payload, unless the option
// "no-masking" has been enabled.
ws := &websocket{compress: compress, maskread: !noMasking}
// Check for X-Forwarded-For header
if cips, ok := r.Header[wsXForwardedForHeader]; ok {
if len(cips) > 0 {
cip := cips[0]
if net.ParseIP(cip) != nil {
ws.clientIP = cip
}
}
}
if kind == CLIENT || kind == MQTT {
// Indicate if this is likely coming from a browser.
if ua := r.Header.Get("User-Agent"); ua != _EMPTY_ && strings.HasPrefix(ua, "Mozilla/") {
ws.browser = true
// Disable fragmentation of compressed frames for Safari browsers.
// Unfortunately, you could be running Chrome on macOS and this
// string will contain "Safari/" (along "Chrome/"). However, what
// I have found is that actual Safari browser also have "Version/".
// So make the combination of the two.
ws.nocompfrag = ws.compress && strings.Contains(ua, "Version/") && strings.Contains(ua, "Safari/")
}
if cookies := r.Cookies(); len(cookies) > 0 {
ows := &opts.Websocket
for _, c := range cookies {
if ows.JWTCookie == c.Name {
ws.cookieJwt = c.Value
} else if ows.UsernameCookie == c.Name {
ws.cookieUsername = c.Value
} else if ows.PasswordCookie == c.Name {
ws.cookiePassword = c.Value
} else if ows.TokenCookie == c.Name {
ws.cookieToken = c.Value
}
}
}
}
return &wsUpgradeResult{conn: conn, ws: ws, kind: kind}, nil
}
// Returns true if the header named `name` contains a token with value `value`.
func wsHeaderContains(header http.Header, name string, value string) bool {
for _, s := range header[name] {
tokens := strings.Split(s, ",")
for _, t := range tokens {
t = strings.Trim(t, " \t")
if strings.EqualFold(t, value) {
return true
}
}
}
return false
}
func wsPMCExtensionSupport(header http.Header, checkPMCOnly bool) (bool, bool) {
for _, extensionList := range header["Sec-Websocket-Extensions"] {
extensions := strings.Split(extensionList, ",")
for _, extension := range extensions {
extension = strings.Trim(extension, " \t")
params := strings.Split(extension, ";")
for i, p := range params {
p = strings.Trim(p, " \t")
if strings.EqualFold(p, wsPMCExtension) {
if checkPMCOnly {
return true, false
}
var snc bool
var cnc bool
for j := i + 1; j < len(params); j++ {
p = params[j]
p = strings.Trim(p, " \t")
if strings.EqualFold(p, wsPMCSrvNoCtx) {
snc = true
} else if strings.EqualFold(p, wsPMCCliNoCtx) {
cnc = true
}
if snc && cnc {
return true, true
}
}
return true, false
}
}
}
}
return false, false
}
// Send an HTTP error with the given `status` to the given http response writer `w`.
// Return an error created based on the `reason` string.
func wsReturnHTTPError(w http.ResponseWriter, r *http.Request, status int, reason string) error {
err := fmt.Errorf("%s - websocket handshake error: %s", r.RemoteAddr, reason)
w.Header().Set("Sec-Websocket-Version", "13")
http.Error(w, http.StatusText(status), status)
return err
}
// If the server is configured to accept any origin, then this function returns
// `nil` without checking if the Origin is present and valid. This is also
// the case if the request does not have the Origin header.
// Otherwise, this will check that the Origin matches the same origin or
// any origin in the allowed list.
func (w *srvWebsocket) checkOrigin(r *http.Request) error {
w.mu.RLock()
checkSame := w.sameOrigin
listEmpty := len(w.allowedOrigins) == 0
w.mu.RUnlock()
if !checkSame && listEmpty {
return nil
}
origin := r.Header.Get("Origin")
if origin == _EMPTY_ {
origin = r.Header.Get("Sec-Websocket-Origin")
}
// If the header is not present, we will accept.
// From https://datatracker.ietf.org/doc/html/rfc6455#section-1.6
// "Naturally, when the WebSocket Protocol is used by a dedicated client
// directly (i.e., not from a web page through a web browser), the origin
// model is not useful, as the client can provide any arbitrary origin string."
if origin == _EMPTY_ {
return nil
}
u, err := url.ParseRequestURI(origin)
if err != nil {
return err
}
oh, op, err := wsGetHostAndPort(u.Scheme == "https", u.Host)
if err != nil {
return err
}
// If checking same origin, compare with the http's request's Host.
if checkSame {
rh, rp, err := wsGetHostAndPort(r.TLS != nil, r.Host)
if err != nil {
return err
}
rs := "http"
if r.TLS != nil {
rs = "https"
}
if oh != rh || op != rp || !strings.EqualFold(u.Scheme, rs) {
return errors.New("not same origin")
}
// I guess it is possible to have cases where one wants to check
// same origin, but also that the origin is in the allowed list.
// So continue with the next check.
}
if !listEmpty {
w.mu.RLock()
origins := w.allowedOrigins[oh]
w.mu.RUnlock()
var allowed bool
for _, ao := range origins {
if u.Scheme == ao.scheme && op == ao.port {
allowed = true
break
}
}
if !allowed {
return errors.New("not in the allowed list")
}
}
return nil
}
func wsGetHostAndPort(tls bool, hostport string) (string, string, error) {
host, port, err := net.SplitHostPort(hostport)
if err != nil {
// If error is missing port, then use defaults based on the scheme
if ae, ok := err.(*net.AddrError); ok && strings.Contains(ae.Err, "missing port") {
err = nil
host = hostport
if tls {
port = "443"
} else {
port = "80"
}
}
}
return strings.ToLower(host), port, err
}
// Concatenate the key sent by the client with the GUID, then computes the SHA1 hash
// and returns it as a based64 encoded string.
func wsAcceptKey(key string) string {
h := sha1.New()
h.Write([]byte(key))
h.Write(wsGUID)
return base64.StdEncoding.EncodeToString(h.Sum(nil))
}
func wsMakeChallengeKey() (string, error) {
p := make([]byte, 16)
if _, err := io.ReadFull(crand.Reader, p); err != nil {
return _EMPTY_, err
}
return base64.StdEncoding.EncodeToString(p), nil
}
// Validate the websocket related options.
func validateWebsocketOptions(o *Options) error {
wo := &o.Websocket
// If no port is defined, we don't care about other options
if wo.Port == 0 {
return nil
}
// Enforce TLS... unless NoTLS is set to true.
if wo.TLSConfig == nil && !wo.NoTLS {
return errors.New("websocket requires TLS configuration")
}
// Make sure that allowed origins, if specified, can be parsed.
for _, ao := range wo.AllowedOrigins {
u, err := url.ParseRequestURI(ao)
if err != nil {
return fmt.Errorf("unable to parse allowed origin: %v", err)
}
if u.Scheme != "http" && u.Scheme != "https" {
return fmt.Errorf("unable to parse allowed origin %q: allowed origins must be absolute URLs with http or https scheme", ao)
}
if u.Host == _EMPTY_ {
return fmt.Errorf("unable to parse allowed origin %q: host is required", ao)
}
if _, _, err := wsGetHostAndPort(u.Scheme == "https", u.Host); err != nil {
return fmt.Errorf("unable to parse allowed origin: %v", err)
}
}
// If there is a NoAuthUser, we need to have Users defined and
// the user to be present.
if wo.NoAuthUser != _EMPTY_ {
if err := validateNoAuthUser(o, wo.NoAuthUser); err != nil {
return err
}
}
// Token/Username not possible if there are users/nkeys
if len(o.Users) > 0 || len(o.Nkeys) > 0 {
if wo.Username != _EMPTY_ {
return fmt.Errorf("websocket authentication username not compatible with presence of users/nkeys")
}
if wo.Token != _EMPTY_ {
return fmt.Errorf("websocket authentication token not compatible with presence of users/nkeys")
}
}
// Using JWT requires Trusted Keys
if wo.JWTCookie != _EMPTY_ {
if len(o.TrustedOperators) == 0 && len(o.TrustedKeys) == 0 {
return fmt.Errorf("trusted operators or trusted keys configuration is required for JWT authentication via cookie %q", wo.JWTCookie)
}
}
if err := validatePinnedCerts(wo.TLSPinnedCerts); err != nil {
return fmt.Errorf("websocket: %v", err)
}
// Check for invalid headers here.
for key := range wo.Headers {
k := strings.ToLower(key)
switch k {
case "host",
"content-length",
"connection",
"upgrade",
"nats-no-masking":
return fmt.Errorf("websocket: invalid header %q not allowed", key)
}
if strings.HasPrefix(k, "sec-websocket-") {
return fmt.Errorf("websocket: invalid header %q, \"Sec-WebSocket-\" prefix not allowed", key)
}
}
return nil
}
// Creates or updates the existing map
func (s *Server) wsSetOriginOptions(o *WebsocketOpts) {
ws := &s.websocket
ws.mu.Lock()
defer ws.mu.Unlock()
// Copy over the option's same origin boolean
ws.sameOrigin = o.SameOrigin
// Reset the map. Will help for config reload if/when we support it.
ws.allowedOrigins = nil
if o.AllowedOrigins == nil {
return
}
for _, ao := range o.AllowedOrigins {
// We have previously checked (during options validation) that the urls
// are parseable, but if we get an error, report and skip.
u, err := url.ParseRequestURI(ao)
if err != nil {
s.Errorf("error parsing allowed origin: %v", err)
continue
}
h, p, _ := wsGetHostAndPort(u.Scheme == "https", u.Host)
if ws.allowedOrigins == nil {
ws.allowedOrigins = make(map[string][]*allowedOrigin, len(o.AllowedOrigins))
}
ws.allowedOrigins[h] = append(ws.allowedOrigins[h], &allowedOrigin{scheme: u.Scheme, port: p})
}
}
// Calculate the raw headers for websocket upgrade response.
func (s *Server) wsSetHeadersOptions(o *WebsocketOpts) {
var sb strings.Builder
for k, v := range o.Headers {
sb.WriteString(k)
sb.WriteString(": ")
sb.WriteString(v)
sb.WriteString(_CRLF_)
}
ws := &s.websocket
ws.mu.Lock()
defer ws.mu.Unlock()
ws.rawHeaders = sb.String()
}
// Given the websocket options, we check if any auth configuration
// has been provided. If so, possibly create users/nkey users and
// store them in s.websocket.users/nkeys.
// Also update a boolean that indicates if auth is required for
// websocket clients.
// Server lock is held on entry.
func (s *Server) wsConfigAuth(opts *WebsocketOpts) {
ws := &s.websocket
// If any of those is specified, we consider that there is an override.
ws.authOverride = opts.Username != _EMPTY_ || opts.Token != _EMPTY_ || opts.NoAuthUser != _EMPTY_
}
func (s *Server) startWebsocketServer() {
if s.isShuttingDown() {
return
}
sopts := s.getOpts()
o := &sopts.Websocket
s.wsSetOriginOptions(o)
s.wsSetHeadersOptions(o)
var hl net.Listener
var proto string
var err error
port := o.Port
if port == -1 {
port = 0
}
hp := net.JoinHostPort(o.Host, strconv.Itoa(port))
// We are enforcing (when validating the options) the use of TLS, but the
// code was originally supporting both modes. The reason for TLS only is
// that we expect users to send JWTs with bearer tokens and we want to
// avoid the possibility of it being "intercepted".
s.mu.Lock()
// Do not check o.NoTLS here. If a TLS configuration is available, use it,
// regardless of NoTLS. If we don't have a TLS config, it means that the
// user has configured NoTLS because otherwise the server would have failed
// to start due to options validation.
if o.TLSConfig != nil {
proto = wsSchemePrefixTLS
config := o.TLSConfig.Clone()
config.GetConfigForClient = s.wsGetTLSConfig
hl, err = tls.Listen("tcp", hp, config)
} else {
proto = wsSchemePrefix
hl, err = net.Listen("tcp", hp)
}
s.websocket.listenerErr = err
if err != nil {
s.mu.Unlock()
s.Fatalf("Unable to listen for websocket connections: %v", err)
return
}
if port == 0 {
o.Port = hl.Addr().(*net.TCPAddr).Port
}
s.Noticef("Listening for websocket clients on %s://%s:%d", proto, o.Host, o.Port)
if proto == wsSchemePrefix {
s.Warnf("Websocket not configured with TLS. DO NOT USE IN PRODUCTION!")
}
// These 3 are immutable and will be accessed without lock by the client
// when generating/sending the INFO protocols.
s.websocket.tls = proto == wsSchemePrefixTLS
s.websocket.host, s.websocket.port = o.Host, o.Port
// This will be updated when/if the cluster changes.
s.websocket.connectURLs, err = s.getConnectURLs(o.Advertise, o.Host, o.Port)
if err != nil {
s.Fatalf("Unable to get websocket connect URLs: %v", err)
hl.Close()
s.mu.Unlock()
return
}
hasLeaf := sopts.LeafNode.Port != 0
mux := http.NewServeMux()
mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
res, err := s.wsUpgrade(w, r)
if err != nil {
s.Errorf(err.Error())
return
}
switch res.kind {
case CLIENT:
s.createWSClient(res.conn, res.ws)
case MQTT:
s.createMQTTClient(res.conn, res.ws)
case LEAF:
if !hasLeaf {
s.Errorf("Not configured to accept leaf node connections")
// Silently close for now. If we want to send an error back, we would
// need to create the leafnode client anyway, so that is handling websocket
// frames, then send the error to the remote.
res.conn.Close()
return
}
s.createLeafNode(res.conn, nil, nil, res.ws)
}
})
hs := &http.Server{
Addr: hp,
Handler: mux,
ReadTimeout: o.HandshakeTimeout,
ErrorLog: log.New(&captureHTTPServerLog{s, "websocket: "}, _EMPTY_, 0),
}
s.websocket.mu.Lock()
s.websocket.server = hs
s.websocket.listener = hl
s.websocket.mu.Unlock()
go func() {
if err := hs.Serve(hl); err != http.ErrServerClosed {
s.Fatalf("websocket listener error: %v", err)
}
if s.isLameDuckMode() {
// Signal that we are not accepting new clients
s.ldmCh <- true
// Now wait for the Shutdown...
<-s.quitCh
return
}
s.done <- true
}()
s.mu.Unlock()
}
// The TLS configuration is passed to the listener when the websocket
// "server" is setup. That prevents TLS configuration updates on reload
// from being used. By setting this function in tls.Config.GetConfigForClient
// we instruct the TLS handshake to ask for the tls configuration to be
// used for a specific client. We don't care which client, we always use
// the same TLS configuration.
func (s *Server) wsGetTLSConfig(_ *tls.ClientHelloInfo) (*tls.Config, error) {
opts := s.getOpts()
return opts.Websocket.TLSConfig, nil
}
// This is similar to createClient() but has some modifications
// specific to handle websocket clients.
// The comments have been kept to minimum to reduce code size.
// Check createClient() for more details.
func (s *Server) createWSClient(conn net.Conn, ws *websocket) *client {
opts := s.getOpts()
maxPay := int32(opts.MaxPayload)
maxSubs := int32(opts.MaxSubs)
if maxSubs == 0 {
maxSubs = -1
}
now := time.Now().UTC()
c := &client{srv: s, nc: conn, opts: defaultOpts, mpay: maxPay, msubs: maxSubs, start: now, last: now, ws: ws}
c.registerWithAccount(s.globalAccount())
var info Info
var authRequired bool
s.mu.Lock()
info = s.copyInfo()
// Check auth, override if applicable.
if !info.AuthRequired {
// Set info.AuthRequired since this is what is sent to the client.
info.AuthRequired = s.websocket.authOverride
}
if s.nonceRequired() {
var raw [nonceLen]byte
nonce := raw[:]
s.generateNonce(nonce)
info.Nonce = string(nonce)
}
c.nonce = []byte(info.Nonce)
authRequired = info.AuthRequired
s.totalClients++
s.mu.Unlock()
c.mu.Lock()
if authRequired {
c.flags.set(expectConnect)
}
c.initClient()
c.Debugf("Client connection created")
c.sendProtoNow(c.generateClientInfoJSON(info))
c.mu.Unlock()
s.mu.Lock()
if !s.isRunning() || s.ldm {
if s.isShuttingDown() {
conn.Close()
}
s.mu.Unlock()
return c
}
if opts.MaxConn > 0 && len(s.clients) >= opts.MaxConn {
s.mu.Unlock()
c.maxConnExceeded()
return nil
}
s.clients[c.cid] = c
s.mu.Unlock()
c.mu.Lock()
// Websocket clients do TLS in the websocket http server.
// So no TLS initiation here...
if _, ok := conn.(*tls.Conn); ok {
c.flags.set(handshakeComplete)
}
if c.isClosed() {
c.mu.Unlock()
c.closeConnection(WriteError)
return nil
}
if authRequired {
timeout := opts.AuthTimeout
// Possibly override with Websocket specific value.
if opts.Websocket.AuthTimeout != 0 {
timeout = opts.Websocket.AuthTimeout
}
c.setAuthTimer(secondsToDuration(timeout))
}
c.setPingTimer()
s.startGoRoutine(func() { c.readLoop(nil) })
s.startGoRoutine(func() { c.writeLoop() })
c.mu.Unlock()
return c
}
func (c *client) wsCollapsePtoNB() (net.Buffers, int64) {
nb := c.out.nb
var mfs int
var usz int
if c.ws.browser {
mfs = wsFrameSizeForBrowsers
}
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).
bufs := c.ws.frames
compress := c.ws.compress
if compress && len(nb) > 0 {
// First, make sure we don't compress for very small cumulative buffers.
for _, b := range nb {
usz += len(b)
}
if usz <= wsCompressThreshold {
compress = false
if cp := c.ws.compressor; cp != nil {
cp.Reset(nil)
}
}
}
if compress && len(nb) > 0 {
// Overwrite mfs if this connection does not support fragmented compressed frames.
if mfs > 0 && c.ws.nocompfrag {
mfs = 0
}
buf := bytes.NewBuffer(nbPoolGet(usz))
cp := c.ws.compressor
if cp == nil {
c.ws.compressor, _ = flate.NewWriter(buf, flate.BestSpeed)
cp = c.ws.compressor
} else {
cp.Reset(buf)
}
var csz int
for _, b := range nb {
for len(b) > 0 {
n, err := cp.Write(b)
if err != nil {
// Whatever this error is, it'll be handled by the cp.Flush()
// call below, as the same error will be returned there.
// Let the outer loop return all the buffers back to the pool
// and fall through naturally.
break
}
b = b[n:]
}
nbPoolPut(b) // No longer needed as contents written to compressor.
}
if err := cp.Flush(); err != nil {
c.Errorf("Error during compression: %v", err)
c.markConnAsClosed(WriteError)
cp.Reset(nil)
return nil, 0
}
b := buf.Bytes()
p := b[:len(b)-4]
if mfs > 0 && len(p) > mfs {
for first, final := true, false; len(p) > 0; first = false {
lp := len(p)
if lp > mfs {
lp = mfs
} else {
final = true
}
// Only the first frame should be marked as compressed, so pass
// `first` for the compressed boolean.
fh := nbPoolGet(wsMaxFrameHeaderSize)[:wsMaxFrameHeaderSize]
n, key := wsFillFrameHeader(fh, mask, first, final, first, wsBinaryMessage, lp)
if mask {
wsMaskBuf(key, p[:lp])
}
bufs = append(bufs, fh[:n], p[:lp])
csz += n + lp
p = p[lp:]
}
} else {
ol := len(p)
h, key := wsCreateFrameHeader(mask, true, wsBinaryMessage, ol)
if mask {
wsMaskBuf(key, p)
}
if ol > 0 {
bufs = append(bufs, h, p)
}
csz = len(h) + ol
}
// Make sure that the compressor no longer holds a reference to
// the bytes.Buffer, so that the underlying memory gets cleaned
// up after flushOutbound/flushAndClose. For this to be safe, we
// always cp.Reset(...) before reusing the compressor again.
cp.Reset(nil)
// Add to pb the compressed data size (including headers), but
// remove the original uncompressed data size that was added
// during the queueing.
c.out.pb += int64(csz) - int64(usz)
c.ws.fs += int64(csz)
} else if len(nb) > 0 {
var total int
if mfs > 0 {
// We are limiting the frame size.
startFrame := func() int {
bufs = append(bufs, nbPoolGet(wsMaxFrameHeaderSize))
return len(bufs) - 1
}
endFrame := func(idx, size int) {
bufs[idx] = bufs[idx][:wsMaxFrameHeaderSize]
n, key := wsFillFrameHeader(bufs[idx], mask, wsFirstFrame, wsFinalFrame, wsUncompressedFrame, wsBinaryMessage, size)
bufs[idx] = bufs[idx][:n]
c.out.pb += int64(n)
c.ws.fs += int64(n + size)
if mask {
wsMaskBufs(key, bufs[idx+1:])
}
}
fhIdx := startFrame()
for i := 0; i < len(nb); i++ {
b := nb[i]
if total+len(b) <= mfs {
buf := nbPoolGet(len(b))
bufs = append(bufs, append(buf, b...))
total += len(b)
nbPoolPut(nb[i])
continue
}
for len(b) > 0 {
endStart := total != 0
if endStart {
endFrame(fhIdx, total)
}
total = len(b)
if total >= mfs {
total = mfs
}
if endStart {
fhIdx = startFrame()
}
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)
}
} else {
// If there is no limit on the frame size, create a single frame for
// all pending buffers.
for _, b := range nb {
total += len(b)
}
wsfh, key := wsCreateFrameHeader(mask, false, wsBinaryMessage, total)
c.out.pb += int64(len(wsfh))
bufs = append(bufs, wsfh)
idx := len(bufs)
bufs = append(bufs, nb...)
if mask {
wsMaskBufs(key, bufs[idx:])
}
c.ws.fs += int64(len(wsfh) + total)
}
}
if len(c.ws.closeMsg) > 0 {
bufs = append(bufs, c.ws.closeMsg)
c.ws.fs += int64(len(c.ws.closeMsg))
c.ws.closeMsg = nil
c.ws.compressor = nil
}
c.ws.frames = nil
return bufs, c.ws.fs
}
func isWSURL(u *url.URL) bool {
return strings.HasPrefix(strings.ToLower(u.Scheme), wsSchemePrefix)
}
func isWSSURL(u *url.URL) bool {
return strings.HasPrefix(strings.ToLower(u.Scheme), wsSchemePrefixTLS)
}