websocket.go 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131
  1. package net
  2. import (
  3. "encoding/binary"
  4. "math/bits"
  5. )
  6. // kanged from https://github.com/nhooyr/websocket/blob/master/frame.go
  7. // License: MIT
  8. // MaskWebSocket applies the WebSocket masking algorithm to p
  9. // with the given key.
  10. // See https://tools.ietf.org/html/rfc6455#section-5.3
  11. //
  12. // The returned value is the correctly rotated key to
  13. // to continue to mask/unmask the message.
  14. //
  15. // It is optimized for LittleEndian and expects the key
  16. // to be in little endian.
  17. //
  18. // See https://github.com/golang/go/issues/31586
  19. func MaskWebSocket(key uint32, b []byte) uint32 {
  20. if len(b) >= 8 {
  21. key64 := uint64(key)<<32 | uint64(key)
  22. // At some point in the future we can clean these unrolled loops up.
  23. // See https://github.com/golang/go/issues/31586#issuecomment-487436401
  24. // Then we xor until b is less than 128 bytes.
  25. for len(b) >= 128 {
  26. v := binary.LittleEndian.Uint64(b)
  27. binary.LittleEndian.PutUint64(b, v^key64)
  28. v = binary.LittleEndian.Uint64(b[8:16])
  29. binary.LittleEndian.PutUint64(b[8:16], v^key64)
  30. v = binary.LittleEndian.Uint64(b[16:24])
  31. binary.LittleEndian.PutUint64(b[16:24], v^key64)
  32. v = binary.LittleEndian.Uint64(b[24:32])
  33. binary.LittleEndian.PutUint64(b[24:32], v^key64)
  34. v = binary.LittleEndian.Uint64(b[32:40])
  35. binary.LittleEndian.PutUint64(b[32:40], v^key64)
  36. v = binary.LittleEndian.Uint64(b[40:48])
  37. binary.LittleEndian.PutUint64(b[40:48], v^key64)
  38. v = binary.LittleEndian.Uint64(b[48:56])
  39. binary.LittleEndian.PutUint64(b[48:56], v^key64)
  40. v = binary.LittleEndian.Uint64(b[56:64])
  41. binary.LittleEndian.PutUint64(b[56:64], v^key64)
  42. v = binary.LittleEndian.Uint64(b[64:72])
  43. binary.LittleEndian.PutUint64(b[64:72], v^key64)
  44. v = binary.LittleEndian.Uint64(b[72:80])
  45. binary.LittleEndian.PutUint64(b[72:80], v^key64)
  46. v = binary.LittleEndian.Uint64(b[80:88])
  47. binary.LittleEndian.PutUint64(b[80:88], v^key64)
  48. v = binary.LittleEndian.Uint64(b[88:96])
  49. binary.LittleEndian.PutUint64(b[88:96], v^key64)
  50. v = binary.LittleEndian.Uint64(b[96:104])
  51. binary.LittleEndian.PutUint64(b[96:104], v^key64)
  52. v = binary.LittleEndian.Uint64(b[104:112])
  53. binary.LittleEndian.PutUint64(b[104:112], v^key64)
  54. v = binary.LittleEndian.Uint64(b[112:120])
  55. binary.LittleEndian.PutUint64(b[112:120], v^key64)
  56. v = binary.LittleEndian.Uint64(b[120:128])
  57. binary.LittleEndian.PutUint64(b[120:128], v^key64)
  58. b = b[128:]
  59. }
  60. // Then we xor until b is less than 64 bytes.
  61. for len(b) >= 64 {
  62. v := binary.LittleEndian.Uint64(b)
  63. binary.LittleEndian.PutUint64(b, v^key64)
  64. v = binary.LittleEndian.Uint64(b[8:16])
  65. binary.LittleEndian.PutUint64(b[8:16], v^key64)
  66. v = binary.LittleEndian.Uint64(b[16:24])
  67. binary.LittleEndian.PutUint64(b[16:24], v^key64)
  68. v = binary.LittleEndian.Uint64(b[24:32])
  69. binary.LittleEndian.PutUint64(b[24:32], v^key64)
  70. v = binary.LittleEndian.Uint64(b[32:40])
  71. binary.LittleEndian.PutUint64(b[32:40], v^key64)
  72. v = binary.LittleEndian.Uint64(b[40:48])
  73. binary.LittleEndian.PutUint64(b[40:48], v^key64)
  74. v = binary.LittleEndian.Uint64(b[48:56])
  75. binary.LittleEndian.PutUint64(b[48:56], v^key64)
  76. v = binary.LittleEndian.Uint64(b[56:64])
  77. binary.LittleEndian.PutUint64(b[56:64], v^key64)
  78. b = b[64:]
  79. }
  80. // Then we xor until b is less than 32 bytes.
  81. for len(b) >= 32 {
  82. v := binary.LittleEndian.Uint64(b)
  83. binary.LittleEndian.PutUint64(b, v^key64)
  84. v = binary.LittleEndian.Uint64(b[8:16])
  85. binary.LittleEndian.PutUint64(b[8:16], v^key64)
  86. v = binary.LittleEndian.Uint64(b[16:24])
  87. binary.LittleEndian.PutUint64(b[16:24], v^key64)
  88. v = binary.LittleEndian.Uint64(b[24:32])
  89. binary.LittleEndian.PutUint64(b[24:32], v^key64)
  90. b = b[32:]
  91. }
  92. // Then we xor until b is less than 16 bytes.
  93. for len(b) >= 16 {
  94. v := binary.LittleEndian.Uint64(b)
  95. binary.LittleEndian.PutUint64(b, v^key64)
  96. v = binary.LittleEndian.Uint64(b[8:16])
  97. binary.LittleEndian.PutUint64(b[8:16], v^key64)
  98. b = b[16:]
  99. }
  100. // Then we xor until b is less than 8 bytes.
  101. for len(b) >= 8 {
  102. v := binary.LittleEndian.Uint64(b)
  103. binary.LittleEndian.PutUint64(b, v^key64)
  104. b = b[8:]
  105. }
  106. }
  107. // Then we xor until b is less than 4 bytes.
  108. for len(b) >= 4 {
  109. v := binary.LittleEndian.Uint32(b)
  110. binary.LittleEndian.PutUint32(b, v^key)
  111. b = b[4:]
  112. }
  113. // xor remaining bytes.
  114. for i := range b {
  115. b[i] ^= byte(key)
  116. key = bits.RotateLeft32(key, -8)
  117. }
  118. return key
  119. }