From 82f5698802ea086a021df1bf789ef866153c53cf Mon Sep 17 00:00:00 2001
From: Graham McIntire <graham@mcintire.me>
Date: Thu, 12 Feb 2026 09:54:31 -0600
Subject: [PATCH] combine WebSocket frame header and payload into single write

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.
---
 websocket.go      | 48 ++++++++++++++++++-----------------------------
 websocket_test.go | 24 ------------------------
 2 files changed, 18 insertions(+), 54 deletions(-)

diff --git a/websocket.go b/websocket.go
index 7bcac82..cabdae2 100644
--- a/websocket.go
+++ b/websocket.go
@@ -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 {
-		return err
-	}
-	if length > 0 {
-		if _, err := ws.conn.Write(masked); err != nil {
-			return err
-		}
-	}
-	return nil
+	_, err := ws.conn.Write(buf)
+	return err
 }
diff --git a/websocket_test.go b/websocket_test.go
index 3a68f04..10d9605 100644
--- a/websocket_test.go
+++ b/websocket_test.go
@@ -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