101 lines
2.8 KiB
Go
101 lines
2.8 KiB
Go
package ws
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
func TestSessionConfigUpperBounds(t *testing.T) {
|
|
for _, c := range []SessionConfig{{MaxMessageBytes: 1<<30 + 1}, {ReadQueueSize: 65537}, {WriteQueueSize: 65537}, {ReadQueueBytes: 1<<30 + 1}, {WriteQueueBytes: 1<<30 + 1}, {CloseTimeout: 25 * time.Hour}} {
|
|
if c.Validate() == nil {
|
|
t.Fatal("oversized config accepted")
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestClientReconnectsRemoteGoingAway(t *testing.T) {
|
|
var handshakes atomic.Int32
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
|
|
if err != nil {
|
|
return
|
|
}
|
|
defer conn.Close()
|
|
if handshakes.Add(1) == 1 {
|
|
_ = conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(1001, ""), time.Now().Add(time.Second))
|
|
}
|
|
for {
|
|
if _, _, err = conn.ReadMessage(); err != nil {
|
|
return
|
|
}
|
|
}
|
|
}))
|
|
defer server.Close()
|
|
client, err := NewClient(ClientConfig{ProxyMode: "direct", Reconnect: ReconnectConfig{Enabled: true, MinBackoff: time.Millisecond, MaxBackoff: time.Millisecond}})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer func() {
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancel()
|
|
if err := client.Shutdown(ctx); err != nil {
|
|
t.Error(err)
|
|
}
|
|
}()
|
|
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
|
defer cancel()
|
|
connected := make(chan struct{}, 1)
|
|
done := make(chan error, 1)
|
|
go func() {
|
|
done <- client.Run(ctx, "ws"+strings.TrimPrefix(server.URL, "http"), Handlers{OnOpen: func(context.Context, *Session) error {
|
|
if handshakes.Load() >= 2 {
|
|
connected <- struct{}{}
|
|
}
|
|
return nil
|
|
}})
|
|
}()
|
|
select {
|
|
case <-connected:
|
|
case err := <-done:
|
|
t.Fatalf("Run ended before reconnect: %v", err)
|
|
case <-ctx.Done():
|
|
t.Fatal(ctx.Err())
|
|
}
|
|
_ = client.Close()
|
|
select {
|
|
case err := <-done:
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
case <-ctx.Done():
|
|
t.Fatal(ctx.Err())
|
|
}
|
|
if stats := client.Stats(); stats.DialAttempts != 2 || stats.Reconnects != 1 || stats.HandshakeFailures != 0 {
|
|
t.Fatal(stats)
|
|
}
|
|
}
|
|
|
|
func TestClientHandshakeStats(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Error(w, "denied", http.StatusUnauthorized) }))
|
|
defer server.Close()
|
|
client, err := NewClient(ClientConfig{ProxyMode: "direct"})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer client.Close()
|
|
_, err = client.Dial(context.Background(), "ws"+strings.TrimPrefix(server.URL, "http"), Handlers{})
|
|
if err == nil {
|
|
t.Fatal("expected rejection")
|
|
}
|
|
if stats := client.Stats(); stats.DialAttempts != 1 || stats.HandshakeFailures != 1 || stats.Reconnects != 0 {
|
|
t.Fatal(stats)
|
|
}
|
|
}
|