prepared.go 2.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102
  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. package websocket
  5. import (
  6. "bytes"
  7. "net"
  8. "sync"
  9. "time"
  10. )
  11. // PreparedMessage caches on the wire representations of a message payload.
  12. // Use PreparedMessage to efficiently send a message payload to multiple
  13. // connections. PreparedMessage is especially useful when compression is used
  14. // because the CPU and memory expensive compression operation can be executed
  15. // once for a given set of compression options.
  16. type PreparedMessage struct {
  17. messageType int
  18. data []byte
  19. mu sync.Mutex
  20. frames map[prepareKey]*preparedFrame
  21. }
  22. // prepareKey defines a unique set of options to cache prepared frames in PreparedMessage.
  23. type prepareKey struct {
  24. isServer bool
  25. compress bool
  26. compressionLevel int
  27. }
  28. // preparedFrame contains data in wire representation.
  29. type preparedFrame struct {
  30. once sync.Once
  31. data []byte
  32. }
  33. // NewPreparedMessage returns an initialized PreparedMessage. You can then send
  34. // it to connection using WritePreparedMessage method. Valid wire
  35. // representation will be calculated lazily only once for a set of current
  36. // connection options.
  37. func NewPreparedMessage(messageType int, data []byte) (*PreparedMessage, error) {
  38. pm := &PreparedMessage{
  39. messageType: messageType,
  40. frames: make(map[prepareKey]*preparedFrame),
  41. data: data,
  42. }
  43. // Prepare a plain server frame.
  44. _, frameData, err := pm.frame(prepareKey{isServer: true, compress: false})
  45. if err != nil {
  46. return nil, err
  47. }
  48. // To protect against caller modifying the data argument, remember the data
  49. // copied to the plain server frame.
  50. pm.data = frameData[len(frameData)-len(data):]
  51. return pm, nil
  52. }
  53. func (pm *PreparedMessage) frame(key prepareKey) (int, []byte, error) {
  54. pm.mu.Lock()
  55. frame, ok := pm.frames[key]
  56. if !ok {
  57. frame = &preparedFrame{}
  58. pm.frames[key] = frame
  59. }
  60. pm.mu.Unlock()
  61. var err error
  62. frame.once.Do(func() {
  63. // Prepare a frame using a 'fake' connection.
  64. // TODO: Refactor code in conn.go to allow more direct construction of
  65. // the frame.
  66. mu := make(chan struct{}, 1)
  67. mu <- struct{}{}
  68. var nc prepareConn
  69. c := &Conn{
  70. conn: &nc,
  71. mu: mu,
  72. isServer: key.isServer,
  73. compressionLevel: key.compressionLevel,
  74. enableWriteCompression: true,
  75. writeBuf: make([]byte, defaultWriteBufferSize+maxFrameHeaderSize),
  76. }
  77. if key.compress {
  78. c.newCompressionWriter = compressNoContextTakeover
  79. }
  80. err = c.WriteMessage(pm.messageType, pm.data)
  81. frame.data = nc.buf.Bytes()
  82. })
  83. return pm.messageType, frame.data, err
  84. }
  85. type prepareConn struct {
  86. buf bytes.Buffer
  87. net.Conn
  88. }
  89. func (pc *prepareConn) Write(p []byte) (int, error) { return pc.buf.Write(p) }
  90. func (pc *prepareConn) SetWriteDeadline(t time.Time) error { return nil }