Skip to content

Commit 083d297

Browse files
nhooyrnhooyr
committed
Refactor and compile masking code again
1 parent d4b8f48 commit 083d297

File tree

8 files changed

+139
-137
lines changed

8 files changed

+139
-137
lines changed

‎frame.go

-123
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ import (
88
"fmt"
99
"io"
1010
"math"
11-
"math/bits"
1211

1312
"nhooyr.io/websocket/internal/errd"
1413
)
@@ -172,125 +171,3 @@ func writeFrameHeader(h header, w *bufio.Writer, buf []byte) (err error) {
172171

173172
return nil
174173
}
175-
176-
// mask applies the WebSocket masking algorithm to p
177-
// with the given key.
178-
// See https://tools.ietf.org/html/rfc6455#section-5.3
179-
//
180-
// The returned value is the correctly rotated key to
181-
// to continue to mask/unmask the message.
182-
//
183-
// It is optimized for LittleEndian and expects the key
184-
// to be in little endian.
185-
//
186-
// See https://github.com/golang/go/issues/31586
187-
func maskGo(key uint32, b []byte) uint32 {
188-
if len(b) >= 8 {
189-
key64 := uint64(key)<<32 | uint64(key)
190-
191-
// At some point in the future we can clean these unrolled loops up.
192-
// See https://github.com/golang/go/issues/31586#issuecomment-487436401
193-
194-
// Then we xor until b is less than 128 bytes.
195-
for len(b) >= 128 {
196-
v := binary.LittleEndian.Uint64(b)
197-
binary.LittleEndian.PutUint64(b, v^key64)
198-
v = binary.LittleEndian.Uint64(b[8:16])
199-
binary.LittleEndian.PutUint64(b[8:16], v^key64)
200-
v = binary.LittleEndian.Uint64(b[16:24])
201-
binary.LittleEndian.PutUint64(b[16:24], v^key64)
202-
v = binary.LittleEndian.Uint64(b[24:32])
203-
binary.LittleEndian.PutUint64(b[24:32], v^key64)
204-
v = binary.LittleEndian.Uint64(b[32:40])
205-
binary.LittleEndian.PutUint64(b[32:40], v^key64)
206-
v = binary.LittleEndian.Uint64(b[40:48])
207-
binary.LittleEndian.PutUint64(b[40:48], v^key64)
208-
v = binary.LittleEndian.Uint64(b[48:56])
209-
binary.LittleEndian.PutUint64(b[48:56], v^key64)
210-
v = binary.LittleEndian.Uint64(b[56:64])
211-
binary.LittleEndian.PutUint64(b[56:64], v^key64)
212-
v = binary.LittleEndian.Uint64(b[64:72])
213-
binary.LittleEndian.PutUint64(b[64:72], v^key64)
214-
v = binary.LittleEndian.Uint64(b[72:80])
215-
binary.LittleEndian.PutUint64(b[72:80], v^key64)
216-
v = binary.LittleEndian.Uint64(b[80:88])
217-
binary.LittleEndian.PutUint64(b[80:88], v^key64)
218-
v = binary.LittleEndian.Uint64(b[88:96])
219-
binary.LittleEndian.PutUint64(b[88:96], v^key64)
220-
v = binary.LittleEndian.Uint64(b[96:104])
221-
binary.LittleEndian.PutUint64(b[96:104], v^key64)
222-
v = binary.LittleEndian.Uint64(b[104:112])
223-
binary.LittleEndian.PutUint64(b[104:112], v^key64)
224-
v = binary.LittleEndian.Uint64(b[112:120])
225-
binary.LittleEndian.PutUint64(b[112:120], v^key64)
226-
v = binary.LittleEndian.Uint64(b[120:128])
227-
binary.LittleEndian.PutUint64(b[120:128], v^key64)
228-
b = b[128:]
229-
}
230-
231-
// Then we xor until b is less than 64 bytes.
232-
for len(b) >= 64 {
233-
v := binary.LittleEndian.Uint64(b)
234-
binary.LittleEndian.PutUint64(b, v^key64)
235-
v = binary.LittleEndian.Uint64(b[8:16])
236-
binary.LittleEndian.PutUint64(b[8:16], v^key64)
237-
v = binary.LittleEndian.Uint64(b[16:24])
238-
binary.LittleEndian.PutUint64(b[16:24], v^key64)
239-
v = binary.LittleEndian.Uint64(b[24:32])
240-
binary.LittleEndian.PutUint64(b[24:32], v^key64)
241-
v = binary.LittleEndian.Uint64(b[32:40])
242-
binary.LittleEndian.PutUint64(b[32:40], v^key64)
243-
v = binary.LittleEndian.Uint64(b[40:48])
244-
binary.LittleEndian.PutUint64(b[40:48], v^key64)
245-
v = binary.LittleEndian.Uint64(b[48:56])
246-
binary.LittleEndian.PutUint64(b[48:56], v^key64)
247-
v = binary.LittleEndian.Uint64(b[56:64])
248-
binary.LittleEndian.PutUint64(b[56:64], v^key64)
249-
b = b[64:]
250-
}
251-
252-
// Then we xor until b is less than 32 bytes.
253-
for len(b) >= 32 {
254-
v := binary.LittleEndian.Uint64(b)
255-
binary.LittleEndian.PutUint64(b, v^key64)
256-
v = binary.LittleEndian.Uint64(b[8:16])
257-
binary.LittleEndian.PutUint64(b[8:16], v^key64)
258-
v = binary.LittleEndian.Uint64(b[16:24])
259-
binary.LittleEndian.PutUint64(b[16:24], v^key64)
260-
v = binary.LittleEndian.Uint64(b[24:32])
261-
binary.LittleEndian.PutUint64(b[24:32], v^key64)
262-
b = b[32:]
263-
}
264-
265-
// Then we xor until b is less than 16 bytes.
266-
for len(b) >= 16 {
267-
v := binary.LittleEndian.Uint64(b)
268-
binary.LittleEndian.PutUint64(b, v^key64)
269-
v = binary.LittleEndian.Uint64(b[8:16])
270-
binary.LittleEndian.PutUint64(b[8:16], v^key64)
271-
b = b[16:]
272-
}
273-
274-
// Then we xor until b is less than 8 bytes.
275-
for len(b) >= 8 {
276-
v := binary.LittleEndian.Uint64(b)
277-
binary.LittleEndian.PutUint64(b, v^key64)
278-
b = b[8:]
279-
}
280-
}
281-
282-
// Then we xor until b is less than 4 bytes.
283-
for len(b) >= 4 {
284-
v := binary.LittleEndian.Uint32(b)
285-
binary.LittleEndian.PutUint32(b, v^key)
286-
b = b[4:]
287-
}
288-
289-
// xor remaining bytes.
290-
for i := range b {
291-
b[i] ^= byte(key)
292-
key = bits.RotateLeft32(key, -8)
293-
}
294-
295-
return key
296-
}

‎internal/examples/go.mod

+2
Original file line numberDiff line numberDiff line change
@@ -8,3 +8,5 @@ require (
88
golang.org/x/time v0.3.0
99
nhooyr.io/websocket v0.0.0-00010101000000-000000000000
1010
)
11+
12+
require golang.org/x/sys v0.13.0 // indirect

‎internal/examples/go.sum

+2
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,4 @@
1+
golang.org/x/sys v0.13.0 h1:Af8nKPmuFypiUBjVoU9V20FiaFXOcuZI21p0ycVYYGE=
2+
golang.org/x/sys v0.13.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
13
golang.org/x/time v0.3.0 h1:rg5rLMjNzMS1RkNLzCG38eapWhnYLFYXDXj2gOlr8j4=
24
golang.org/x/time v0.3.0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=

‎internal/thirdparty/go.mod

+1-1
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ require (
3434
golang.org/x/arch v0.3.0 // indirect
3535
golang.org/x/crypto v0.9.0 // indirect
3636
golang.org/x/net v0.10.0 // indirect
37-
golang.org/x/sys v0.8.0 // indirect
37+
golang.org/x/sys v0.13.0 // indirect
3838
golang.org/x/text v0.9.0 // indirect
3939
google.golang.org/protobuf v1.30.0 // indirect
4040
gopkg.in/yaml.v3 v3.0.1 // indirect

‎internal/thirdparty/go.sum

+2-2
Original file line numberDiff line numberDiff line change
@@ -76,8 +76,8 @@ golang.org/x/net v0.10.0 h1:X2//UzNDwYmtCLn7To6G58Wr6f5ahEAQgKNzv9Y951M=
7676
golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg=
7777
golang.org/x/sys v0.0.0-20220704084225-05e143d24a9e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
7878
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
79-
golang.org/x/sys v0.8.0 h1:EBmGv8NaZBZTWvrbjNoL6HVt+IVy3QDQpJs7VRIw3tU=
80-
golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
79+
golang.org/x/sys v0.13.0 h1:Af8nKPmuFypiUBjVoU9V20FiaFXOcuZI21p0ycVYYGE=
80+
golang.org/x/sys v0.13.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
8181
golang.org/x/text v0.9.0 h1:2sjJmO8cDvYveuX97RDLsxlyUxLl+GHoLxBiRdHllBE=
8282
golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
8383
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=

‎mask.go

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

‎mask_asm.go

+2-4
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,4 @@
1-
//go:build !appengine && (amd64 || arm64)
2-
// +build !appengine
3-
// +build amd64 arm64
1+
//go:build amd64 || arm64
42

53
package websocket
64

@@ -13,7 +11,7 @@ func mask(key uint32, b []byte) uint32 {
1311
return key
1412
}
1513

16-
var useAVX2 = cpu.X86.HasAVX2
14+
var useAVX2 = cpu.X86.HasAVX2 //lint:ignore U1000 mask_amd64.s
1715

1816
//go:noescape
1917
func maskAsm(b *byte, len int, key uint32) uint32

‎mask_generic.go

-7
This file was deleted.

0 commit comments

Comments
 (0)