327 lines
10 KiB
Go
327 lines
10 KiB
Go
package ws
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
func managedServer(t testing.TB, cfg ServerConfig, h Handlers) (*Server, *httptest.Server, string) {
|
|
t.Helper()
|
|
s, err := NewServer(cfg)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
httpServer := httptest.NewServer(s.Handler(h))
|
|
t.Cleanup(func() {
|
|
httpServer.Close()
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancel()
|
|
if err := s.Shutdown(ctx); err != nil {
|
|
t.Error(err)
|
|
}
|
|
})
|
|
return s, httpServer, "ws" + strings.TrimPrefix(httpServer.URL, "http")
|
|
}
|
|
|
|
func TestServerOriginAndProtocolPolicy(t *testing.T) {
|
|
_, host, endpoint := managedServer(t, ServerConfig{Subprotocols: []string{"test.v1"}, RequireSubprotocol: true}, Handlers{})
|
|
for _, tt := range []struct {
|
|
name, origin string
|
|
code int
|
|
}{
|
|
{"service", "", 101}, {"same", host.URL, 101}, {"foreign", "https://foreign.invalid", 403}, {"null", "null", 403}, {"path", host.URL + "/", 403}, {"comma", host.URL + ", https://evil.invalid", 403},
|
|
} {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
d := websocket.Dialer{Subprotocols: []string{"test.v1"}}
|
|
h := http.Header{}
|
|
if tt.origin != "" {
|
|
h.Set("Origin", tt.origin)
|
|
}
|
|
c, r, err := d.Dial(endpoint, h)
|
|
if c != nil {
|
|
defer c.Close()
|
|
}
|
|
if r == nil || r.StatusCode != tt.code {
|
|
t.Fatalf("status=%v err=%v", r, err)
|
|
}
|
|
if tt.code == 101 && (err != nil || c.Subprotocol() != "test.v1") {
|
|
t.Fatal("negotiation failed", err)
|
|
}
|
|
})
|
|
}
|
|
_, r, err := websocket.DefaultDialer.Dial(endpoint, nil)
|
|
if err == nil || r.StatusCode != 400 {
|
|
t.Fatal("missing subprotocol accepted")
|
|
}
|
|
d := websocket.Dialer{Subprotocols: []string{"test.v1"}}
|
|
_, r, err = d.Dial(endpoint, http.Header{"Origin": []string{host.URL, host.URL}})
|
|
if err == nil || r.StatusCode != 403 {
|
|
t.Fatal("duplicate origin accepted")
|
|
}
|
|
}
|
|
|
|
func TestServerAllowedOriginsAndAuthentication(t *testing.T) {
|
|
s, _, endpoint := managedServer(t, ServerConfig{AllowedOrigins: []string{"https://allowed.invalid"}, RequireOrigin: true, Prepare: func(ctx context.Context, r *http.Request) (Handlers, error) {
|
|
if r.Header.Get("Authorization") != "Bearer local-test" {
|
|
return Handlers{}, &RejectError{StatusCode: 403}
|
|
}
|
|
return Handlers{OnMessage: func(ctx context.Context, s *Session, m Message) error { return s.Send(ctx, m) }}, nil
|
|
}}, Handlers{})
|
|
for _, h := range []http.Header{nil, {"Origin": []string{"https://allowed.invalid"}}, {"Authorization": []string{"Bearer local-test"}}, {"Origin": []string{"https://wrong.invalid"}, "Authorization": []string{"Bearer local-test"}}} {
|
|
c, r, e := websocket.DefaultDialer.Dial(endpoint, h)
|
|
if c != nil {
|
|
c.Close()
|
|
}
|
|
if e == nil || r.StatusCode != 403 {
|
|
t.Fatal("unauthorized handshake accepted")
|
|
}
|
|
}
|
|
c, _, err := websocket.DefaultDialer.Dial(endpoint, http.Header{"Origin": []string{"https://allowed.invalid"}, "Authorization": []string{"Bearer local-test"}})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer c.Close()
|
|
if err = c.WriteMessage(websocket.TextMessage, []byte("echo")); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
c.SetReadDeadline(time.Now().Add(time.Second))
|
|
_, body, err := c.ReadMessage()
|
|
if err != nil || string(body) != "echo" {
|
|
t.Fatal(string(body), err)
|
|
}
|
|
if s.Stats().Accepted != 1 {
|
|
t.Fatal(s.Stats())
|
|
}
|
|
}
|
|
|
|
func TestServerStopDuringHandshake(t *testing.T) {
|
|
entered, release := make(chan struct{}), make(chan struct{})
|
|
s, _, endpoint := managedServer(t, ServerConfig{Prepare: func(context.Context, *http.Request) (Handlers, error) {
|
|
close(entered)
|
|
<-release
|
|
return Handlers{}, nil
|
|
}}, Handlers{})
|
|
dialDone := make(chan struct{})
|
|
go func() {
|
|
defer close(dialDone)
|
|
c, _, _ := websocket.DefaultDialer.Dial(endpoint, nil)
|
|
if c != nil {
|
|
defer c.Close()
|
|
c.ReadMessage()
|
|
}
|
|
}()
|
|
<-entered
|
|
_ = s.Stop(context.Background())
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
|
|
defer cancel()
|
|
if err := s.Drain(ctx); !errors.Is(err, context.DeadlineExceeded) {
|
|
t.Fatal("handshake was not tracked", err)
|
|
}
|
|
close(release)
|
|
select {
|
|
case <-dialDone:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("accepted connection escaped stop")
|
|
}
|
|
ctx2, cancel2 := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancel2()
|
|
if err := s.Drain(ctx2); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
func TestServerLimitsAndShutdownCallbacks(t *testing.T) {
|
|
entered, release := make(chan struct{}), make(chan struct{})
|
|
s, _, endpoint := managedServer(t, ServerConfig{MaxConnections: 1, Session: SessionConfig{CloseTimeout: 20 * time.Millisecond}}, Handlers{OnMessage: func(context.Context, *Session, Message) error { close(entered); <-release; return nil }})
|
|
c, _, err := websocket.DefaultDialer.Dial(endpoint, nil)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer c.Close()
|
|
_, r, err := websocket.DefaultDialer.Dial(endpoint, nil)
|
|
if err == nil || r.StatusCode != 503 {
|
|
t.Fatal("connection cap failed")
|
|
}
|
|
if err = c.WriteMessage(websocket.TextMessage, []byte("hold")); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
<-entered
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
|
|
defer cancel()
|
|
if err = s.Shutdown(ctx); !errors.Is(err, context.DeadlineExceeded) {
|
|
t.Fatal("blocked callback incorrectly drained", err)
|
|
}
|
|
close(release)
|
|
ctx2, cancel2 := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancel2()
|
|
if err = s.Drain(ctx2); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
_, r, err = websocket.DefaultDialer.Dial(endpoint, nil)
|
|
if err == nil || r.StatusCode != 503 {
|
|
t.Fatal("server admitted after shutdown")
|
|
}
|
|
}
|
|
|
|
func TestServerHandshakeRateAndPanic(t *testing.T) {
|
|
_, _, endpoint := managedServer(t, ServerConfig{HandshakeRate: 0.001, HandshakeBurst: 1, Prepare: func(context.Context, *http.Request) (Handlers, error) { panic("SECRET") }}, Handlers{})
|
|
_, r, e := websocket.DefaultDialer.Dial(endpoint, nil)
|
|
if e == nil || r.StatusCode != 401 {
|
|
t.Fatal("prepare panic not contained", e)
|
|
}
|
|
_, r, e = websocket.DefaultDialer.Dial(endpoint, nil)
|
|
if e == nil || r.StatusCode != 429 {
|
|
t.Fatal("rate cap not enforced", e)
|
|
}
|
|
}
|
|
|
|
func TestServerUnsupportedWriter(t *testing.T) {
|
|
s, err := NewServer(ServerConfig{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
r := httptest.NewRequest("GET", "http://localhost/ws", nil)
|
|
r.Header.Set("Upgrade", "websocket")
|
|
r.Header.Set("Connection", "Upgrade")
|
|
w := httptest.NewRecorder()
|
|
_, err = s.Accept(w, r, Handlers{})
|
|
if !errors.Is(err, ErrUnsupportedTransport) || w.Code != 501 {
|
|
t.Fatal(w.Code, err)
|
|
}
|
|
if err = s.Shutdown(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
func TestServerConcurrentShutdown(t *testing.T) {
|
|
s, _, endpoint := managedServer(t, ServerConfig{Session: SessionConfig{CloseTimeout: 10 * time.Millisecond}}, Handlers{})
|
|
var wg sync.WaitGroup
|
|
for i := 0; i < 40; i++ {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
c, _, _ := websocket.DefaultDialer.Dial(endpoint, nil)
|
|
if c != nil {
|
|
defer c.Close()
|
|
c.SetReadDeadline(time.Now().Add(time.Second))
|
|
c.ReadMessage()
|
|
}
|
|
}()
|
|
}
|
|
_ = s.Stop(context.Background())
|
|
wg.Wait()
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancel()
|
|
if err := s.Drain(ctx); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if stats := s.Stats(); stats.ActiveConnections != 0 || stats.Handshakes != 0 {
|
|
t.Fatal(stats)
|
|
}
|
|
}
|
|
|
|
func TestServerConfigValidation(t *testing.T) {
|
|
for _, cfg := range []ServerConfig{{AllowedOrigins: []string{"*"}}, {AllowedOrigins: []string{"https://user:secret@host"}}, {AllowedOrigins: []string{"https://host/path"}}, {HandshakeRate: -1}, {HandshakeBurst: 1}, {RequireSubprotocol: true}, {Subprotocols: []string{"bad\r\n"}}, {MaxConnections: -1}} {
|
|
if err := cfg.Validate(); err == nil {
|
|
t.Fatal("bad config accepted")
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestServerParallelConnections(t *testing.T) {
|
|
s, _, endpoint := managedServer(t, ServerConfig{MaxConnections: 256}, Handlers{OnMessage: func(ctx context.Context, s *Session, m Message) error { return s.Send(ctx, m) }})
|
|
var workers sync.WaitGroup
|
|
for i := 0; i < 128; i++ {
|
|
workers.Add(1)
|
|
go func(id int) {
|
|
defer workers.Done()
|
|
c, _, err := websocket.DefaultDialer.Dial(endpoint, nil)
|
|
if err != nil {
|
|
t.Error(err)
|
|
return
|
|
}
|
|
defer c.Close()
|
|
c.SetReadDeadline(time.Now().Add(15 * time.Second))
|
|
c.SetWriteDeadline(time.Now().Add(15 * time.Second))
|
|
for n := 0; n < 50; n++ {
|
|
value := fmt.Sprintf("%d:%d", id, n)
|
|
if err := c.WriteMessage(websocket.TextMessage, []byte(value)); err != nil {
|
|
t.Error(err)
|
|
return
|
|
}
|
|
_, got, err := c.ReadMessage()
|
|
if err != nil || string(got) != value {
|
|
t.Errorf("session isolation: %s %v", got, err)
|
|
return
|
|
}
|
|
}
|
|
}(i)
|
|
}
|
|
workers.Wait()
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancel()
|
|
if err := s.Shutdown(ctx); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if s.Stats().Accepted != 128 {
|
|
t.Fatal(s.Stats())
|
|
}
|
|
}
|
|
|
|
func TestServerTypedNilRejection(t *testing.T) {
|
|
_, _, endpoint := managedServer(t, ServerConfig{Prepare: func(context.Context, *http.Request) (Handlers, error) { var err *RejectError; return Handlers{}, err }}, Handlers{})
|
|
_, r, err := websocket.DefaultDialer.Dial(endpoint, nil)
|
|
if err == nil || r == nil || r.StatusCode != 401 {
|
|
t.Fatal(r, err)
|
|
}
|
|
}
|
|
|
|
func FuzzOrigin(f *testing.F) {
|
|
for _, v := range []string{"https://example.com", "http://localhost:8080", "null", "https://[::1]:443", "https://host:invalid", "https://evil@safe"} {
|
|
f.Add(v)
|
|
}
|
|
f.Fuzz(func(t *testing.T, v string) {
|
|
a, ok := canonicalOrigin(v)
|
|
if ok {
|
|
b, ok2 := canonicalOrigin(a)
|
|
if !ok2 || a != b {
|
|
t.Fatal("not canonical")
|
|
}
|
|
}
|
|
})
|
|
}
|
|
|
|
func BenchmarkServerRoundTrip(b *testing.B) {
|
|
for _, size := range []int{64, 1024, 64 << 10} {
|
|
b.Run(fmt.Sprint(size), func(b *testing.B) {
|
|
_, _, endpoint := managedServer(b, ServerConfig{}, Handlers{OnMessage: func(ctx context.Context, s *Session, m Message) error { return s.Send(ctx, m) }})
|
|
c, _, err := websocket.DefaultDialer.Dial(endpoint, nil)
|
|
if err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
defer c.Close()
|
|
payload := make([]byte, size)
|
|
b.SetBytes(int64(size) * 2)
|
|
b.ReportAllocs()
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
if err = c.WriteMessage(websocket.BinaryMessage, payload); err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
if _, _, err = c.ReadMessage(); err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|