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 5 package websocket
6 7 import (
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.
19 type 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.
27 type prepareKey struct {
28 isServer bool
29 compress bool
30 compressionLevel int
31 }
32 33 // preparedFrame contains data in wire representation.
34 type 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.
43 func 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 62 func (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 96 type prepareConn struct {
97 buf bytes.Buffer
98 net.Conn
99 }
100 101 func (pc *prepareConn) Write(p []byte) (int, error) { return pc.buf.Write(p) }
102 func (pc *prepareConn) SetWriteDeadline(t time.Time) error { return nil }
103