Skip to content

Commit

Permalink
mask.go: Cleanup assembly and add nbio benchmark
Browse files Browse the repository at this point in the history
  • Loading branch information
@nhooyr
nhooyr committed Oct 19, 2023
1 parent 9d1b1f6 commit 32d0aa1
Show file tree
Hide file tree
Showing 11 changed files with 79 additions and 29 deletions.
2 changes: 1 addition & 1 deletion frame.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -186,7 +186,7 @@ func writeFrameHeader(h header, w *bufio.Writer, buf []byte) (err error) {
// See https://github.com/golang/go/issues/31586
//
//lint:ignore U1000 mask.go
func maskGo(key uint32, b []byte) uint32 {
func maskGo(b []byte, key uint32) uint32 {
if len(b) >= 8 {
key64 := uint64(key)<<32 | uint64(key)

Expand Down
2 changes: 1 addition & 1 deletion frame_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ func Test_mask(t *testing.T) {
key := []byte{0xa, 0xb, 0xc, 0xff}
key32 := binary.LittleEndian.Uint32(key)
p := []byte{0xa, 0xb, 0xc, 0xf2, 0xc}
gotKey32 := mask(key32, p)
gotKey32 := mask(p, key32)

expP := []byte{0, 0, 0, 0x0d, 0x6}
assert.Equal(t, "p", expP, p)
Expand Down
50 changes: 32 additions & 18 deletions internal/thirdparty/frame_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,38 +8,43 @@ import (

"github.com/gobwas/ws"
_ "github.com/gorilla/websocket"
_ "github.com/lesismal/nbio/nbhttp/websocket"

_ "nhooyr.io/websocket"
)

func basicMask(maskKey [4]byte, pos int, b []byte) int {
func basicMask(b []byte, maskKey [4]byte, pos int) int {
for i := range b {
b[i] ^= maskKey[pos&3]
pos++
}
return pos & 3
}

//go:linkname gorillaMaskBytes github.com/gorilla/websocket.maskBytes
func gorillaMaskBytes(key [4]byte, pos int, b []byte) int
//go:linkname maskGo nhooyr.io/websocket.maskGo
func maskGo(b []byte, key32 uint32) int

//go:linkname mask nhooyr.io/websocket.mask
func mask(key32 uint32, b []byte) int
//go:linkname maskAsm nhooyr.io/websocket.maskAsm
func maskAsm(b *byte, len int, key32 uint32) uint32

//go:linkname maskGo nhooyr.io/websocket.maskGo
func maskGo(key32 uint32, b []byte) int
//go:linkname nbioMaskBytes github.com/lesismal/nbio/nbhttp/websocket.maskXOR
func nbioMaskBytes(b, key []byte) int

//go:linkname gorillaMaskBytes github.com/gorilla/websocket.maskBytes
func gorillaMaskBytes(key [4]byte, pos int, b []byte) int

func Benchmark_mask(b *testing.B) {
sizes := []int{
2,
3,
4,
8,
16,
32,
128,
256,
512,
1024,
2048,
4096,
8192,
16384,
}

Expand All @@ -51,7 +56,7 @@ func Benchmark_mask(b *testing.B) {
name: "basic",
fn: func(b *testing.B, key [4]byte, p []byte) {
for i := 0; i < b.N; i++ {
basicMask(key, 0, p)
basicMask(p, key, 0)
}
},
},
Expand All @@ -63,7 +68,7 @@ func Benchmark_mask(b *testing.B) {
b.ResetTimer()

for i := 0; i < b.N; i++ {
maskGo(key32, p)
maskGo(p, key32)
}
},
},
Expand All @@ -74,7 +79,7 @@ func Benchmark_mask(b *testing.B) {
b.ResetTimer()

for i := 0; i < b.N; i++ {
mask(key32, p)
maskAsm(&p[0], len(p), key32)
}
},
},
Expand All @@ -95,16 +100,25 @@ func Benchmark_mask(b *testing.B) {
}
},
},
{
name: "nbio",
fn: func(b *testing.B, key [4]byte, p []byte) {
keyb := key[:]
for i := 0; i < b.N; i++ {
nbioMaskBytes(p, keyb)
}
},
},
}

key := [4]byte{1, 2, 3, 4}

for _, size := range sizes {
p := make([]byte, size)
for _, fn := range fns {
b.Run(fn.name, func(b *testing.B) {
for _, size := range sizes {
p := make([]byte, size)

b.Run(strconv.Itoa(size), func(b *testing.B) {
for _, fn := range fns {
b.Run(fn.name, func(b *testing.B) {
b.Run(strconv.Itoa(size), func(b *testing.B) {
b.SetBytes(int64(size))

fn.fn(b, key, p)
Expand Down
2 changes: 2 additions & 0 deletions internal/thirdparty/go.mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ require (
github.com/gin-gonic/gin v1.9.1
github.com/gobwas/ws v1.3.0
github.com/gorilla/websocket v1.5.0
github.com/lesismal/nbio v1.3.18
nhooyr.io/websocket v0.0.0-00010101000000-000000000000
)

Expand All @@ -25,6 +26,7 @@ require (
github.com/json-iterator/go v1.1.12 // indirect
github.com/klauspost/cpuid/v2 v2.2.4 // indirect
github.com/leodido/go-urn v1.2.4 // indirect
github.com/lesismal/llib v1.1.12 // indirect
github.com/mattn/go-isatty v0.0.19 // indirect
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
github.com/modern-go/reflect2 v1.0.2 // indirect
Expand Down
36 changes: 36 additions & 0 deletions internal/thirdparty/go.sum
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,10 @@ github.com/klauspost/cpuid/v2 v2.2.4 h1:acbojRNwl3o09bUq+yDCtZFc1aiwaAAxtcn8YkZX
github.com/klauspost/cpuid/v2 v2.2.4/go.mod h1:RVVoqg1df56z8g3pUjL/3lE5UfnlrJX8tyFgg4nqhuY=
github.com/leodido/go-urn v1.2.4 h1:XlAE/cm/ms7TE/VMVoduSpNBoyc2dOxHs5MZSwAN63Q=
github.com/leodido/go-urn v1.2.4/go.mod h1:7ZrI8mTSeBSHl/UaRyKQW1qZeMgak41ANeCNaVckg+4=
github.com/lesismal/llib v1.1.12 h1:KJFB8bL02V+QGIvILEw/w7s6bKj9Ps9Px97MZP2EOk0=
github.com/lesismal/llib v1.1.12/go.mod h1:70tFXXe7P1FZ02AU9l8LgSOK7d7sRrpnkUr3rd3gKSg=
github.com/lesismal/nbio v1.3.18 h1:kmJZlxjQpVfuCPYcXdv0Biv9LHVViJZet5K99Xs3RAs=
github.com/lesismal/nbio v1.3.18/go.mod h1:KWlouFT5cgDdW5sMX8RsHASUMGniea9X0XIellZ0B38=
github.com/mattn/go-isatty v0.0.19 h1:JITubQf0MOLdlGRuRq+jtsDlekdYPia9ZFsB8h/APPA=
github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
Expand All @@ -67,19 +71,51 @@ github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS
github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08=
github.com/ugorji/go/codec v1.2.11 h1:BMaWp1Bb6fHwEtbplGBGJ498wD+LKlNSl25MjdZY4dU=
github.com/ugorji/go/codec v1.2.11/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg=
github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
golang.org/x/arch v0.0.0-20210923205945-b76863e36670/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8=
golang.org/x/arch v0.3.0 h1:02VY4/ZcO/gBOH6PUaoiptASxtXU10jazRCP865E97k=
golang.org/x/arch v0.3.0/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8=
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
golang.org/x/crypto v0.0.0-20210513122933-cd7d49e622d5/go.mod h1:P+XmwS30IXTQdn5tA2iutPOUgjI07+tq3H3K9MVA1s8=
golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
golang.org/x/crypto v0.6.0/go.mod h1:OFC/31mSvZgRz0V1QTNCzfAI1aIRzbiufJtkMIlEp58=
golang.org/x/crypto v0.9.0 h1:LF6fAI+IutBocDJ2OT0Q1g8plpYljMZ4+lty+dsqw3g=
golang.org/x/crypto v0.9.0/go.mod h1:yrmDGqONDYtNj3tH8X9dzUun2m2lzPa9ngI6/RUPGR0=
golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
golang.org/x/net v0.0.0-20210510120150-4163338589ed/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
golang.org/x/net v0.6.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs=
golang.org/x/net v0.10.0 h1:X2//UzNDwYmtCLn7To6G58Wr6f5ahEAQgKNzv9Y951M=
golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg=
golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.0.0-20220704084225-05e143d24a9e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.13.0 h1:Af8nKPmuFypiUBjVoU9V20FiaFXOcuZI21p0ycVYYGE=
golang.org/x/sys v0.13.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k=
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
golang.org/x/text v0.9.0 h1:2sjJmO8cDvYveuX97RDLsxlyUxLl+GHoLxBiRdHllBE=
golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
Expand Down
4 changes: 2 additions & 2 deletions mask.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,6 @@

package websocket

func mask(key uint32, b []byte) uint32 {
return maskGo(key, b)
func mask(b []byte, key uint32) uint32 {
return maskGo(b, key)
}
2 changes: 0 additions & 2 deletions mask_amd64.s
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,6 @@ TEXT ·maskAsm(SB), NOSPLIT, $0-28

CMPQ CX, $8
JL less_than_8
CMPQ CX, $64
JL less_than_64
CMPQ CX, $512
JLE sse
TESTQ $31, AX
Expand Down
2 changes: 1 addition & 1 deletion mask_arm64.s
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,8 @@
TEXT ·maskAsm(SB), NOSPLIT, $0-28
// R0 = b
// R1 = len
// R2 = uint64(key)<<32 | uint64(key)
// R3 = key (uint32)
// R2 = uint64(key)<<32 | uint64(key)
MOVD b_ptr+0(FP), R0
MOVD b_len+8(FP), R1
MOVWU key+16(FP), R3
Expand Down
2 changes: 1 addition & 1 deletion mask_asm.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

package websocket

func mask(key uint32, b []byte) uint32 {
func mask(b []byte, key uint32) uint32 {
if len(b) > 0 {
return maskAsm(&b[0], len(b), key)
}
Expand Down
4 changes: 2 additions & 2 deletions read.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -289,7 +289,7 @@ func (c *Conn) handleControl(ctx context.Context, h header) (err error) {
}

if h.masked {
mask(h.maskKey, b)
mask(b, h.maskKey)
}

switch h.opcode {
Expand Down Expand Up @@ -453,7 +453,7 @@ func (mr *msgReader) read(p []byte) (int, error) {
mr.payloadLength -= int64(n)

if !mr.c.client {
mr.maskKey = mask(mr.maskKey, p)
mr.maskKey = mask(p, mr.maskKey)
}

return n, nil
Expand Down
2 changes: 1 addition & 1 deletion write.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -365,7 +365,7 @@ func (c *Conn) writeFramePayload(p []byte) (n int, err error) {
return n, err
}

maskKey = mask(maskKey, c.writeBuf[i:c.bw.Buffered()])
maskKey = mask(c.writeBuf[i:c.bw.Buffered()], maskKey)

p = p[j:]
n += j
Expand Down

0 comments on commit 32d0aa1

Please sign in to comment.