devplacepy/nodes/client/ws_test.go

package main

import (
	"bytes"
	"context"
	"errors"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

func upgrade(w http.ResponseWriter, r *http.Request) (*Conn, error) {
	hijacker := w.(http.Hijacker)
	rwc, brw, err := hijacker.Hijack()
	if err != nil {
		return nil, err
	}
	answer := "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " +
		acceptKey(r.Header.Get("Sec-WebSocket-Key")) + "\r\n\r\n"
	if _, err := rwc.Write([]byte(answer)); err != nil {
		return nil, err
	}
	conn := newConn(rwc, brw.Reader, false)
	conn.maxMessage = MaxResultMessage
	return conn, nil
}

func wsServer(t *testing.T, handler func(*Conn)) *httptest.Server {
	t.Helper()
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrade(w, r)
		if err != nil {
			t.Errorf("upgrade: %v", err)
			return
		}
		defer conn.Close()
		handler(conn)
	}))
	t.Cleanup(server.Close)
	return server
}

func dialTest(t *testing.T, server *httptest.Server) *Conn {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	conn, err := Dial(ctx, "ws"+strings.TrimPrefix(server.URL, "http")+"/socket")
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	t.Cleanup(func() { conn.Close() })
	_ = conn.SetReadDeadline(time.Now().Add(10 * time.Second))
	return conn
}

func TestAcceptKeyVector(t *testing.T) {
	if got := acceptKey("dGhlIHNhbXBsZSBub25jZQ=="); got != "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=" {
		t.Fatalf("accept key %s", got)
	}
}

func TestRoundTripAndFragmentation(t *testing.T) {
	sizes := []int{0, 1, 125, 126, 65535, 65536, MaxWSFramePayload, MaxWSFramePayload + 1, 700 * 1024}
	server := wsServer(t, func(c *Conn) {
		for range sizes {
			message, err := c.ReadMessage()
			if err != nil {
				t.Errorf("server read: %v", err)
				return
			}
			if len(message) <= MaxInboundMessage {
				if err := c.WriteText(string(message)); err != nil {
					t.Errorf("server write: %v", err)
					return
				}
			} else if err := c.WriteText("big"); err != nil {
				return
			}
		}
	})
	conn := dialTest(t, server)
	for _, size := range sizes {
		text := strings.Repeat("a", size)
		if err := conn.WriteText(text); err != nil {
			t.Fatalf("write %d: %v", size, err)
		}
		got, err := conn.ReadMessage()
		if err != nil {
			t.Fatalf("read %d: %v", size, err)
		}
		want := text
		if size > MaxInboundMessage {
			want = "big"
		}
		if string(got) != want {
			t.Fatalf("size %d: got %d bytes", size, len(got))
		}
	}
}

func TestPingAnsweredWithPong(t *testing.T) {
	done := make(chan []byte, 1)
	server := wsServer(t, func(c *Conn) {
		_ = c.writeControl(opPing, []byte("probe"))
		_ = c.WriteText("after")
		_, op, payload, err := c.readFrame()
		if err != nil || op != opPong {
			t.Errorf("expected pong, got op %d err %v", op, err)
		}
		done <- payload
	})
	conn := dialTest(t, server)
	message, err := conn.ReadMessage()
	if err != nil || string(message) != "after" {
		t.Fatalf("message %q err %v", message, err)
	}
	if got := <-done; !bytes.Equal(got, []byte("probe")) {
		t.Fatalf("pong payload %q", got)
	}
}

func TestClientFramesAreMasked(t *testing.T) {
	masked := make(chan bool, 1)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrade(w, r)
		if err != nil {
			return
		}
		defer conn.Close()
		head := make([]byte, 2)
		if _, err := conn.br.Read(head); err != nil {
			return
		}
		masked <- head[1]&0x80 != 0
	}))
	defer server.Close()
	conn := dialTest(t, server)
	if err := conn.WriteText("x"); err != nil {
		t.Fatal(err)
	}
	if !<-masked {
		t.Fatal("client frame was not masked")
	}
}

func TestCloseCodeSurfaces(t *testing.T) {
	server := wsServer(t, func(c *Conn) { _ = c.WriteClose(CloseRetryOwner, "owner") })
	conn := dialTest(t, server)
	_, err := conn.ReadMessage()
	var closeErr *CloseError
	if !errors.As(err, &closeErr) || closeErr.Code != CloseRetryOwner || closeErr.Reason != "owner" {
		t.Fatalf("got %v", err)
	}
}

func TestOversizeFrameRejectedWithClose(t *testing.T) {
	server := wsServer(t, func(c *Conn) {
		_ = c.writeFrame(true, opText, []byte(strings.Repeat("b", MaxWSFramePayload+1)))
		_, _, _, _ = c.readFrame()
	})
	conn := dialTest(t, server)
	_, err := conn.ReadMessage()
	var closeErr *CloseError
	if !errors.As(err, &closeErr) || closeErr.Code != CloseTooLarge {
		t.Fatalf("got %v", err)
	}
}

func TestOversizeMessageAcrossFragmentsRejected(t *testing.T) {
	server := wsServer(t, func(c *Conn) {
		chunk := []byte(strings.Repeat("c", MaxWSFramePayload))
		_ = c.writeFrame(false, opText, chunk)
		_ = c.writeFrame(true, opContinuation, chunk)
		_, _, _, _ = c.readFrame()
	})
	conn := dialTest(t, server)
	_, err := conn.ReadMessage()
	var closeErr *CloseError
	if !errors.As(err, &closeErr) || closeErr.Code != CloseTooLarge {
		t.Fatalf("got %v", err)
	}
}

func TestBinaryAndBadUTF8Rejected(t *testing.T) {
	cases := []struct {
		name    string
		op      byte
		payload []byte
		code    int
	}{
		{"binary", opBinary, []byte("x"), CloseUnsupported},
		{"utf8", opText, []byte{0xff, 0xfe}, CloseInvalidPayload},
		{"continuation", opContinuation, []byte("x"), CloseProtocolError},
	}
	for _, item := range cases {
		server := wsServer(t, func(c *Conn) {
			_ = c.writeFrame(true, item.op, item.payload)
			_, _, _, _ = c.readFrame()
		})
		conn := dialTest(t, server)
		_, err := conn.ReadMessage()
		var closeErr *CloseError
		if !errors.As(err, &closeErr) || closeErr.Code != item.code {
			t.Fatalf("%s: got %v", item.name, err)
		}
	}
}

func TestHandshakeRefusal(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "no", http.StatusForbidden)
	}))
	defer server.Close()
	_, err := Dial(context.Background(), "ws"+strings.TrimPrefix(server.URL, "http"))
	if err == nil || !strings.Contains(err.Error(), "403") {
		t.Fatalf("got %v", err)
	}
}