blob: c854225e9676fa71c7e1d7e33b1c23a6f138b5af [file] [log] [blame]
khenaidoo59ce9dd2019-11-11 13:05:32 -05001// 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
khenaidoo59ce9dd2019-11-11 13:05:32 -050022 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.
khenaidoo26721882021-08-11 17:42:52 -040076 mu := make(chan struct{}, 1)
77 mu <- struct{}{}
khenaidoo59ce9dd2019-11-11 13:05:32 -050078 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 }