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 }