prepared.go

Documentation: github.com/gorilla/websocket

     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 }
   103
View as plain text