github.com/simonmittag/ws@v1.1.0-rc.5.0.20210419231947-82b846128245/util.go (about)

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