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

     1  package ws
     2  
     3  import (
     4  	"encoding/binary"
     5  	"fmt"
     6  	"io"
     7  )
     8  
     9  // Errors used by frame reader.
    10  var (
    11  	ErrHeaderLengthMSB        = fmt.Errorf("header error: the most significant bit must be 0")
    12  	ErrHeaderLengthUnexpected = fmt.Errorf("header error: unexpected payload length bits")
    13  )
    14  
    15  // ReadHeader reads a frame header from r.
    16  func ReadHeader(r io.Reader) (h Header, err error) {
    17  	// Make slice of bytes with capacity 12 that could hold any header.
    18  	//
    19  	// The maximum header size is 14, but due to the 2 hop reads,
    20  	// after first hop that reads first 2 constant bytes, we could reuse 2 bytes.
    21  	// So 14 - 2 = 12.
    22  	bts := make([]byte, 2, MaxHeaderSize-2)
    23  
    24  	// Prepare to hold first 2 bytes to choose size of next read.
    25  	_, err = io.ReadFull(r, bts)
    26  	if err != nil {
    27  		return
    28  	}
    29  
    30  	h.Fin = bts[0]&bit0 != 0
    31  	h.Rsv = (bts[0] & 0x70) >> 4
    32  	h.OpCode = OpCode(bts[0] & 0x0f)
    33  
    34  	var extra int
    35  
    36  	if bts[1]&bit0 != 0 {
    37  		h.Masked = true
    38  		extra += 4
    39  	}
    40  
    41  	length := bts[1] & 0x7f
    42  	switch {
    43  	case length < 126:
    44  		h.Length = int64(length)
    45  
    46  	case length == 126:
    47  		extra += 2
    48  
    49  	case length == 127:
    50  		extra += 8
    51  
    52  	default:
    53  		err = ErrHeaderLengthUnexpected
    54  		return
    55  	}
    56  
    57  	if extra == 0 {
    58  		return
    59  	}
    60  
    61  	// Increase len of bts to extra bytes need to read.
    62  	// Overwrite first 2 bytes that was read before.
    63  	bts = bts[:extra]
    64  	_, err = io.ReadFull(r, bts)
    65  	if err != nil {
    66  		return
    67  	}
    68  
    69  	switch {
    70  	case length == 126:
    71  		h.Length = int64(binary.BigEndian.Uint16(bts[:2]))
    72  		bts = bts[2:]
    73  
    74  	case length == 127:
    75  		if bts[0]&0x80 != 0 {
    76  			err = ErrHeaderLengthMSB
    77  			return
    78  		}
    79  		h.Length = int64(binary.BigEndian.Uint64(bts[:8]))
    80  		bts = bts[8:]
    81  	}
    82  
    83  	if h.Masked {
    84  		copy(h.Mask[:], bts)
    85  	}
    86  
    87  	return
    88  }
    89  
    90  // ReadFrame reads a frame from r.
    91  // It is not designed for high optimized use case cause it makes allocation
    92  // for frame.Header.Length size inside to read frame payload into.
    93  //
    94  // Note that ReadFrame does not unmask payload.
    95  func ReadFrame(r io.Reader) (f Frame, err error) {
    96  	f.Header, err = ReadHeader(r)
    97  	if err != nil {
    98  		return
    99  	}
   100  
   101  	if f.Header.Length > 0 {
   102  		// int(f.Header.Length) is safe here cause we have
   103  		// checked it for overflow above in ReadHeader.
   104  		f.Payload = make([]byte, int(f.Header.Length))
   105  		_, err = io.ReadFull(r, f.Payload)
   106  	}
   107  
   108  	return
   109  }
   110  
   111  // MustReadFrame is like ReadFrame but panics if frame can not be read.
   112  func MustReadFrame(r io.Reader) Frame {
   113  	f, err := ReadFrame(r)
   114  	if err != nil {
   115  		panic(err)
   116  	}
   117  	return f
   118  }
   119  
   120  // ParseCloseFrameData parses close frame status code and closure reason if any provided.
   121  // If there is no status code in the payload
   122  // the empty status code is returned (code.Empty()) with empty string as a reason.
   123  func ParseCloseFrameData(payload []byte) (code StatusCode, reason string) {
   124  	if len(payload) < 2 {
   125  		// We returning empty StatusCode here, preventing the situation
   126  		// when endpoint really sent code 1005 and we should return ProtocolError on that.
   127  		//
   128  		// In other words, we ignoring this rule [RFC6455:7.1.5]:
   129  		//   If this Close control frame contains no status code, _The WebSocket
   130  		//   Connection Close Code_ is considered to be 1005.
   131  		return
   132  	}
   133  	code = StatusCode(binary.BigEndian.Uint16(payload))
   134  	reason = string(payload[2:])
   135  	return
   136  }
   137  
   138  // ParseCloseFrameDataUnsafe is like ParseCloseFrameData except the thing
   139  // that it does not copies payload bytes into reason, but prepares unsafe cast.
   140  func ParseCloseFrameDataUnsafe(payload []byte) (code StatusCode, reason string) {
   141  	if len(payload) < 2 {
   142  		return
   143  	}
   144  	code = StatusCode(binary.BigEndian.Uint16(payload))
   145  	reason = btsToString(payload[2:])
   146  	return
   147  }