frame.go raw

   1  //go:build !js
   2  
   3  package websocket
   4  
   5  import (
   6  	"bufio"
   7  	"encoding/binary"
   8  	"fmt"
   9  	"io"
  10  	"math"
  11  
  12  	"github.com/coder/websocket/internal/errd"
  13  )
  14  
  15  // opcode represents a WebSocket opcode.
  16  type opcode int
  17  
  18  // https://tools.ietf.org/html/rfc6455#section-11.8.
  19  const (
  20  	opContinuation opcode = iota
  21  	opText
  22  	opBinary
  23  	// 3 - 7 are reserved for further non-control frames.
  24  	_
  25  	_
  26  	_
  27  	_
  28  	_
  29  	opClose
  30  	opPing
  31  	opPong
  32  	// 11-16 are reserved for further control frames.
  33  )
  34  
  35  // header represents a WebSocket frame header.
  36  // See https://tools.ietf.org/html/rfc6455#section-5.2.
  37  type header struct {
  38  	fin    bool
  39  	rsv1   bool
  40  	rsv2   bool
  41  	rsv3   bool
  42  	opcode opcode
  43  
  44  	payloadLength int64
  45  
  46  	masked  bool
  47  	maskKey uint32
  48  }
  49  
  50  // readFrameHeader reads a header from the reader.
  51  // See https://tools.ietf.org/html/rfc6455#section-5.2.
  52  func readFrameHeader(r *bufio.Reader, readBuf []byte) (h header, err error) {
  53  	defer errd.Wrap(&err, "failed to read frame header")
  54  
  55  	b, err := r.ReadByte()
  56  	if err != nil {
  57  		return header{}, err
  58  	}
  59  
  60  	h.fin = b&(1<<7) != 0
  61  	h.rsv1 = b&(1<<6) != 0
  62  	h.rsv2 = b&(1<<5) != 0
  63  	h.rsv3 = b&(1<<4) != 0
  64  
  65  	h.opcode = opcode(b & 0xf)
  66  
  67  	b, err = r.ReadByte()
  68  	if err != nil {
  69  		return header{}, err
  70  	}
  71  
  72  	h.masked = b&(1<<7) != 0
  73  
  74  	payloadLength := b &^ (1 << 7)
  75  	switch {
  76  	case payloadLength < 126:
  77  		h.payloadLength = int64(payloadLength)
  78  	case payloadLength == 126:
  79  		_, err = io.ReadFull(r, readBuf[:2])
  80  		h.payloadLength = int64(binary.BigEndian.Uint16(readBuf))
  81  	case payloadLength == 127:
  82  		_, err = io.ReadFull(r, readBuf)
  83  		h.payloadLength = int64(binary.BigEndian.Uint64(readBuf))
  84  	}
  85  	if err != nil {
  86  		return header{}, err
  87  	}
  88  
  89  	if h.payloadLength < 0 {
  90  		return header{}, fmt.Errorf("received negative payload length: %v", h.payloadLength)
  91  	}
  92  
  93  	if h.masked {
  94  		_, err = io.ReadFull(r, readBuf[:4])
  95  		if err != nil {
  96  			return header{}, err
  97  		}
  98  		h.maskKey = binary.LittleEndian.Uint32(readBuf)
  99  	}
 100  
 101  	return h, nil
 102  }
 103  
 104  // maxControlPayload is the maximum length of a control frame payload.
 105  // See https://tools.ietf.org/html/rfc6455#section-5.5.
 106  const maxControlPayload = 125
 107  
 108  // writeFrameHeader writes the bytes of the header to w.
 109  // See https://tools.ietf.org/html/rfc6455#section-5.2
 110  func writeFrameHeader(h header, w *bufio.Writer, buf []byte) (err error) {
 111  	defer errd.Wrap(&err, "failed to write frame header")
 112  
 113  	var b byte
 114  	if h.fin {
 115  		b |= 1 << 7
 116  	}
 117  	if h.rsv1 {
 118  		b |= 1 << 6
 119  	}
 120  	if h.rsv2 {
 121  		b |= 1 << 5
 122  	}
 123  	if h.rsv3 {
 124  		b |= 1 << 4
 125  	}
 126  
 127  	b |= byte(h.opcode)
 128  
 129  	err = w.WriteByte(b)
 130  	if err != nil {
 131  		return err
 132  	}
 133  
 134  	lengthByte := byte(0)
 135  	if h.masked {
 136  		lengthByte |= 1 << 7
 137  	}
 138  
 139  	switch {
 140  	case h.payloadLength > math.MaxUint16:
 141  		lengthByte |= 127
 142  	case h.payloadLength > 125:
 143  		lengthByte |= 126
 144  	case h.payloadLength >= 0:
 145  		lengthByte |= byte(h.payloadLength)
 146  	}
 147  	err = w.WriteByte(lengthByte)
 148  	if err != nil {
 149  		return err
 150  	}
 151  
 152  	switch {
 153  	case h.payloadLength > math.MaxUint16:
 154  		binary.BigEndian.PutUint64(buf, uint64(h.payloadLength))
 155  		_, err = w.Write(buf)
 156  	case h.payloadLength > 125:
 157  		binary.BigEndian.PutUint16(buf, uint16(h.payloadLength))
 158  		_, err = w.Write(buf[:2])
 159  	}
 160  	if err != nil {
 161  		return err
 162  	}
 163  
 164  	if h.masked {
 165  		binary.LittleEndian.PutUint32(buf, h.maskKey)
 166  		_, err = w.Write(buf[:4])
 167  		if err != nil {
 168  			return err
 169  		}
 170  	}
 171  
 172  	return nil
 173  }
 174