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 }
103