util.go 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214
  1. package ws
  2. import (
  3. "bufio"
  4. "bytes"
  5. "fmt"
  6. "reflect"
  7. "unsafe"
  8. "github.com/gobwas/httphead"
  9. )
  10. // SelectFromSlice creates accept function that could be used as Protocol/Extension
  11. // select during upgrade.
  12. func SelectFromSlice(accept []string) func(string) bool {
  13. if len(accept) > 16 {
  14. mp := make(map[string]struct{}, len(accept))
  15. for _, p := range accept {
  16. mp[p] = struct{}{}
  17. }
  18. return func(p string) bool {
  19. _, ok := mp[p]
  20. return ok
  21. }
  22. }
  23. return func(p string) bool {
  24. for _, ok := range accept {
  25. if p == ok {
  26. return true
  27. }
  28. }
  29. return false
  30. }
  31. }
  32. // SelectEqual creates accept function that could be used as Protocol/Extension
  33. // select during upgrade.
  34. func SelectEqual(v string) func(string) bool {
  35. return func(p string) bool {
  36. return v == p
  37. }
  38. }
  39. func strToBytes(str string) (bts []byte) {
  40. s := (*reflect.StringHeader)(unsafe.Pointer(&str))
  41. b := (*reflect.SliceHeader)(unsafe.Pointer(&bts))
  42. b.Data = s.Data
  43. b.Len = s.Len
  44. b.Cap = s.Len
  45. return
  46. }
  47. func btsToString(bts []byte) (str string) {
  48. return *(*string)(unsafe.Pointer(&bts))
  49. }
  50. // asciiToInt converts bytes to int.
  51. func asciiToInt(bts []byte) (ret int, err error) {
  52. // ASCII numbers all start with the high-order bits 0011.
  53. // If you see that, and the next bits are 0-9 (0000 - 1001) you can grab those
  54. // bits and interpret them directly as an integer.
  55. var n int
  56. if n = len(bts); n < 1 {
  57. return 0, fmt.Errorf("converting empty bytes to int")
  58. }
  59. for i := 0; i < n; i++ {
  60. if bts[i]&0xf0 != 0x30 {
  61. return 0, fmt.Errorf("%s is not a numeric character", string(bts[i]))
  62. }
  63. ret += int(bts[i]&0xf) * pow(10, n-i-1)
  64. }
  65. return ret, nil
  66. }
  67. // pow for integers implementation.
  68. // See Donald Knuth, The Art of Computer Programming, Volume 2, Section 4.6.3
  69. func pow(a, b int) int {
  70. p := 1
  71. for b > 0 {
  72. if b&1 != 0 {
  73. p *= a
  74. }
  75. b >>= 1
  76. a *= a
  77. }
  78. return p
  79. }
  80. func bsplit3(bts []byte, sep byte) (b1, b2, b3 []byte) {
  81. a := bytes.IndexByte(bts, sep)
  82. b := bytes.IndexByte(bts[a+1:], sep)
  83. if a == -1 || b == -1 {
  84. return bts, nil, nil
  85. }
  86. b += a + 1
  87. return bts[:a], bts[a+1 : b], bts[b+1:]
  88. }
  89. func btrim(bts []byte) []byte {
  90. var i, j int
  91. for i = 0; i < len(bts) && (bts[i] == ' ' || bts[i] == '\t'); {
  92. i++
  93. }
  94. for j = len(bts); j > i && (bts[j-1] == ' ' || bts[j-1] == '\t'); {
  95. j--
  96. }
  97. return bts[i:j]
  98. }
  99. func strHasToken(header, token string) (has bool) {
  100. return btsHasToken(strToBytes(header), strToBytes(token))
  101. }
  102. func btsHasToken(header, token []byte) (has bool) {
  103. httphead.ScanTokens(header, func(v []byte) bool {
  104. has = bytes.EqualFold(v, token)
  105. return !has
  106. })
  107. return
  108. }
  109. const (
  110. toLower = 'a' - 'A' // for use with OR.
  111. toUpper = ^byte(toLower) // for use with AND.
  112. toLower8 = uint64(toLower) |
  113. uint64(toLower)<<8 |
  114. uint64(toLower)<<16 |
  115. uint64(toLower)<<24 |
  116. uint64(toLower)<<32 |
  117. uint64(toLower)<<40 |
  118. uint64(toLower)<<48 |
  119. uint64(toLower)<<56
  120. )
  121. // Algorithm below is like standard textproto/CanonicalMIMEHeaderKey, except
  122. // that it operates with slice of bytes and modifies it inplace without copying.
  123. func canonicalizeHeaderKey(k []byte) {
  124. upper := true
  125. for i, c := range k {
  126. if upper && 'a' <= c && c <= 'z' {
  127. k[i] &= toUpper
  128. } else if !upper && 'A' <= c && c <= 'Z' {
  129. k[i] |= toLower
  130. }
  131. upper = c == '-'
  132. }
  133. }
  134. // readLine reads line from br. It reads until '\n' and returns bytes without
  135. // '\n' or '\r\n' at the end.
  136. // It returns err if and only if line does not end in '\n'. Note that read
  137. // bytes returned in any case of error.
  138. //
  139. // It is much like the textproto/Reader.ReadLine() except the thing that it
  140. // returns raw bytes, instead of string. That is, it avoids copying bytes read
  141. // from br.
  142. //
  143. // textproto/Reader.ReadLineBytes() is also makes copy of resulting bytes to be
  144. // safe with future I/O operations on br.
  145. //
  146. // We could control I/O operations on br and do not need to make additional
  147. // copy for safety.
  148. //
  149. // NOTE: it may return copied flag to notify that returned buffer is safe to
  150. // use.
  151. func readLine(br *bufio.Reader) ([]byte, error) {
  152. var line []byte
  153. for {
  154. bts, err := br.ReadSlice('\n')
  155. if err == bufio.ErrBufferFull {
  156. // Copy bytes because next read will discard them.
  157. line = append(line, bts...)
  158. continue
  159. }
  160. // Avoid copy of single read.
  161. if line == nil {
  162. line = bts
  163. } else {
  164. line = append(line, bts...)
  165. }
  166. if err != nil {
  167. return line, err
  168. }
  169. // Size of line is at least 1.
  170. // In other case bufio.ReadSlice() returns error.
  171. n := len(line)
  172. // Cut '\n' or '\r\n'.
  173. if n > 1 && line[n-2] == '\r' {
  174. line = line[:n-2]
  175. } else {
  176. line = line[:n-1]
  177. }
  178. return line, nil
  179. }
  180. }
  181. func min(a, b int) int {
  182. if a < b {
  183. return a
  184. }
  185. return b
  186. }
  187. func nonZero(a, b int) int {
  188. if a != 0 {
  189. return a
  190. }
  191. return b
  192. }