feat(ws): release generic WebSocket client and server v0.31.0-preview
This commit is contained in:
392
ws/session_test.go
Normal file
392
ws/session_test.go
Normal file
@@ -0,0 +1,392 @@
|
||||
package ws
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
func sessionPair(t testing.TB, cfg SessionConfig, h Handlers, start, compression bool) (*Session, *websocket.Conn) {
|
||||
t.Helper()
|
||||
if err := cfg.Validate(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
accepted := make(chan *Session, 1)
|
||||
up := websocket.Upgrader{EnableCompression: compression}
|
||||
host := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
conn, err := up.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
s := newSession(conn, cfg, h, context.Background())
|
||||
accepted <- s
|
||||
if start {
|
||||
s.start()
|
||||
}
|
||||
}))
|
||||
dialer := websocket.Dialer{EnableCompression: compression}
|
||||
peer, _, err := dialer.Dial("ws"+strings.TrimPrefix(host.URL, "http"), nil)
|
||||
if err != nil {
|
||||
host.Close()
|
||||
t.Fatal(err)
|
||||
}
|
||||
s := <-accepted
|
||||
t.Cleanup(func() {
|
||||
s.Abort()
|
||||
s.start()
|
||||
_ = peer.Close()
|
||||
host.Close()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
_ = s.Wait(ctx)
|
||||
})
|
||||
return s, peer
|
||||
}
|
||||
|
||||
func sessionTestConfig() SessionConfig {
|
||||
return SessionConfig{DisableHeartbeat: true, CloseTimeout: 100 * time.Millisecond, WriteTimeout: time.Second}
|
||||
}
|
||||
func waitSession(t *testing.T, s *Session) error {
|
||||
t.Helper()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
err := s.Wait(ctx)
|
||||
if errors.Is(err, context.DeadlineExceeded) && ctx.Err() != nil {
|
||||
t.Fatal("session did not finish")
|
||||
}
|
||||
return err
|
||||
}
|
||||
func peerClose(t *testing.T, peer *websocket.Conn) {
|
||||
t.Helper()
|
||||
_ = peer.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(1000, ""), time.Now().Add(time.Second))
|
||||
_ = peer.SetReadDeadline(time.Now().Add(time.Second))
|
||||
_, _, _ = peer.ReadMessage()
|
||||
}
|
||||
|
||||
func TestSessionEchoAndCompletion(t *testing.T) {
|
||||
var opened, closed atomic.Bool
|
||||
s, peer := sessionPair(t, sessionTestConfig(), Handlers{
|
||||
OnOpen: func(context.Context, *Session) error { opened.Store(true); return nil },
|
||||
OnMessage: func(ctx context.Context, s *Session, msg Message) error {
|
||||
if !opened.Load() {
|
||||
return errors.New("message before open")
|
||||
}
|
||||
return s.Send(ctx, msg)
|
||||
},
|
||||
OnClose: func(*Session, error) { closed.Store(true) },
|
||||
}, true, false)
|
||||
for _, kind := range []int{websocket.TextMessage, websocket.BinaryMessage} {
|
||||
if err := peer.WriteMessage(kind, []byte("hello")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, data, err := peer.ReadMessage()
|
||||
if err != nil || got != kind || string(data) != "hello" {
|
||||
t.Fatalf("echo: %d %q %v", got, data, err)
|
||||
}
|
||||
}
|
||||
peerClose(t, peer)
|
||||
if err := waitSession(t, s); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !closed.Load() {
|
||||
t.Fatal("Done preceded OnClose")
|
||||
}
|
||||
stats := s.Stats()
|
||||
if stats.MessagesReceived != 2 || stats.MessagesSent != 2 || stats.QueuedMessages != 0 || stats.ReadQueuedMessages != 0 {
|
||||
t.Fatalf("stats: %+v", stats)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionQueueBoundsCopiesAndGracefulDrain(t *testing.T) {
|
||||
cfg := sessionTestConfig()
|
||||
cfg.WriteQueueSize = 2
|
||||
cfg.WriteQueueBytes = 4
|
||||
s, peer := sessionPair(t, cfg, Handlers{}, false, false)
|
||||
data := []byte("ab")
|
||||
if err := s.TrySend(Message{Type: Text, Data: data}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
data[0] = 'z'
|
||||
if err := s.TrySend(Message{Type: Text, Data: []byte("cd")}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.TrySend(Message{Type: Text, Data: []byte("e")}); !errors.Is(err, ErrQueueFull) {
|
||||
t.Fatalf("overflow: %v", err)
|
||||
}
|
||||
if err := s.Close(1000, ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.TrySend(Message{Type: Text}); !errors.Is(err, ErrClosed) {
|
||||
t.Fatalf("closing send: %v", err)
|
||||
}
|
||||
s.start()
|
||||
for _, want := range []string{"ab", "cd"} {
|
||||
_, got, err := peer.ReadMessage()
|
||||
if err != nil || string(got) != want {
|
||||
t.Fatalf("drain: %q %v", got, err)
|
||||
}
|
||||
}
|
||||
_, _, err := peer.ReadMessage()
|
||||
var ce *websocket.CloseError
|
||||
if !errors.As(err, &ce) || ce.Code != 1000 {
|
||||
t.Fatalf("close: %v", err)
|
||||
}
|
||||
if err := waitSession(t, s); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionCanceledQueueDoesNotCloseConnection(t *testing.T) {
|
||||
s, peer := sessionPair(t, sessionTestConfig(), Handlers{}, false, false)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
result := make(chan error, 1)
|
||||
go func() { result <- s.Send(ctx, Message{Type: Text, Data: []byte("skip")}) }()
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for s.Stats().QueuedMessages == 0 {
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatal("send never admitted")
|
||||
}
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
cancel()
|
||||
if err := <-result; !errors.Is(err, context.Canceled) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.TrySend(Message{Type: Text, Data: []byte("next")}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s.start()
|
||||
_, data, err := peer.ReadMessage()
|
||||
if err != nil || string(data) != "next" {
|
||||
t.Fatalf("got %q %v", data, err)
|
||||
}
|
||||
peerClose(t, peer)
|
||||
if err := waitSession(t, s); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionConcurrentSendAndJSON(t *testing.T) {
|
||||
s, peer := sessionPair(t, sessionTestConfig(), Handlers{}, true, false)
|
||||
const count = 24
|
||||
results := make(chan error, count)
|
||||
for i := 0; i < count; i++ {
|
||||
go func(i int) { results <- s.SendJSON(context.Background(), i) }(i)
|
||||
}
|
||||
seen := make(map[string]bool)
|
||||
for i := 0; i < count; i++ {
|
||||
_, data, err := peer.ReadMessage()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if seen[string(data)] {
|
||||
t.Fatal("duplicate write")
|
||||
}
|
||||
seen[string(data)] = true
|
||||
}
|
||||
for i := 0; i < count; i++ {
|
||||
if err := <-results; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
peerClose(t, peer)
|
||||
if err := waitSession(t, s); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionBoundedDispatcherDoesNotBlockPing(t *testing.T) {
|
||||
entered := make(chan struct{})
|
||||
release := make(chan struct{})
|
||||
var once sync.Once
|
||||
s, peer := sessionPair(t, sessionTestConfig(), Handlers{OnMessage: func(ctx context.Context, _ *Session, _ Message) error {
|
||||
once.Do(func() { close(entered) })
|
||||
select {
|
||||
case <-release:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
return nil
|
||||
}}, true, false)
|
||||
t.Cleanup(func() { close(release) })
|
||||
if err := peer.WriteMessage(websocket.TextMessage, []byte("block callback")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
<-entered
|
||||
pong := make(chan struct{}, 1)
|
||||
peer.SetPongHandler(func(data string) error {
|
||||
if data == "probe" {
|
||||
pong <- struct{}{}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
readDone := make(chan struct{})
|
||||
go func() { defer close(readDone); _, _, _ = peer.ReadMessage() }()
|
||||
if err := peer.WriteControl(websocket.PingMessage, []byte("probe"), time.Now().Add(time.Second)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
select {
|
||||
case <-pong:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("blocked callback starved control frames")
|
||||
}
|
||||
_ = s.Close(1000, "")
|
||||
<-readDone
|
||||
if err := waitSession(t, s); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionInboundQueueFullCancelsCallback(t *testing.T) {
|
||||
cfg := sessionTestConfig()
|
||||
cfg.ReadQueueSize = 1
|
||||
cfg.ReadQueueBytes = 8
|
||||
entered := make(chan struct{})
|
||||
s, peer := sessionPair(t, cfg, Handlers{OnMessage: func(ctx context.Context, _ *Session, _ Message) error { close(entered); <-ctx.Done(); return nil }}, true, false)
|
||||
_ = peer.WriteMessage(websocket.TextMessage, []byte("one"))
|
||||
<-entered
|
||||
_ = peer.WriteMessage(websocket.TextMessage, []byte("two"))
|
||||
_, _, err := peer.ReadMessage()
|
||||
var ce *websocket.CloseError
|
||||
if !errors.As(err, &ce) || ce.Code != 1008 {
|
||||
t.Fatalf("close: %v", err)
|
||||
}
|
||||
if err := waitSession(t, s); !errors.Is(err, ErrQueueFull) {
|
||||
t.Fatalf("session: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionCompressedReadLimitAndInvalidUTF8(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
data []byte
|
||||
code int
|
||||
cause error
|
||||
}{
|
||||
{"compressed", []byte(strings.Repeat("x", 4096)), 1009, ErrMessageTooLarge},
|
||||
{"utf8", []byte{0xff}, 1007, ErrInvalidMessage},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
cfg := sessionTestConfig()
|
||||
cfg.MaxMessageBytes = 128
|
||||
s, peer := sessionPair(t, cfg, Handlers{}, true, true)
|
||||
if err := peer.WriteMessage(websocket.TextMessage, tc.data); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, _, err := peer.ReadMessage()
|
||||
var ce *websocket.CloseError
|
||||
if !errors.As(err, &ce) || ce.Code != tc.code {
|
||||
t.Fatalf("close: %v", err)
|
||||
}
|
||||
if err := waitSession(t, s); !errors.Is(err, tc.cause) {
|
||||
t.Fatalf("limit: %v", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionCallbackPanicAndCloseBarrier(t *testing.T) {
|
||||
closeEntered := make(chan struct{})
|
||||
release := make(chan struct{})
|
||||
s, peer := sessionPair(t, sessionTestConfig(), Handlers{
|
||||
OnMessage: func(context.Context, *Session, Message) error { panic("secret peer data") },
|
||||
OnClose: func(*Session, error) { close(closeEntered); <-release },
|
||||
}, true, false)
|
||||
_ = peer.WriteMessage(websocket.TextMessage, []byte("panic"))
|
||||
_, _, _ = peer.ReadMessage()
|
||||
<-closeEntered
|
||||
select {
|
||||
case <-s.Done():
|
||||
t.Fatal("Done ignored running OnClose")
|
||||
default:
|
||||
}
|
||||
close(release)
|
||||
if err := waitSession(t, s); err == nil || strings.Contains(err.Error(), "secret") {
|
||||
t.Fatalf("panic error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionHeartbeatDetectsUnresponsivePeer(t *testing.T) {
|
||||
cfg := SessionConfig{PingInterval: 10 * time.Millisecond, PongTimeout: 30 * time.Millisecond, CloseTimeout: 30 * time.Millisecond}
|
||||
s, _ := sessionPair(t, cfg, Handlers{}, true, false)
|
||||
if err := waitSession(t, s); err == nil {
|
||||
t.Fatal("peer without reader remained healthy")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionWriteCancellationStopsConnection(t *testing.T) {
|
||||
cfg := sessionTestConfig()
|
||||
cfg.MaxMessageBytes = 16 << 20
|
||||
cfg.WriteQueueBytes = 16 << 20
|
||||
s, _ := sessionPair(t, cfg, Handlers{}, true, false)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
||||
defer cancel()
|
||||
err := s.Send(ctx, Message{Type: Binary, Data: make([]byte, 16<<20)})
|
||||
if !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Fatalf("send: %v", err)
|
||||
}
|
||||
if err := waitSession(t, s); !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Fatalf("write cancellation: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionValidationAndRedactedErrors(t *testing.T) {
|
||||
for _, cfg := range []SessionConfig{{MaxMessageBytes: -1}, {ReadQueueSize: -1}, {WriteQueueBytes: -1}, {WriteTimeout: -1}, {PingInterval: time.Second, PongTimeout: time.Second}} {
|
||||
if cfg.Validate() == nil {
|
||||
t.Fatalf("invalid config accepted: %+v", cfg)
|
||||
}
|
||||
}
|
||||
s, _ := sessionPair(t, sessionTestConfig(), Handlers{}, false, false)
|
||||
for _, msg := range []Message{{Type: 9}, {Type: Text, Data: []byte{0xff}}} {
|
||||
if !errors.Is(s.TrySend(msg), ErrInvalidMessage) {
|
||||
t.Fatal("invalid message accepted")
|
||||
}
|
||||
}
|
||||
if !errors.Is(s.TrySend(Message{Type: Binary, Data: make([]byte, (1<<20)+1)}), ErrMessageTooLarge) {
|
||||
t.Fatal("large outbound accepted")
|
||||
}
|
||||
for _, code := range []int{1005, 1006, 1015, 5000} {
|
||||
if s.Close(code, "") == nil {
|
||||
t.Fatalf("invalid close %d", code)
|
||||
}
|
||||
}
|
||||
secret := fmt.Errorf("https://user:password@example.test?token=secret")
|
||||
for _, err := range []error{&OpError{Op: "dial", Cause: secret}, &HandshakeError{StatusCode: 401, Cause: secret}, &CloseError{Code: 1008, Cause: secret}} {
|
||||
if strings.Contains(err.Error(), "secret") || !errors.Is(err, secret) {
|
||||
t.Fatalf("error wrapping: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkSessionSend(b *testing.B) {
|
||||
s, peer := sessionPair(b, sessionTestConfig(), Handlers{}, true, false)
|
||||
readerDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(readerDone)
|
||||
for {
|
||||
if _, _, err := peer.ReadMessage(); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
msg := Message{Type: Binary, Data: make([]byte, 256)}
|
||||
b.ReportAllocs()
|
||||
b.SetBytes(256)
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
if err := s.Send(context.Background(), msg); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
b.StopTimer()
|
||||
_ = s.Close(1000, "")
|
||||
<-readerDone
|
||||
}
|
||||
Reference in New Issue
Block a user