flvx/go-backend/internal/ws/server_test.go

package ws

import (
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"sync"
	"testing"

	"go-backend/internal/auth"
	"go-backend/internal/store/repo"

	"github.com/gorilla/websocket"
)

func TestServeHTTPRejectsDisabledAdminToken(t *testing.T) {
	secret := "unit-test-secret"
	token, err := auth.GenerateToken(1, "admin_user", 0, secret)
	if err != nil {
		t.Fatalf("generate token: %v", err)
	}

	server := NewServer(nil, secret)
	server.SetUserAuthStateLookup(func(userID int64) (*auth.UserAuthState, error) {
		return &auth.UserAuthState{ID: userID, RoleID: 0, Status: 0, PasswordChangedAt: 0}, nil
	})

	req := httptest.NewRequest(http.MethodGet, "/system-info?type=0&secret="+url.QueryEscape(token), nil)
	rec := httptest.NewRecorder()

	server.ServeHTTP(rec, req)

	if rec.Code != http.StatusForbidden {
		t.Fatalf("expected forbidden for disabled admin token, got %d", rec.Code)
	}
}

func TestServeHTTPRejectsNonAdminToken(t *testing.T) {
	secret := "unit-test-secret"
	token, err := auth.GenerateToken(2, "normal_user", 1, secret)
	if err != nil {
		t.Fatalf("generate token: %v", err)
	}

	server := NewServer(nil, secret)
	server.SetUserAuthStateLookup(func(userID int64) (*auth.UserAuthState, error) {
		return &auth.UserAuthState{ID: userID, RoleID: 1, Status: 1, PasswordChangedAt: 0}, nil
	})

	req := httptest.NewRequest(http.MethodGet, "/system-info?type=0&secret="+url.QueryEscape(token), nil)
	rec := httptest.NewRecorder()

	server.ServeHTTP(rec, req)

	if rec.Code != http.StatusForbidden {
		t.Fatalf("expected forbidden for non-admin token, got %d", rec.Code)
	}
}

func TestValidateAdminSessionRejectsAuthStateChanges(t *testing.T) {
	secret := "unit-test-secret"
	token, err := auth.GenerateToken(1, "admin_user", 0, secret)
	if err != nil {
		t.Fatalf("generate token: %v", err)
	}
	claims, err := auth.ParseClaims(token, secret)
	if err != nil {
		t.Fatalf("parse claims: %v", err)
	}

	tests := []struct {
		name  string
		state *auth.UserAuthState
	}{
		{name: "disabled", state: &auth.UserAuthState{ID: 1, RoleID: 0, Status: 0, PasswordChangedAt: 0}},
		{name: "role changed", state: &auth.UserAuthState{ID: 1, RoleID: 1, Status: 1, PasswordChangedAt: 0}},
		{name: "password changed", state: &auth.UserAuthState{ID: 1, RoleID: 0, Status: 1, PasswordChangedAt: claims.IatMs + 1}},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			server := NewServer(nil, secret)
			server.SetUserAuthStateLookup(func(userID int64) (*auth.UserAuthState, error) {
				return tt.state, nil
			})

			if ok := server.validateAdminSession(1, claims); ok {
				t.Fatalf("expected session validation to fail for %s state", tt.name)
			}
		})
	}
}

func TestValidateAdminSessionRejectsExpiredToken(t *testing.T) {
	secret := "unit-test-secret"
	token, err := auth.GenerateToken(1, "admin_user", 0, secret)
	if err != nil {
		t.Fatalf("generate token: %v", err)
	}
	claims, err := auth.ParseClaims(token, secret)
	if err != nil {
		t.Fatalf("parse claims: %v", err)
	}
	claims.Exp = 1

	server := NewServer(nil, secret)
	server.SetUserAuthStateLookup(func(userID int64) (*auth.UserAuthState, error) {
		return &auth.UserAuthState{ID: userID, RoleID: 0, Status: 1, PasswordChangedAt: 0}, nil
	})

	if ok := server.validateAdminSession(1, claims); ok {
		t.Fatal("expected expired token to be rejected")
	}
}

func TestServeHTTPAllowsMonitorTokenWithPermission(t *testing.T) {
	secret := "unit-test-secret"
	token, err := auth.GenerateToken(2, "normal_user", 1, secret)
	if err != nil {
		t.Fatalf("generate token: %v", err)
	}

	r, err := repo.Open(t.TempDir() + "/monitor.db")
	if err != nil {
		t.Fatalf("open repo: %v", err)
	}
	defer r.Close()
	if err := r.InsertMonitorPermission(2, 123); err != nil {
		t.Fatalf("insert permission: %v", err)
	}

	server := NewServer(r, secret)
	server.SetUserAuthStateLookup(func(userID int64) (*auth.UserAuthState, error) {
		return &auth.UserAuthState{ID: userID, RoleID: 1, Status: 1, PasswordChangedAt: 0}, nil
	})

	ts := httptest.NewServer(server)
	defer ts.Close()

	conn, resp, err := websocket.DefaultDialer.Dial(
		"ws"+strings.TrimPrefix(ts.URL, "http")+"/system-info?type=0&secret="+url.QueryEscape(token),
		nil,
	)
	if err != nil {
		if resp != nil {
			t.Fatalf("dial websocket error = %v, status=%d", err, resp.StatusCode)
		}
		t.Fatalf("dial websocket error = %v", err)
	}
	_ = conn.Close()
}

func TestConnWrapSerializesConcurrentWrites(t *testing.T) {
	serverConn, clientConn := websocketTestPipe(t)
	defer serverConn.Close()
	defer clientConn.Close()

	cw := &connWrap{conn: serverConn}
	readerDone := make(chan struct{})
	go func() {
		defer close(readerDone)
		for i := 0; i < 64; i++ {
			if _, _, err := clientConn.ReadMessage(); err != nil {
				return
			}
		}
	}()

	var wg sync.WaitGroup
	for i := 0; i < 64; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if err := writeWSMessage(cw, websocket.TextMessage, []byte("x")); err != nil {
				t.Errorf("writeWSMessage() error = %v", err)
			}
		}()
	}

	wg.Wait()
	_ = clientConn.Close()
	<-readerDone
}

func websocketTestPipe(t *testing.T) (*websocket.Conn, *websocket.Conn) {
	t.Helper()

	upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
	serverConnCh := make(chan *websocket.Conn, 1)
	serverErrCh := make(chan error, 1)
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			serverErrCh <- err
			return
		}
		serverConnCh <- conn
	}))
	t.Cleanup(ts.Close)

	clientConn, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(ts.URL, "http"), nil)
	if err != nil {
		t.Fatalf("dial websocket: %v", err)
	}

	select {
	case err := <-serverErrCh:
		t.Fatalf("upgrade websocket: %v", err)
	case serverConn := <-serverConnCh:
		return serverConn, clientConn
	}

	return nil, nil
}