go-omni/ws/observe/observe_test.go

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)
	}
}