vendor/github.com/gorilla/websocket/prepared.go

   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 bool, 1)
  77                 mu <- true
  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 }