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	err         error
 23	mu          sync.Mutex
 24	frames      map[prepareKey]*preparedFrame
 25}
 26
 27// prepareKey defines a unique set of options to cache prepared frames in PreparedMessage.
 28type prepareKey struct {
 29	isServer         bool
 30	compress         bool
 31	compressionLevel int
 32}
 33
 34// preparedFrame contains data in wire representation.
 35type preparedFrame struct {
 36	once sync.Once
 37	data []byte
 38}
 39
 40// NewPreparedMessage returns an initialized PreparedMessage. You can then send
 41// it to connection using WritePreparedMessage method. Valid wire
 42// representation will be calculated lazily only once for a set of current
 43// connection options.
 44func NewPreparedMessage(messageType int, data []byte) (*PreparedMessage, error) {
 45	pm := &PreparedMessage{
 46		messageType: messageType,
 47		frames:      make(map[prepareKey]*preparedFrame),
 48		data:        data,
 49	}
 50
 51	// Prepare a plain server frame.
 52	_, frameData, err := pm.frame(prepareKey{isServer: true, compress: false})
 53	if err != nil {
 54		return nil, err
 55	}
 56
 57	// To protect against caller modifying the data argument, remember the data
 58	// copied to the plain server frame.
 59	pm.data = frameData[len(frameData)-len(data):]
 60	return pm, nil
 61}
 62
 63func (pm *PreparedMessage) frame(key prepareKey) (int, []byte, error) {
 64	pm.mu.Lock()
 65	frame, ok := pm.frames[key]
 66	if !ok {
 67		frame = &preparedFrame{}
 68		pm.frames[key] = frame
 69	}
 70	pm.mu.Unlock()
 71
 72	var err error
 73	frame.once.Do(func() {
 74		// Prepare a frame using a 'fake' connection.
 75		// TODO: Refactor code in conn.go to allow more direct construction of
 76		// the frame.
 77		mu := make(chan bool, 1)
 78		mu <- true
 79		var nc prepareConn
 80		c := &Conn{
 81			conn:                   &nc,
 82			mu:                     mu,
 83			isServer:               key.isServer,
 84			compressionLevel:       key.compressionLevel,
 85			enableWriteCompression: true,
 86			writeBuf:               make([]byte, defaultWriteBufferSize+maxFrameHeaderSize),
 87		}
 88		if key.compress {
 89			c.newCompressionWriter = compressNoContextTakeover
 90		}
 91		err = c.WriteMessage(pm.messageType, pm.data)
 92		frame.data = nc.buf.Bytes()
 93	})
 94	return pm.messageType, frame.data, err
 95}
 96
 97type prepareConn struct {
 98	buf bytes.Buffer
 99	net.Conn
100}
101
102func (pc *prepareConn) Write(p []byte) (int, error)        { return pc.buf.Write(p) }
103func (pc *prepareConn) SetWriteDeadline(t time.Time) error { return nil }