graham/towerops-agent
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

combine WebSocket frame header and payload into single write

Browse source 
Build the entire frame (header + mask key + masked payload) in one
buffer and write it with a single syscall instead of two separate
writes. Removes the failOnNthWrite test that tested the old
two-write implementation detail.
This commit is contained in:
Graham McIntire 2026-02-12 09:54:31 -06:00
parent fbaf7a07dc
commit 82f5698802
No known key found for this signature in database
2 changed files with 18 additions and 54 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

46
websocket.go
View file

@ -180,44 +180,32 @@ func (ws *WSConn) writeFrame(opcode int, payload []byte) error {
defer ws.mu.Unlock()
length := len(payload)
// Max header: 2 + 8 + 4 (mask) = 14 bytes
header := make([]byte, 2, 14)
header[0] = 0x80 | byte(opcode) // FIN + opcode
header[1] = 0x80 // masked (client must mask)
// Single buffer: max header (2+8+4=14) + payload
buf := make([]byte, 0, 14+length)
// Header
buf = append(buf, 0x80|byte(opcode)) // FIN + opcode
switch {
case length <= 125:
header[1] |= byte(length)
buf = append(buf, 0x80|byte(length))
case length <= 65535:
header[1] |= 126
ext := make([]byte, 2)
binary.BigEndian.PutUint16(ext, uint16(length))
header = append(header, ext...)
buf = append(buf, 0x80|126, byte(length>>8), byte(length))
default:
header[1] |= 127
ext := make([]byte, 8)
binary.BigEndian.PutUint64(ext, uint64(length))
header = append(header, ext...)
buf = append(buf, 0x80|127,
byte(length>>56), byte(length>>48), byte(length>>40), byte(length>>32),
byte(length>>24), byte(length>>16), byte(length>>8), byte(length))
}
// Generate mask key
maskKey := make([]byte, 4)
_, _ = rand.Read(maskKey)
header = append(header, maskKey...)
// Mask key
var maskKey [4]byte
_, _ = rand.Read(maskKey[:])
buf = append(buf, maskKey[:]...)
// Mask payload
masked := make([]byte, length)
for i := range payload {
masked[i] = payload[i] ^ maskKey[i%4]
// Masked payload
for i, b := range payload {
buf = append(buf, b^maskKey[i%4])
}
if _, err := ws.conn.Write(header); err != nil {
_, err := ws.conn.Write(buf)
return err
}
if length > 0 {
if _, err := ws.conn.Write(masked); err != nil {
return err
}
}
return nil
}

24
websocket_test.go
View file

@ -283,30 +283,6 @@ func TestWriteFrameHeaderError(t *testing.T) {
}
}
func TestWriteFramePayloadError(t *testing.T) {
// Writer that succeeds for header but fails for payload
ws := testWSConn(&failOnNthWrite{failAfter: 1})
err := ws.writeFrame(opText, []byte("hello"))
if err == nil {
t.Error("expected write error on payload")
}
}
// failOnNthWrite fails after N successful writes.
type failOnNthWrite struct {
failAfter int
count int
}
func (f *failOnNthWrite) Read(p []byte) (int, error) { return 0, io.EOF }
func (f *failOnNthWrite) Write(p []byte) (int, error) {
f.count++
if f.count > f.failAfter {
return 0, fmt.Errorf("write failed")
}
return len(p), nil
}
func (f *failOnNthWrite) Close() error { return nil }
func TestReadMessageError(t *testing.T) {
// Empty buffer causes immediate EOF on readFrame