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