frame.go 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420
  1. package ws
  2. import (
  3. "bytes"
  4. "encoding/binary"
  5. "math/rand"
  6. )
  7. // Constants defined by specification.
  8. const (
  9. // All control frames MUST have a payload length of 125 bytes or less and MUST NOT be fragmented.
  10. MaxControlFramePayloadSize = 125
  11. )
  12. // OpCode represents operation code.
  13. type OpCode byte
  14. // Operation codes defined by specification.
  15. // See https://tools.ietf.org/html/rfc6455#section-5.2
  16. const (
  17. OpContinuation OpCode = 0x0
  18. OpText OpCode = 0x1
  19. OpBinary OpCode = 0x2
  20. OpClose OpCode = 0x8
  21. OpPing OpCode = 0x9
  22. OpPong OpCode = 0xa
  23. )
  24. // IsControl checks whether the c is control operation code.
  25. // See https://tools.ietf.org/html/rfc6455#section-5.5
  26. func (c OpCode) IsControl() bool {
  27. // RFC6455: Control frames are identified by opcodes where
  28. // the most significant bit of the opcode is 1.
  29. //
  30. // Note that OpCode is only 4 bit length.
  31. return c&0x8 != 0
  32. }
  33. // IsData checks whether the c is data operation code.
  34. // See https://tools.ietf.org/html/rfc6455#section-5.6
  35. func (c OpCode) IsData() bool {
  36. // RFC6455: Data frames (e.g., non-control frames) are identified by opcodes
  37. // where the most significant bit of the opcode is 0.
  38. //
  39. // Note that OpCode is only 4 bit length.
  40. return c&0x8 == 0
  41. }
  42. // IsReserved checks whether the c is reserved operation code.
  43. // See https://tools.ietf.org/html/rfc6455#section-5.2
  44. func (c OpCode) IsReserved() bool {
  45. // RFC6455:
  46. // %x3-7 are reserved for further non-control frames
  47. // %xB-F are reserved for further control frames
  48. return (0x3 <= c && c <= 0x7) || (0xb <= c && c <= 0xf)
  49. }
  50. // StatusCode represents the encoded reason for closure of websocket connection.
  51. //
  52. // There are few helper methods on StatusCode that helps to define a range in
  53. // which given code is lay in. accordingly to ranges defined in specification.
  54. //
  55. // See https://tools.ietf.org/html/rfc6455#section-7.4
  56. type StatusCode uint16
  57. // StatusCodeRange describes range of StatusCode values.
  58. type StatusCodeRange struct {
  59. Min, Max StatusCode
  60. }
  61. // Status code ranges defined by specification.
  62. // See https://tools.ietf.org/html/rfc6455#section-7.4.2
  63. var (
  64. StatusRangeNotInUse = StatusCodeRange{0, 999}
  65. StatusRangeProtocol = StatusCodeRange{1000, 2999}
  66. StatusRangeApplication = StatusCodeRange{3000, 3999}
  67. StatusRangePrivate = StatusCodeRange{4000, 4999}
  68. )
  69. // Status codes defined by specification.
  70. // See https://tools.ietf.org/html/rfc6455#section-7.4.1
  71. const (
  72. StatusNormalClosure StatusCode = 1000
  73. StatusGoingAway StatusCode = 1001
  74. StatusProtocolError StatusCode = 1002
  75. StatusUnsupportedData StatusCode = 1003
  76. StatusNoMeaningYet StatusCode = 1004
  77. StatusInvalidFramePayloadData StatusCode = 1007
  78. StatusPolicyViolation StatusCode = 1008
  79. StatusMessageTooBig StatusCode = 1009
  80. StatusMandatoryExt StatusCode = 1010
  81. StatusInternalServerError StatusCode = 1011
  82. StatusTLSHandshake StatusCode = 1015
  83. // StatusAbnormalClosure is a special code designated for use in
  84. // applications.
  85. StatusAbnormalClosure StatusCode = 1006
  86. // StatusNoStatusRcvd is a special code designated for use in applications.
  87. StatusNoStatusRcvd StatusCode = 1005
  88. )
  89. // In reports whether the code is defined in given range.
  90. func (s StatusCode) In(r StatusCodeRange) bool {
  91. return r.Min <= s && s <= r.Max
  92. }
  93. // Empty reports whether the code is empty.
  94. // Empty code has no any meaning neither app level codes nor other.
  95. // This method is useful just to check that code is golang default value 0.
  96. func (s StatusCode) Empty() bool {
  97. return s == 0
  98. }
  99. // IsNotUsed reports whether the code is predefined in not used range.
  100. func (s StatusCode) IsNotUsed() bool {
  101. return s.In(StatusRangeNotInUse)
  102. }
  103. // IsApplicationSpec reports whether the code should be defined by
  104. // application, framework or libraries specification.
  105. func (s StatusCode) IsApplicationSpec() bool {
  106. return s.In(StatusRangeApplication)
  107. }
  108. // IsPrivateSpec reports whether the code should be defined privately.
  109. func (s StatusCode) IsPrivateSpec() bool {
  110. return s.In(StatusRangePrivate)
  111. }
  112. // IsProtocolSpec reports whether the code should be defined by protocol specification.
  113. func (s StatusCode) IsProtocolSpec() bool {
  114. return s.In(StatusRangeProtocol)
  115. }
  116. // IsProtocolDefined reports whether the code is already defined by protocol specification.
  117. func (s StatusCode) IsProtocolDefined() bool {
  118. switch s {
  119. case StatusNormalClosure,
  120. StatusGoingAway,
  121. StatusProtocolError,
  122. StatusUnsupportedData,
  123. StatusInvalidFramePayloadData,
  124. StatusPolicyViolation,
  125. StatusMessageTooBig,
  126. StatusMandatoryExt,
  127. StatusInternalServerError,
  128. StatusNoStatusRcvd,
  129. StatusAbnormalClosure,
  130. StatusTLSHandshake:
  131. return true
  132. }
  133. return false
  134. }
  135. // IsProtocolReserved reports whether the code is defined by protocol specification
  136. // to be reserved only for application usage purpose.
  137. func (s StatusCode) IsProtocolReserved() bool {
  138. switch s {
  139. // [RFC6455]: {1005,1006,1015} is a reserved value and MUST NOT be set as a status code in a
  140. // Close control frame by an endpoint.
  141. case StatusNoStatusRcvd, StatusAbnormalClosure, StatusTLSHandshake:
  142. return true
  143. default:
  144. return false
  145. }
  146. }
  147. // Compiled control frames for common use cases.
  148. // For construct-serialize optimizations.
  149. var (
  150. CompiledPing = MustCompileFrame(NewPingFrame(nil))
  151. CompiledPong = MustCompileFrame(NewPongFrame(nil))
  152. CompiledClose = MustCompileFrame(NewCloseFrame(nil))
  153. CompiledCloseNormalClosure = MustCompileFrame(closeFrameNormalClosure)
  154. CompiledCloseGoingAway = MustCompileFrame(closeFrameGoingAway)
  155. CompiledCloseProtocolError = MustCompileFrame(closeFrameProtocolError)
  156. CompiledCloseUnsupportedData = MustCompileFrame(closeFrameUnsupportedData)
  157. CompiledCloseNoMeaningYet = MustCompileFrame(closeFrameNoMeaningYet)
  158. CompiledCloseInvalidFramePayloadData = MustCompileFrame(closeFrameInvalidFramePayloadData)
  159. CompiledClosePolicyViolation = MustCompileFrame(closeFramePolicyViolation)
  160. CompiledCloseMessageTooBig = MustCompileFrame(closeFrameMessageTooBig)
  161. CompiledCloseMandatoryExt = MustCompileFrame(closeFrameMandatoryExt)
  162. CompiledCloseInternalServerError = MustCompileFrame(closeFrameInternalServerError)
  163. CompiledCloseTLSHandshake = MustCompileFrame(closeFrameTLSHandshake)
  164. )
  165. // Header represents websocket frame header.
  166. // See https://tools.ietf.org/html/rfc6455#section-5.2
  167. type Header struct {
  168. Fin bool
  169. Rsv byte
  170. OpCode OpCode
  171. Masked bool
  172. Mask [4]byte
  173. Length int64
  174. }
  175. // Rsv1 reports whether the header has first rsv bit set.
  176. func (h Header) Rsv1() bool { return h.Rsv&bit5 != 0 }
  177. // Rsv2 reports whether the header has second rsv bit set.
  178. func (h Header) Rsv2() bool { return h.Rsv&bit6 != 0 }
  179. // Rsv3 reports whether the header has third rsv bit set.
  180. func (h Header) Rsv3() bool { return h.Rsv&bit7 != 0 }
  181. // Rsv creates rsv byte representation from bits.
  182. func Rsv(r1, r2, r3 bool) (rsv byte) {
  183. if r1 {
  184. rsv |= bit5
  185. }
  186. if r2 {
  187. rsv |= bit6
  188. }
  189. if r3 {
  190. rsv |= bit7
  191. }
  192. return rsv
  193. }
  194. // RsvBits returns rsv bits from bytes representation.
  195. func RsvBits(rsv byte) (r1, r2, r3 bool) {
  196. r1 = rsv&bit5 != 0
  197. r2 = rsv&bit6 != 0
  198. r3 = rsv&bit7 != 0
  199. return
  200. }
  201. // Frame represents websocket frame.
  202. // See https://tools.ietf.org/html/rfc6455#section-5.2
  203. type Frame struct {
  204. Header Header
  205. Payload []byte
  206. }
  207. // NewFrame creates frame with given operation code,
  208. // flag of completeness and payload bytes.
  209. func NewFrame(op OpCode, fin bool, p []byte) Frame {
  210. return Frame{
  211. Header: Header{
  212. Fin: fin,
  213. OpCode: op,
  214. Length: int64(len(p)),
  215. },
  216. Payload: p,
  217. }
  218. }
  219. // NewTextFrame creates text frame with p as payload.
  220. // Note that p is not copied.
  221. func NewTextFrame(p []byte) Frame {
  222. return NewFrame(OpText, true, p)
  223. }
  224. // NewBinaryFrame creates binary frame with p as payload.
  225. // Note that p is not copied.
  226. func NewBinaryFrame(p []byte) Frame {
  227. return NewFrame(OpBinary, true, p)
  228. }
  229. // NewPingFrame creates ping frame with p as payload.
  230. // Note that p is not copied.
  231. // Note that p must have length of MaxControlFramePayloadSize bytes or less due
  232. // to RFC.
  233. func NewPingFrame(p []byte) Frame {
  234. return NewFrame(OpPing, true, p)
  235. }
  236. // NewPongFrame creates pong frame with p as payload.
  237. // Note that p is not copied.
  238. // Note that p must have length of MaxControlFramePayloadSize bytes or less due
  239. // to RFC.
  240. func NewPongFrame(p []byte) Frame {
  241. return NewFrame(OpPong, true, p)
  242. }
  243. // NewCloseFrame creates close frame with given close body.
  244. // Note that p is not copied.
  245. // Note that p must have length of MaxControlFramePayloadSize bytes or less due
  246. // to RFC.
  247. func NewCloseFrame(p []byte) Frame {
  248. return NewFrame(OpClose, true, p)
  249. }
  250. // NewCloseFrameBody encodes a closure code and a reason into a binary
  251. // representation.
  252. //
  253. // It returns slice which is at most MaxControlFramePayloadSize bytes length.
  254. // If the reason is too big it will be cropped to fit the limit defined by the
  255. // spec.
  256. //
  257. // See https://tools.ietf.org/html/rfc6455#section-5.5
  258. func NewCloseFrameBody(code StatusCode, reason string) []byte {
  259. n := min(2+len(reason), MaxControlFramePayloadSize)
  260. p := make([]byte, n)
  261. crop := min(MaxControlFramePayloadSize-2, len(reason))
  262. PutCloseFrameBody(p, code, reason[:crop])
  263. return p
  264. }
  265. // PutCloseFrameBody encodes code and reason into buf.
  266. //
  267. // It will panic if the buffer is too small to accommodate a code or a reason.
  268. //
  269. // PutCloseFrameBody does not check buffer to be RFC compliant, but note that
  270. // by RFC it must be at most MaxControlFramePayloadSize.
  271. func PutCloseFrameBody(p []byte, code StatusCode, reason string) {
  272. _ = p[1+len(reason)]
  273. binary.BigEndian.PutUint16(p, uint16(code))
  274. copy(p[2:], reason)
  275. }
  276. // MaskFrame masks frame and returns frame with masked payload and Mask header's field set.
  277. // Note that it copies f payload to prevent collisions.
  278. // For less allocations you could use MaskFrameInPlace or construct frame manually.
  279. func MaskFrame(f Frame) Frame {
  280. return MaskFrameWith(f, NewMask())
  281. }
  282. // MaskFrameWith masks frame with given mask and returns frame
  283. // with masked payload and Mask header's field set.
  284. // Note that it copies f payload to prevent collisions.
  285. // For less allocations you could use MaskFrameInPlaceWith or construct frame manually.
  286. func MaskFrameWith(f Frame, mask [4]byte) Frame {
  287. // TODO(gobwas): check CopyCipher ws copy() Cipher().
  288. p := make([]byte, len(f.Payload))
  289. copy(p, f.Payload)
  290. f.Payload = p
  291. return MaskFrameInPlaceWith(f, mask)
  292. }
  293. // MaskFrameInPlace masks frame and returns frame with masked payload and Mask
  294. // header's field set.
  295. // Note that it applies xor cipher to f.Payload without copying, that is, it
  296. // modifies f.Payload inplace.
  297. func MaskFrameInPlace(f Frame) Frame {
  298. return MaskFrameInPlaceWith(f, NewMask())
  299. }
  300. var zeroMask [4]byte
  301. // UnmaskFrame unmasks frame and returns frame with unmasked payload and Mask
  302. // header's field cleared.
  303. // Note that it copies f payload.
  304. func UnmaskFrame(f Frame) Frame {
  305. p := make([]byte, len(f.Payload))
  306. copy(p, f.Payload)
  307. f.Payload = p
  308. return UnmaskFrameInPlace(f)
  309. }
  310. // UnmaskFrameInPlace unmasks frame and returns frame with unmasked payload and
  311. // Mask header's field cleared.
  312. // Note that it applies xor cipher to f.Payload without copying, that is, it
  313. // modifies f.Payload inplace.
  314. func UnmaskFrameInPlace(f Frame) Frame {
  315. Cipher(f.Payload, f.Header.Mask, 0)
  316. f.Header.Masked = false
  317. f.Header.Mask = zeroMask
  318. return f
  319. }
  320. // MaskFrameInPlaceWith masks frame with given mask and returns frame
  321. // with masked payload and Mask header's field set.
  322. // Note that it applies xor cipher to f.Payload without copying, that is, it
  323. // modifies f.Payload inplace.
  324. func MaskFrameInPlaceWith(f Frame, m [4]byte) Frame {
  325. f.Header.Masked = true
  326. f.Header.Mask = m
  327. Cipher(f.Payload, m, 0)
  328. return f
  329. }
  330. // NewMask creates new random mask.
  331. func NewMask() (ret [4]byte) {
  332. binary.BigEndian.PutUint32(ret[:], rand.Uint32())
  333. return
  334. }
  335. // CompileFrame returns byte representation of given frame.
  336. // In terms of memory consumption it is useful to precompile static frames
  337. // which are often used.
  338. func CompileFrame(f Frame) (bts []byte, err error) {
  339. buf := bytes.NewBuffer(make([]byte, 0, 16))
  340. err = WriteFrame(buf, f)
  341. bts = buf.Bytes()
  342. return
  343. }
  344. // MustCompileFrame is like CompileFrame but panics if frame can not be
  345. // encoded.
  346. func MustCompileFrame(f Frame) []byte {
  347. bts, err := CompileFrame(f)
  348. if err != nil {
  349. panic(err)
  350. }
  351. return bts
  352. }
  353. func makeCloseFrame(code StatusCode) Frame {
  354. return NewCloseFrame(NewCloseFrameBody(code, ""))
  355. }
  356. var (
  357. closeFrameNormalClosure = makeCloseFrame(StatusNormalClosure)
  358. closeFrameGoingAway = makeCloseFrame(StatusGoingAway)
  359. closeFrameProtocolError = makeCloseFrame(StatusProtocolError)
  360. closeFrameUnsupportedData = makeCloseFrame(StatusUnsupportedData)
  361. closeFrameNoMeaningYet = makeCloseFrame(StatusNoMeaningYet)
  362. closeFrameInvalidFramePayloadData = makeCloseFrame(StatusInvalidFramePayloadData)
  363. closeFramePolicyViolation = makeCloseFrame(StatusPolicyViolation)
  364. closeFrameMessageTooBig = makeCloseFrame(StatusMessageTooBig)
  365. closeFrameMandatoryExt = makeCloseFrame(StatusMandatoryExt)
  366. closeFrameInternalServerError = makeCloseFrame(StatusInternalServerError)
  367. closeFrameTLSHandshake = makeCloseFrame(StatusTLSHandshake)
  368. )