gtsocial-umbx

Unnamed repository; edit this file 'description' to name the repository.
Log | Files | Refs | README | LICENSE

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 }