From a55f8c36c8ac29a277ef9c1ddbf517e2ac99fd2f Mon Sep 17 00:00:00 2001 From: Molecule AI Core-BE Date: Tue, 12 May 2026 07:40:17 +0000 Subject: [PATCH 01/10] =?UTF-8?q?test(handlers/socket):=20add=20socket=5Ft?= =?UTF-8?q?est.go=20=E2=80=94=206=20cases=20covering=20Phase=2030.1/30.2?= =?UTF-8?q?=20auth=20gate?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit HandleConnect has two branches: 1. Canvas clients (no X-Workspace-ID): auth gate bypassed entirely 2. Workspace agents (X-Workspace-ID present): Phase 30.1/30.2 bearer token enforcement — HasAnyLiveToken gates ValidateToken. 6 cases: - DB error on HasAnyLiveToken → 500 - hasLive=true, no Bearer header → 401 - hasLive=true, invalid Bearer → 401 - hasLive=true, empty Bearer → 401 (ValidateToken ErrInvalidToken) - hasLive=true, valid Bearer → auth passed (upgrade fails in httptest; verified by absence of 401/500) - canvas client (no X-Workspace-ID) → auth bypassed WebSocket upgrade itself not testable in httptest; covered by the auth-pass cases which verify the upgrade is reached without returning an auth error. --- .../internal/handlers/socket_test.go | 243 ++++++++++++++++++ 1 file changed, 243 insertions(+) create mode 100644 workspace-server/internal/handlers/socket_test.go diff --git a/workspace-server/internal/handlers/socket_test.go b/workspace-server/internal/handlers/socket_test.go new file mode 100644 index 00000000..63af3682 --- /dev/null +++ b/workspace-server/internal/handlers/socket_test.go @@ -0,0 +1,243 @@ +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/db" + "github.com/Molecule-AI/molecule-monorepo/platform/internal/ws" + "github.com/gin-gonic/gin" +) + +// newSocketHandlerWithDB creates a SocketHandler backed by the given mock DB +// and a Hub with buffered channels so Register/Unregister don't block. +func newSocketHandlerWithDB(t *testing.T, mockDB *sql.DB, hub *ws.Hub) *SocketHandler { + db.DB = mockDB + 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, mock, 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, mock, 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, mock, 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, mock, 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, mock, 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, mock, 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[:] +} -- 2.45.2 From 0733a2815c9040136f8104ecb1185c871aa67b58 Mon Sep 17 00:00:00 2001 From: core-devops Date: Tue, 12 May 2026 20:52:06 +0000 Subject: [PATCH 02/10] ci: rerun after mc#724 all-required fix lands -- 2.45.2 From 4fa992a6413c76c73e440a73bc5c82449e7a7a5e Mon Sep 17 00:00:00 2001 From: claude-ceo-assistant Date: Tue, 12 May 2026 21:16:49 +0000 Subject: [PATCH 03/10] ci: rerun after concurrency-block clear -- 2.45.2 From 1e4e49d1494b1203966bb89a0eeead7d6624d19c Mon Sep 17 00:00:00 2001 From: claude-ceo-assistant Date: Tue, 12 May 2026 21:25:31 +0000 Subject: [PATCH 04/10] ci: clean-slate rerun -- 2.45.2 From 37ff6b72983cfb609071da01c70188c2fff811fb Mon Sep 17 00:00:00 2001 From: claude-ceo-assistant Date: Tue, 12 May 2026 21:30:27 +0000 Subject: [PATCH 05/10] ci: post-restart rerun -- 2.45.2 From 94ec46c89f8c68a45a3bb69fa8921cbe6cb8be83 Mon Sep 17 00:00:00 2001 From: core-lead Date: Tue, 12 May 2026 21:35:12 +0000 Subject: [PATCH 06/10] ci: clean-slate rerun v2 -- 2.45.2 From debd8e4d100eb8fa03b37ae9ee63703f2257fda7 Mon Sep 17 00:00:00 2001 From: core-lead Date: Tue, 12 May 2026 21:44:43 +0000 Subject: [PATCH 07/10] ci: global-zombie-purge rerun -- 2.45.2 From d66ef046032b7d108f54a3ed2ac36c98de487f64 Mon Sep 17 00:00:00 2001 From: core-lead Date: Tue, 12 May 2026 21:48:31 +0000 Subject: [PATCH 08/10] ci: post-full-purge rerun -- 2.45.2 From 30fcf9cb45ebf436c3898285b28c7a8dee366fdb Mon Sep 17 00:00:00 2001 From: core-devops Date: Tue, 12 May 2026 22:07:24 +0000 Subject: [PATCH 09/10] ci: post-purge rerun -- 2.45.2 From d2661bb0cb1342e06c1ea93d13108727f4706bab Mon Sep 17 00:00:00 2001 From: Molecule AI Core-Security Date: Tue, 12 May 2026 18:40:42 -0700 Subject: [PATCH 10/10] =?UTF-8?q?fix(test/handlers):=20correct=20newSocket?= =?UTF-8?q?HandlerWithDB=20signature=20=E2=80=94=20drop=20*sql.DB=20param?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit setupTestDB already sets db.DB globally; passing sqlmock.Sqlmock as *sql.DB caused a build failure. Remove the redundant parameter and update callers. Co-Authored-By: claude-sonnet-4-6 --- .../internal/handlers/socket_test.go | 21 +++++++++---------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/workspace-server/internal/handlers/socket_test.go b/workspace-server/internal/handlers/socket_test.go index 63af3682..eb889fe4 100644 --- a/workspace-server/internal/handlers/socket_test.go +++ b/workspace-server/internal/handlers/socket_test.go @@ -8,15 +8,14 @@ import ( "testing" "github.com/DATA-DOG/go-sqlmock" - "github.com/Molecule-AI/molecule-monorepo/platform/internal/db" "github.com/Molecule-AI/molecule-monorepo/platform/internal/ws" "github.com/gin-gonic/gin" ) -// newSocketHandlerWithDB creates a SocketHandler backed by the given mock DB -// and a Hub with buffered channels so Register/Unregister don't block. -func newSocketHandlerWithDB(t *testing.T, mockDB *sql.DB, hub *ws.Hub) *SocketHandler { - db.DB = mockDB +// 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), @@ -44,7 +43,7 @@ func socketRequest(method, path, workspaceID, authHeader string) *http.Request { func TestSocketHandler_AuthGate_DBError_Returns500(t *testing.T) { mock := setupTestDB(t) - handler := newSocketHandlerWithDB(t, mock, nil) + handler := newSocketHandlerWithDB(t, nil) // HasAnyLiveToken issues a query; make it return an error. mock.ExpectQuery("SELECT COUNT"). @@ -71,7 +70,7 @@ func TestSocketHandler_AuthGate_DBError_Returns500(t *testing.T) { func TestSocketHandler_AuthGate_HasLiveToken_MissingBearer_Returns401(t *testing.T) { mock := setupTestDB(t) - handler := newSocketHandlerWithDB(t, mock, nil) + handler := newSocketHandlerWithDB(t, nil) // HasAnyLiveToken succeeds → workspace has a live token. mock.ExpectQuery("SELECT COUNT"). @@ -98,7 +97,7 @@ func TestSocketHandler_AuthGate_HasLiveToken_MissingBearer_Returns401(t *testing func TestSocketHandler_AuthGate_HasLiveToken_InvalidBearer_Returns401(t *testing.T) { mock := setupTestDB(t) - handler := newSocketHandlerWithDB(t, mock, nil) + handler := newSocketHandlerWithDB(t, nil) wsID := "ws-invalid-token" badToken := "not-a-valid-token" @@ -138,7 +137,7 @@ func TestSocketHandler_AuthGate_HasLiveToken_InvalidBearer_Returns401(t *testing func TestSocketHandler_AuthGate_HasLiveToken_ValidBearer_AuthPassed(t *testing.T) { mock := setupTestDB(t) - handler := newSocketHandlerWithDB(t, mock, nil) + handler := newSocketHandlerWithDB(t, nil) wsID := "ws-valid-token" goodToken := "valid-ws-token-123" @@ -180,7 +179,7 @@ func TestSocketHandler_AuthGate_HasLiveToken_ValidBearer_AuthPassed(t *testing.T func TestSocketHandler_CanvasClient_NoAuthGate(t *testing.T) { mock := setupTestDB(t) - handler := newSocketHandlerWithDB(t, mock, nil) + handler := newSocketHandlerWithDB(t, nil) // No X-Workspace-ID header → no auth check → no DB queries expected. w := httptest.NewRecorder() @@ -208,7 +207,7 @@ func TestSocketHandler_CanvasClient_NoAuthGate(t *testing.T) { func TestSocketHandler_AuthGate_HasLiveToken_EmptyBearer_Returns401(t *testing.T) { mock := setupTestDB(t) - handler := newSocketHandlerWithDB(t, mock, nil) + handler := newSocketHandlerWithDB(t, nil) wsID := "ws-has-live-token-empty-bearer" -- 2.45.2