393 lines
11 KiB
Go
393 lines
11 KiB
Go
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
|
|
}
|