util.go 4.4 KB

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