github.com/ezoic/ws@v1.0.4-0.20220713205711-5c1d69e074c5/write.go (about)

     1  package ws
     2  
     3  import (
     4  	"encoding/binary"
     5  	"io"
     6  )
     7  
     8  // Header size length bounds in bytes.
     9  const (
    10  	MaxHeaderSize = 14
    11  	MinHeaderSize = 2
    12  )
    13  
    14  const (
    15  	bit0 = 0x80
    16  	bit1 = 0x40
    17  	bit2 = 0x20
    18  	bit3 = 0x10
    19  	bit4 = 0x08
    20  	bit5 = 0x04
    21  	bit6 = 0x02
    22  	bit7 = 0x01
    23  
    24  	len7  = int64(125)
    25  	len16 = int64(^(uint16(0)))
    26  	len64 = int64(^(uint64(0)) >> 1)
    27  )
    28  
    29  // HeaderSize returns number of bytes that are needed to encode given header.
    30  // It returns -1 if header is malformed.
    31  func HeaderSize(h Header) (n int) {
    32  	switch {
    33  	case h.Length < 126:
    34  		n = 2
    35  	case h.Length <= len16:
    36  		n = 4
    37  	case h.Length <= len64:
    38  		n = 10
    39  	default:
    40  		return -1
    41  	}
    42  	if h.Masked {
    43  		n += len(h.Mask)
    44  	}
    45  	return n
    46  }
    47  
    48  // WriteHeader writes header binary representation into w.
    49  func WriteHeader(w io.Writer, h Header) error {
    50  	// Make slice of bytes with capacity 14 that could hold any header.
    51  	bts := make([]byte, MaxHeaderSize)
    52  
    53  	if h.Fin {
    54  		bts[0] |= bit0
    55  	}
    56  	bts[0] |= h.Rsv << 4
    57  	bts[0] |= byte(h.OpCode)
    58  
    59  	var n int
    60  	switch {
    61  	case h.Length <= len7:
    62  		bts[1] = byte(h.Length)
    63  		n = 2
    64  
    65  	case h.Length <= len16:
    66  		bts[1] = 126
    67  		binary.BigEndian.PutUint16(bts[2:4], uint16(h.Length))
    68  		n = 4
    69  
    70  	case h.Length <= len64:
    71  		bts[1] = 127
    72  		binary.BigEndian.PutUint64(bts[2:10], uint64(h.Length))
    73  		n = 10
    74  
    75  	default:
    76  		return ErrHeaderLengthUnexpected
    77  	}
    78  
    79  	if h.Masked {
    80  		bts[1] |= bit0
    81  		n += copy(bts[n:], h.Mask[:])
    82  	}
    83  
    84  	_, err := w.Write(bts[:n])
    85  
    86  	return err
    87  }
    88  
    89  // WriteFrame writes frame binary representation into w.
    90  func WriteFrame(w io.Writer, f Frame) error {
    91  	err := WriteHeader(w, f.Header)
    92  	if err != nil {
    93  		return err
    94  	}
    95  	_, err = w.Write(f.Payload)
    96  	return err
    97  }
    98  
    99  // MustWriteFrame is like WriteFrame but panics if frame can not be read.
   100  func MustWriteFrame(w io.Writer, f Frame) {
   101  	if err := WriteFrame(w, f); err != nil {
   102  		panic(err)
   103  	}
   104  }