139 lines
3.4 KiB
Go
139 lines
3.4 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/binary"
|
|
"testing"
|
|
)
|
|
|
|
func TestWriteFrameMasked(t *testing.T) {
|
|
// Verify that writeFrame produces a valid masked client frame
|
|
var buf bytes.Buffer
|
|
ws := &WSConn{conn: &nopCloser{readWriter: &buf}}
|
|
|
|
payload := []byte("hello")
|
|
if err := ws.writeFrame(opText, payload); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
frame := buf.Bytes()
|
|
|
|
// First byte: FIN + opcode
|
|
if frame[0] != 0x81 { // 0x80 (FIN) | 0x01 (text)
|
|
t.Errorf("first byte: got %#x, want 0x81", frame[0])
|
|
}
|
|
|
|
// Second byte: MASK + length
|
|
if frame[1] != 0x85 { // 0x80 (mask) | 5 (length)
|
|
t.Errorf("second byte: got %#x, want 0x85", frame[1])
|
|
}
|
|
|
|
// Mask key is bytes 2-5
|
|
maskKey := frame[2:6]
|
|
maskedPayload := frame[6:]
|
|
|
|
// Unmask and verify
|
|
for i := range maskedPayload {
|
|
maskedPayload[i] ^= maskKey[i%4]
|
|
}
|
|
if string(maskedPayload) != "hello" {
|
|
t.Errorf("unmasked payload: got %q, want %q", maskedPayload, "hello")
|
|
}
|
|
}
|
|
|
|
func TestWriteFrameExtendedLength(t *testing.T) {
|
|
// Test 16-bit extended length (126-65535 bytes)
|
|
var buf bytes.Buffer
|
|
ws := &WSConn{conn: &nopCloser{readWriter: &buf}}
|
|
|
|
payload := make([]byte, 300) // > 125, uses 2-byte extended
|
|
if err := ws.writeFrame(opBinary, payload); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
frame := buf.Bytes()
|
|
if frame[1]&0x7F != 126 {
|
|
t.Errorf("expected 126 length marker, got %d", frame[1]&0x7F)
|
|
}
|
|
extLen := binary.BigEndian.Uint16(frame[2:4])
|
|
if extLen != 300 {
|
|
t.Errorf("extended length: got %d, want 300", extLen)
|
|
}
|
|
}
|
|
|
|
func TestReadFrame(t *testing.T) {
|
|
// Build an unmasked server frame
|
|
var buf bytes.Buffer
|
|
payload := []byte("world")
|
|
buf.WriteByte(0x81) // FIN + text
|
|
buf.WriteByte(byte(len(payload)))
|
|
buf.Write(payload)
|
|
|
|
ws := &WSConn{conn: &nopCloser{readWriter: &buf}}
|
|
opcode, data, err := ws.readFrame()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if opcode != opText {
|
|
t.Errorf("opcode: got %d, want %d", opcode, opText)
|
|
}
|
|
if string(data) != "world" {
|
|
t.Errorf("data: got %q, want %q", data, "world")
|
|
}
|
|
}
|
|
|
|
func TestReadFramePingPong(t *testing.T) {
|
|
// Server sends a ping, ReadMessage should auto-respond with pong and continue
|
|
var buf bytes.Buffer
|
|
|
|
// Ping frame
|
|
buf.WriteByte(0x80 | byte(opPing))
|
|
buf.WriteByte(0) // no payload
|
|
|
|
// Then a text frame
|
|
text := []byte("data")
|
|
buf.WriteByte(0x81)
|
|
buf.WriteByte(byte(len(text)))
|
|
buf.Write(text)
|
|
|
|
rw := &captureWriter{Reader: &buf}
|
|
ws := &WSConn{conn: &nopCloser{readWriter: rw}}
|
|
|
|
data, _, err := ws.ReadMessage()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if string(data) != "data" {
|
|
t.Errorf("got %q, want %q", data, "data")
|
|
}
|
|
|
|
// Verify pong was written
|
|
if len(rw.written) == 0 {
|
|
t.Error("expected pong frame to be written")
|
|
}
|
|
}
|
|
|
|
// nopCloser wraps a ReadWriter with a no-op Close.
|
|
type nopCloser struct {
|
|
readWriter interface {
|
|
Read([]byte) (int, error)
|
|
Write([]byte) (int, error)
|
|
}
|
|
}
|
|
|
|
func (n *nopCloser) Read(p []byte) (int, error) { return n.readWriter.Read(p) }
|
|
func (n *nopCloser) Write(p []byte) (int, error) { return n.readWriter.Write(p) }
|
|
func (n *nopCloser) Close() error { return nil }
|
|
|
|
// captureWriter captures written data while reading from a separate Reader.
|
|
type captureWriter struct {
|
|
Reader *bytes.Buffer
|
|
written []byte
|
|
}
|
|
|
|
func (c *captureWriter) Read(p []byte) (int, error) { return c.Reader.Read(p) }
|
|
func (c *captureWriter) Write(p []byte) (int, error) {
|
|
c.written = append(c.written, p...)
|
|
return len(p), nil
|
|
}
|
|
func (c *captureWriter) Close() error { return nil }
|