diff --git a/workspace-server/internal/handlers/socket_test.go b/workspace-server/internal/handlers/socket_test.go new file mode 100644 index 00000000..eb889fe4 --- /dev/null +++ b/workspace-server/internal/handlers/socket_test.go @@ -0,0 +1,242 @@ +package handlers + +import ( + "crypto/sha256" + "database/sql" + "net/http" + "net/http/httptest" + "testing" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/Molecule-AI/molecule-monorepo/platform/internal/ws" + "github.com/gin-gonic/gin" +) + +// newSocketHandlerWithDB creates a SocketHandler with buffered Hub channels. +// The DB is set up via setupTestDB (called before this function in each test). +func newSocketHandlerWithDB(t *testing.T, hub *ws.Hub) *SocketHandler { + t.Helper() + if hub == nil { + hub = &ws.Hub{ + Register: make(chan *ws.Client, 1), + Unregister: make(chan *ws.Client, 1), + } + } + return NewSocketHandler(hub) +} + +// socketRequest builds a test request for the WebSocket connect endpoint. +func socketRequest(method, path, workspaceID, authHeader string) *http.Request { + req := httptest.NewRequest(method, path, nil) + if workspaceID != "" { + req.Header.Set("X-Workspace-ID", workspaceID) + } + if authHeader != "" { + req.Header.Set("Authorization", authHeader) + } + return req +} + +// ───────────────────────────────────────────────────────────────────────────── +// Auth gate: DB error on HasAnyLiveToken → 500 +// ───────────────────────────────────────────────────────────────────────────── + +func TestSocketHandler_AuthGate_DBError_Returns500(t *testing.T) { + mock := setupTestDB(t) + handler := newSocketHandlerWithDB(t, nil) + + // HasAnyLiveToken issues a query; make it return an error. + mock.ExpectQuery("SELECT COUNT"). + WithArgs("ws-auth-db-err"). + WillReturnError(sql.ErrConnDone) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = socketRequest("GET", "/ws", "ws-auth-db-err", "") + + handler.HandleConnect(c) + + if w.Code != http.StatusInternalServerError { + t.Errorf("DB error: expected 500, got %d", w.Code) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet mock expectations: %v", err) + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// Auth gate: workspace HAS live token, missing Bearer → 401 +// ───────────────────────────────────────────────────────────────────────────── + +func TestSocketHandler_AuthGate_HasLiveToken_MissingBearer_Returns401(t *testing.T) { + mock := setupTestDB(t) + handler := newSocketHandlerWithDB(t, nil) + + // HasAnyLiveToken succeeds → workspace has a live token. + mock.ExpectQuery("SELECT COUNT"). + WithArgs("ws-has-token-no-bearer"). + WillReturnRows(sqlmock.NewRows([]string{"n"}).AddRow(1)) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = socketRequest("GET", "/ws", "ws-has-token-no-bearer", "") + + handler.HandleConnect(c) + + if w.Code != http.StatusUnauthorized { + t.Errorf("hasLive but no bearer: expected 401, got %d", w.Code) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet mock expectations: %v", err) + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// Auth gate: workspace HAS live token, invalid Bearer → 401 +// ───────────────────────────────────────────────────────────────────────────── + +func TestSocketHandler_AuthGate_HasLiveToken_InvalidBearer_Returns401(t *testing.T) { + mock := setupTestDB(t) + handler := newSocketHandlerWithDB(t, nil) + + wsID := "ws-invalid-token" + badToken := "not-a-valid-token" + + // HasAnyLiveToken: workspace has a live token. + mock.ExpectQuery("SELECT COUNT"). + WithArgs(wsID). + WillReturnRows(sqlmock.NewRows([]string{"n"}).AddRow(1)) + + // ValidateToken: lookupTokenByHash returns ErrNoRows for an unknown token. + // Any token hash is fine since the token doesn't exist — use AnyArg. + mock.ExpectQuery(`SELECT t\.id, t\.workspace_id.*FROM workspace_auth_tokens t.*JOIN`). + WithArgs(sqlmock.AnyArg()). + WillReturnError(sql.ErrNoRows) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = socketRequest("GET", "/ws", wsID, "Bearer "+badToken) + + handler.HandleConnect(c) + + if w.Code != http.StatusUnauthorized { + t.Errorf("hasLive but invalid bearer: expected 401, got %d", w.Code) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet mock expectations: %v", err) + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// Auth gate: workspace HAS live token, VALID Bearer → upgrade attempted. +// The WebSocket upgrade itself will fail in httptest (gorilla/websocket +// cannot write a real HTTP/1.1 handshake to httptest.ResponseRecorder), but +// the auth gate is passed so we verify no 401/500 was returned before the +// upgrade failure. This is the canvas-client success path. +// ───────────────────────────────────────────────────────────────────────────── + +func TestSocketHandler_AuthGate_HasLiveToken_ValidBearer_AuthPassed(t *testing.T) { + mock := setupTestDB(t) + handler := newSocketHandlerWithDB(t, nil) + + wsID := "ws-valid-token" + goodToken := "valid-ws-token-123" + + // HasAnyLiveToken: workspace has a live token. + mock.ExpectQuery("SELECT COUNT"). + WithArgs(wsID). + WillReturnRows(sqlmock.NewRows([]string{"n"}).AddRow(1)) + + // ValidateToken: token found and workspace is not removed. + // sha256TokenHash returns []byte; rational matcher compares as string. + mock.ExpectQuery(`SELECT t\.id, t\.workspace_id.*FROM workspace_auth_tokens t.*JOIN`). + WithArgs(sha256TokenHash(goodToken)). + WillReturnRows(sqlmock.NewRows([]string{"token_id", "workspace_id"}). + AddRow("tok-abc", wsID)) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = socketRequest("GET", "/ws", wsID, "Bearer "+goodToken) + + handler.HandleConnect(c) + + // The WebSocket upgrade fails in httptest (httptest.ResponseRecorder is not + // a real TCP connection), but the auth gate itself succeeded — we should + // NOT see a 401 or 500 response code. The actual code depends on the + // upgrade error handling; the critical assertion is that auth passed. + if w.Code == http.StatusUnauthorized || w.Code == http.StatusInternalServerError { + t.Errorf("valid token: auth should have passed; got %d", w.Code) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet mock expectations: %v", err) + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// Canvas client (no X-Workspace-ID): auth gate bypassed, upgrade attempted. +// Same httptest limitation as above — we verify no 401/500 before the upgrade. +// ───────────────────────────────────────────────────────────────────────────── + +func TestSocketHandler_CanvasClient_NoAuthGate(t *testing.T) { + mock := setupTestDB(t) + handler := newSocketHandlerWithDB(t, nil) + + // No X-Workspace-ID header → no auth check → no DB queries expected. + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = socketRequest("GET", "/ws", "", "") // no workspace ID + + handler.HandleConnect(c) + + // No auth gate hit → no 401/500. The WebSocket upgrade itself will fail + // in httptest, but that's expected (see TestSocketHandler_AuthGate_HasLiveToken_ValidBearer_AuthPassed). + if w.Code == http.StatusUnauthorized || w.Code == http.StatusInternalServerError { + t.Errorf("canvas client: expected no auth error; got %d", w.Code) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet mock expectations: %v", err) + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// Legacy workspace: HAS live token flag but workspace exists AND ValidateToken +// is called. Since the workspace has a live token, the handler MUST validate +// the presented token (not grandfather through). This is the Phase 30.1/30.2 +// contract — a workspace with tokens on file is NOT grandfathered. +// ───────────────────────────────────────────────────────────────────────────── + +func TestSocketHandler_AuthGate_HasLiveToken_EmptyBearer_Returns401(t *testing.T) { + mock := setupTestDB(t) + handler := newSocketHandlerWithDB(t, nil) + + wsID := "ws-has-live-token-empty-bearer" + + // HasAnyLiveToken: workspace has a live token. + mock.ExpectQuery("SELECT COUNT"). + WithArgs(wsID). + WillReturnRows(sqlmock.NewRows([]string{"n"}).AddRow(1)) + + // Authorization header is "Bearer " (empty token after "Bearer "). + // wsauth.BearerTokenFromHeader strips "Bearer " and gets "". + // ValidateToken is called with "" → returns ErrInvalidToken before DB hit. + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = socketRequest("GET", "/ws", wsID, "Bearer ") + + handler.HandleConnect(c) + + if w.Code != http.StatusUnauthorized { + t.Errorf("empty bearer after Bearer prefix: expected 401, got %d", w.Code) + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// helpers +// ───────────────────────────────────────────────────────────────────────────── + +// sha256TokenHash returns the SHA256 hash of a plaintext token, matching what +// wsauth.ValidateToken does internally before querying the DB. +func sha256TokenHash(plaintext string) []byte { + h := sha256.Sum256([]byte(plaintext)) + return h[:] +}