github.com/mdaxf/iac@v0.0.0-20240519030858-58a061660378/vendor_skip/nhooyr.io/websocket/compress_notjs.go (about)

     1  // +build !js
     2  
     3  package websocket
     4  
     5  import (
     6  	"io"
     7  	"net/http"
     8  	"sync"
     9  
    10  	"github.com/klauspost/compress/flate"
    11  )
    12  
    13  func (m CompressionMode) opts() *compressionOptions {
    14  	return &compressionOptions{
    15  		clientNoContextTakeover: m == CompressionNoContextTakeover,
    16  		serverNoContextTakeover: m == CompressionNoContextTakeover,
    17  	}
    18  }
    19  
    20  type compressionOptions struct {
    21  	clientNoContextTakeover bool
    22  	serverNoContextTakeover bool
    23  }
    24  
    25  func (copts *compressionOptions) setHeader(h http.Header) {
    26  	s := "permessage-deflate"
    27  	if copts.clientNoContextTakeover {
    28  		s += "; client_no_context_takeover"
    29  	}
    30  	if copts.serverNoContextTakeover {
    31  		s += "; server_no_context_takeover"
    32  	}
    33  	h.Set("Sec-WebSocket-Extensions", s)
    34  }
    35  
    36  // These bytes are required to get flate.Reader to return.
    37  // They are removed when sending to avoid the overhead as
    38  // WebSocket framing tell's when the message has ended but then
    39  // we need to add them back otherwise flate.Reader keeps
    40  // trying to return more bytes.
    41  const deflateMessageTail = "\x00\x00\xff\xff"
    42  
    43  type trimLastFourBytesWriter struct {
    44  	w    io.Writer
    45  	tail []byte
    46  }
    47  
    48  func (tw *trimLastFourBytesWriter) reset() {
    49  	if tw != nil && tw.tail != nil {
    50  		tw.tail = tw.tail[:0]
    51  	}
    52  }
    53  
    54  func (tw *trimLastFourBytesWriter) Write(p []byte) (int, error) {
    55  	if tw.tail == nil {
    56  		tw.tail = make([]byte, 0, 4)
    57  	}
    58  
    59  	extra := len(tw.tail) + len(p) - 4
    60  
    61  	if extra <= 0 {
    62  		tw.tail = append(tw.tail, p...)
    63  		return len(p), nil
    64  	}
    65  
    66  	// Now we need to write as many extra bytes as we can from the previous tail.
    67  	if extra > len(tw.tail) {
    68  		extra = len(tw.tail)
    69  	}
    70  	if extra > 0 {
    71  		_, err := tw.w.Write(tw.tail[:extra])
    72  		if err != nil {
    73  			return 0, err
    74  		}
    75  
    76  		// Shift remaining bytes in tail over.
    77  		n := copy(tw.tail, tw.tail[extra:])
    78  		tw.tail = tw.tail[:n]
    79  	}
    80  
    81  	// If p is less than or equal to 4 bytes,
    82  	// all of it is is part of the tail.
    83  	if len(p) <= 4 {
    84  		tw.tail = append(tw.tail, p...)
    85  		return len(p), nil
    86  	}
    87  
    88  	// Otherwise, only the last 4 bytes are.
    89  	tw.tail = append(tw.tail, p[len(p)-4:]...)
    90  
    91  	p = p[:len(p)-4]
    92  	n, err := tw.w.Write(p)
    93  	return n + 4, err
    94  }
    95  
    96  var flateReaderPool sync.Pool
    97  
    98  func getFlateReader(r io.Reader, dict []byte) io.Reader {
    99  	fr, ok := flateReaderPool.Get().(io.Reader)
   100  	if !ok {
   101  		return flate.NewReaderDict(r, dict)
   102  	}
   103  	fr.(flate.Resetter).Reset(r, dict)
   104  	return fr
   105  }
   106  
   107  func putFlateReader(fr io.Reader) {
   108  	flateReaderPool.Put(fr)
   109  }
   110  
   111  type slidingWindow struct {
   112  	buf []byte
   113  }
   114  
   115  var swPoolMu sync.RWMutex
   116  var swPool = map[int]*sync.Pool{}
   117  
   118  func slidingWindowPool(n int) *sync.Pool {
   119  	swPoolMu.RLock()
   120  	p, ok := swPool[n]
   121  	swPoolMu.RUnlock()
   122  	if ok {
   123  		return p
   124  	}
   125  
   126  	p = &sync.Pool{}
   127  
   128  	swPoolMu.Lock()
   129  	swPool[n] = p
   130  	swPoolMu.Unlock()
   131  
   132  	return p
   133  }
   134  
   135  func (sw *slidingWindow) init(n int) {
   136  	if sw.buf != nil {
   137  		return
   138  	}
   139  
   140  	if n == 0 {
   141  		n = 32768
   142  	}
   143  
   144  	p := slidingWindowPool(n)
   145  	buf, ok := p.Get().([]byte)
   146  	if ok {
   147  		sw.buf = buf[:0]
   148  	} else {
   149  		sw.buf = make([]byte, 0, n)
   150  	}
   151  }
   152  
   153  func (sw *slidingWindow) close() {
   154  	if sw.buf == nil {
   155  		return
   156  	}
   157  
   158  	swPoolMu.Lock()
   159  	swPool[cap(sw.buf)].Put(sw.buf)
   160  	swPoolMu.Unlock()
   161  	sw.buf = nil
   162  }
   163  
   164  func (sw *slidingWindow) write(p []byte) {
   165  	if len(p) >= cap(sw.buf) {
   166  		sw.buf = sw.buf[:cap(sw.buf)]
   167  		p = p[len(p)-cap(sw.buf):]
   168  		copy(sw.buf, p)
   169  		return
   170  	}
   171  
   172  	left := cap(sw.buf) - len(sw.buf)
   173  	if left < len(p) {
   174  		// We need to shift spaceNeeded bytes from the end to make room for p at the end.
   175  		spaceNeeded := len(p) - left
   176  		copy(sw.buf, sw.buf[spaceNeeded:])
   177  		sw.buf = sw.buf[:len(sw.buf)-spaceNeeded]
   178  	}
   179  
   180  	sw.buf = append(sw.buf, p...)
   181  }