read.go 3.5 KB

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