prepared.go (2969B)
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 }