1// Copyright 2017 The Gorilla WebSocket Authors. All rights reserved.
  2// Use of this source code is governed by a BSD-style
  3// license that can be found in the LICENSE file.
  4
  5package websocket
  6
  7import (
  8	"bytes"
  9	"net"
 10	"sync"
 11	"time"
 12)
 13
 14// PreparedMessage caches on the wire representations of a message payload.
 15// Use PreparedMessage to efficiently send a message payload to multiple
 16// connections. PreparedMessage is especially useful when compression is used
 17// because the CPU and memory expensive compression operation can be executed
 18// once for a given set of compression options.
 19type PreparedMessage struct {
 20	messageType int
 21	data        []byte
 22	mu          sync.Mutex
 23	frames      map[prepareKey]*preparedFrame
 24}
 25
 26// prepareKey defines a unique set of options to cache prepared frames in PreparedMessage.
 27type prepareKey struct {
 28	isServer         bool
 29	compress         bool
 30	compressionLevel int
 31}
 32
 33// preparedFrame contains data in wire representation.
 34type preparedFrame struct {
 35	once sync.Once
 36	data []byte
 37}
 38
 39// NewPreparedMessage returns an initialized PreparedMessage. You can then send
 40// it to connection using WritePreparedMessage method. Valid wire
 41// representation will be calculated lazily only once for a set of current
 42// connection options.
 43func NewPreparedMessage(messageType int, data []byte) (*PreparedMessage, error) {
 44	pm := &PreparedMessage{
 45		messageType: messageType,
 46		frames:      make(map[prepareKey]*preparedFrame),
 47		data:        data,
 48	}
 49
 50	// Prepare a plain server frame.
 51	_, frameData, err := pm.frame(prepareKey{isServer: true, compress: false})
 52	if err != nil {
 53		return nil, err
 54	}
 55
 56	// To protect against caller modifying the data argument, remember the data
 57	// copied to the plain server frame.
 58	pm.data = frameData[len(frameData)-len(data):]
 59	return pm, nil
 60}
 61
 62func (pm *PreparedMessage) frame(key prepareKey) (int, []byte, error) {
 63	pm.mu.Lock()
 64	frame, ok := pm.frames[key]
 65	if !ok {
 66		frame = &preparedFrame{}
 67		pm.frames[key] = frame
 68	}
 69	pm.mu.Unlock()
 70
 71	var err error
 72	frame.once.Do(func() {
 73		// Prepare a frame using a 'fake' connection.
 74		// TODO: Refactor code in conn.go to allow more direct construction of
 75		// the frame.
 76		mu := make(chan struct{}, 1)
 77		mu <- struct{}{}
 78		var nc prepareConn
 79		c := &Conn{
 80			conn:                   &nc,
 81			mu:                     mu,
 82			isServer:               key.isServer,
 83			compressionLevel:       key.compressionLevel,
 84			enableWriteCompression: true,
 85			writeBuf:               make([]byte, defaultWriteBufferSize+maxFrameHeaderSize),
 86		}
 87		if key.compress {
 88			c.newCompressionWriter = compressNoContextTakeover
 89		}
 90		err = c.WriteMessage(pm.messageType, pm.data)
 91		frame.data = nc.buf.Bytes()
 92	})
 93	return pm.messageType, frame.data, err
 94}
 95
 96type prepareConn struct {
 97	buf bytes.Buffer
 98	net.Conn
 99}
100
101func (pc *prepareConn) Write(p []byte) (int, error)        { return pc.buf.Write(p) }
102func (pc *prepareConn) SetWriteDeadline(t time.Time) error { return nil }