feat(ws): release generic WebSocket client and server v0.31.0-preview
This commit is contained in:
300
ws/observe/observe_test.go
Normal file
300
ws/observe/observe_test.go
Normal file
@@ -0,0 +1,300 @@
|
||||
package observe
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.yuntuo.hk/k2safe/go-omni/logx"
|
||||
"git.yuntuo.hk/k2safe/go-omni/ws"
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
func echo(w http.ResponseWriter, r *http.Request) {
|
||||
conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Close()
|
||||
for {
|
||||
kind, data, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if err = conn.WriteMessage(kind, data); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
func testClient(t *testing.T) *ws.Client {
|
||||
t.Helper()
|
||||
c, err := ws.NewClient(ws.ClientConfig{ProxyMode: "direct", Session: ws.SessionConfig{CloseTimeout: 50 * time.Millisecond}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
if err := c.Shutdown(ctx); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
})
|
||||
return c
|
||||
}
|
||||
func endpoint(s *httptest.Server) string { return "ws" + strings.TrimPrefix(s.URL, "http") }
|
||||
func wait(t *testing.T, ch <-chan struct{}) {
|
||||
t.Helper()
|
||||
select {
|
||||
case <-ch:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("timeout")
|
||||
}
|
||||
}
|
||||
|
||||
func TestObserverConcurrentSessionsAndTotals(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(echo))
|
||||
defer server.Close()
|
||||
client := testClient(t)
|
||||
o := New(nil)
|
||||
const connections = 8
|
||||
const count = 10
|
||||
var workers sync.WaitGroup
|
||||
stopSamples := make(chan struct{})
|
||||
samplesDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(samplesDone)
|
||||
for {
|
||||
select {
|
||||
case <-stopSamples:
|
||||
return
|
||||
default:
|
||||
_ = o.Snapshot()
|
||||
}
|
||||
}
|
||||
}()
|
||||
for i := 0; i < connections; i++ {
|
||||
workers.Add(1)
|
||||
go func() {
|
||||
defer workers.Done()
|
||||
opened := make(chan struct{})
|
||||
received := make(chan struct{}, count)
|
||||
s, err := client.Dial(context.Background(), endpoint(server), o.Wrap(ws.Handlers{OnOpen: func(context.Context, *ws.Session) error { close(opened); return nil }, OnMessage: func(context.Context, *ws.Session, ws.Message) error { received <- struct{}{}; return nil }}))
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
return
|
||||
}
|
||||
wait(t, opened)
|
||||
for j := 0; j < count; j++ {
|
||||
if err = s.Send(context.Background(), ws.Message{Type: ws.Text, Data: []byte("test")}); err != nil {
|
||||
t.Error(err)
|
||||
return
|
||||
}
|
||||
wait(t, received)
|
||||
}
|
||||
_ = s.Close(1000, "")
|
||||
wait(t, s.Done())
|
||||
}()
|
||||
}
|
||||
workers.Wait()
|
||||
close(stopSamples)
|
||||
wait(t, samplesDone)
|
||||
got := o.Snapshot()
|
||||
if got.Opened != connections || got.Closed != connections || got.Active != 0 || got.MessagesSent != connections*count || got.MessagesReceived != connections*count || got.BytesSent != 4*connections*count || got.BytesReceived != 4*connections*count {
|
||||
t.Fatalf("snapshot: %+v", got)
|
||||
}
|
||||
if again := o.Snapshot(); again != got {
|
||||
t.Fatalf("closed totals changed: %+v -> %+v", got, again)
|
||||
}
|
||||
}
|
||||
|
||||
func TestObserverCloseCallbackBoundaryAndPanic(t *testing.T) {
|
||||
for _, panics := range []bool{false, true} {
|
||||
t.Run(fmt.Sprint(panics), func(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(echo))
|
||||
defer server.Close()
|
||||
client := testClient(t)
|
||||
o := New(nil)
|
||||
opened := make(chan struct{})
|
||||
closing := make(chan struct{})
|
||||
release := make(chan struct{})
|
||||
s, err := client.Dial(context.Background(), endpoint(server), o.Wrap(ws.Handlers{OnOpen: func(context.Context, *ws.Session) error { close(opened); return nil }, OnClose: func(*ws.Session, error) {
|
||||
close(closing)
|
||||
<-release
|
||||
if panics {
|
||||
panic("SECRET CALLBACK PANIC")
|
||||
}
|
||||
}}))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
wait(t, opened)
|
||||
_ = s.Close(1000, "")
|
||||
wait(t, closing)
|
||||
if got := o.Snapshot(); got.Active != 1 || got.Closed != 0 {
|
||||
t.Fatalf("closed before callback finished: %+v", got)
|
||||
}
|
||||
select {
|
||||
case <-s.Done():
|
||||
t.Fatal("Done closed before user callback")
|
||||
default:
|
||||
}
|
||||
close(release)
|
||||
wait(t, s.Done())
|
||||
got := o.Snapshot()
|
||||
if got.Active != 0 || got.Closed != 1 {
|
||||
t.Fatalf("callback did not release observation: %+v", got)
|
||||
}
|
||||
if panics && got.CloseCallbackPanics != 1 {
|
||||
t.Fatalf("panic not counted: %+v", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestObserverPayloadFreeLogsAndIndependentDurations(t *testing.T) {
|
||||
var output bytes.Buffer
|
||||
logger, err := logx.New(logx.Config{Writer: &output, Format: "json", Level: "info"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer logger.Close()
|
||||
o := New(logger)
|
||||
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 _, _, err = conn.ReadMessage(); err == nil {
|
||||
_ = conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(1008, "SECRET CLOSE REASON"), time.Now().Add(time.Second))
|
||||
_, _, _ = conn.ReadMessage()
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
client := testClient(t)
|
||||
var sessions []*ws.Session
|
||||
for i := 0; i < 2; i++ {
|
||||
opened := make(chan struct{})
|
||||
s, err := client.Dial(context.Background(), endpoint(server)+"/?token=SECRET-URL", o.Wrap(ws.Handlers{OnOpen: func(context.Context, *ws.Session) error { close(opened); return nil }}))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
wait(t, opened)
|
||||
sessions = append(sessions, s)
|
||||
if i == 0 {
|
||||
time.Sleep(40 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
for _, s := range sessions {
|
||||
if err = s.Send(context.Background(), ws.Message{Type: ws.Text, Data: []byte("SECRET MESSAGE")}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
wait(t, s.Done())
|
||||
}
|
||||
if strings.Contains(output.String(), "SECRET") {
|
||||
t.Fatalf("log leaked source data: %s", output.String())
|
||||
}
|
||||
var durations []float64
|
||||
for _, line := range strings.Split(strings.TrimSpace(output.String()), "\n") {
|
||||
var row map[string]any
|
||||
if err = json.Unmarshal([]byte(line), &row); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for key := range row {
|
||||
switch key {
|
||||
case "time", "level", "msg", "sessionId", "phase", "durationMs", "closeCode":
|
||||
default:
|
||||
t.Fatalf("unexpected log field %q", key)
|
||||
}
|
||||
}
|
||||
if row["msg"] == "ws.close" {
|
||||
if row["closeCode"] != float64(1008) {
|
||||
t.Fatalf("close code: %#v", row)
|
||||
}
|
||||
durations = append(durations, row["durationMs"].(float64))
|
||||
}
|
||||
}
|
||||
if len(durations) != 2 || durations[0] < 30 {
|
||||
t.Fatalf("timers not per session: %v", durations)
|
||||
}
|
||||
}
|
||||
|
||||
func TestObserverQueueSnapshotsAndSampledPeaks(t *testing.T) {
|
||||
startSending := make(chan struct{})
|
||||
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()
|
||||
<-startSending
|
||||
for i := 0; i < 5; i++ {
|
||||
if conn.WriteMessage(websocket.TextMessage, []byte("queue")) != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
for {
|
||||
if _, _, err = conn.ReadMessage(); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
client := testClient(t)
|
||||
o := New(nil)
|
||||
opened := make(chan struct{})
|
||||
release := make(chan struct{})
|
||||
s, err := client.Dial(context.Background(), endpoint(server), o.Wrap(ws.Handlers{OnOpen: func(context.Context, *ws.Session) error { close(opened); return nil }, OnMessage: func(context.Context, *ws.Session, ws.Message) error { <-release; return nil }}))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
wait(t, opened)
|
||||
close(startSending)
|
||||
deadline := time.After(2 * time.Second)
|
||||
var sampled Snapshot
|
||||
for {
|
||||
sampled = o.Snapshot()
|
||||
if sampled.MessagesReceived == 5 {
|
||||
break
|
||||
}
|
||||
select {
|
||||
case <-deadline:
|
||||
close(release)
|
||||
t.Fatal("receive timeout")
|
||||
default:
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
}
|
||||
if sampled.ReadQueuedMessages != 5 || sampled.ReadQueuedBytes != 25 || sampled.SampledPeakReadQueuedMessages != 5 {
|
||||
t.Fatalf("queue snapshot: %+v", sampled)
|
||||
}
|
||||
close(release)
|
||||
_ = s.Close(1000, "")
|
||||
wait(t, s.Done())
|
||||
final := o.Snapshot()
|
||||
if final.ReadQueuedMessages != 0 || final.SampledPeakReadQueuedMessages != 5 {
|
||||
t.Fatalf("final: %+v", final)
|
||||
}
|
||||
}
|
||||
|
||||
type opaqueError struct{}
|
||||
|
||||
func (opaqueError) Error() string { panic("must not expand error") }
|
||||
func TestObserverCloseFieldsNeverExpandsErrors(t *testing.T) {
|
||||
for _, err := range []error{opaqueError{}, &ws.OpError{Op: "SECRET STAGE", Cause: opaqueError{}}, &ws.CloseError{Code: 1008, Cause: opaqueError{}}} {
|
||||
_, phase := closeFields(err, false)
|
||||
if strings.Contains(phase, "SECRET") {
|
||||
t.Fatal("untrusted phase")
|
||||
}
|
||||
}
|
||||
if code, phase := closeFields(opaqueError{}, true); code != 0 || phase != "on_close_panic" {
|
||||
t.Fatalf("panic fields: %d %s", code, phase)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user