test(handlers/socket): add socket_test.go — 6 cases covering Phase 30.1/30.2 auth gate
Some checks failed
Block internal-flavored paths / Block forbidden paths (pull_request) Successful in 18s
CI / Detect changes (pull_request) Successful in 48s
E2E API Smoke Test / detect-changes (pull_request) Successful in 52s
E2E Staging Canvas (Playwright) / detect-changes (pull_request) Successful in 56s
Harness Replays / detect-changes (pull_request) Successful in 28s
Handlers Postgres Integration / detect-changes (pull_request) Successful in 47s
Secret scan / Scan diff for credential-shaped strings (pull_request) Successful in 21s
qa-review / approved (pull_request) Failing after 22s
gate-check-v3 / gate-check (pull_request) Successful in 40s
sop-checklist / all-items-acked (pull_request) acked: 0/7 — missing: comprehensive-testing, local-postgres-e2e, staging-smoke, +4 — body-unfilled: 7
security-review / approved (pull_request) Failing after 22s
sop-checklist-gate / gate (pull_request) Successful in 23s
Runtime PR-Built Compatibility / detect-changes (pull_request) Successful in 52s
sop-tier-check / tier-check (pull_request) Successful in 22s
lint-required-no-paths / lint-required-no-paths (pull_request) Successful in 1m34s
CI / Canvas (Next.js) (pull_request) Successful in 14s
CI / Shellcheck (E2E scripts) (pull_request) Successful in 8s
CI / Python Lint & Test (pull_request) Successful in 9s
Harness Replays / Harness Replays (pull_request) Successful in 12s
E2E Staging Canvas (Playwright) / Canvas tabs E2E (pull_request) Successful in 11s
Runtime PR-Built Compatibility / PR-built wheel + import smoke (pull_request) Successful in 9s
CI / Canvas Deploy Reminder (pull_request) Has been skipped
E2E API Smoke Test / E2E API Smoke Test (pull_request) Failing after 5m51s
Handlers Postgres Integration / Handlers Postgres Integration (pull_request) Failing after 5m56s
CI / Platform (Go) (pull_request) Failing after 11m26s
CI / all-required (pull_request) Failing after 4s
Some checks failed
Block internal-flavored paths / Block forbidden paths (pull_request) Successful in 18s
CI / Detect changes (pull_request) Successful in 48s
E2E API Smoke Test / detect-changes (pull_request) Successful in 52s
E2E Staging Canvas (Playwright) / detect-changes (pull_request) Successful in 56s
Harness Replays / detect-changes (pull_request) Successful in 28s
Handlers Postgres Integration / detect-changes (pull_request) Successful in 47s
Secret scan / Scan diff for credential-shaped strings (pull_request) Successful in 21s
qa-review / approved (pull_request) Failing after 22s
gate-check-v3 / gate-check (pull_request) Successful in 40s
sop-checklist / all-items-acked (pull_request) acked: 0/7 — missing: comprehensive-testing, local-postgres-e2e, staging-smoke, +4 — body-unfilled: 7
security-review / approved (pull_request) Failing after 22s
sop-checklist-gate / gate (pull_request) Successful in 23s
Runtime PR-Built Compatibility / detect-changes (pull_request) Successful in 52s
sop-tier-check / tier-check (pull_request) Successful in 22s
lint-required-no-paths / lint-required-no-paths (pull_request) Successful in 1m34s
CI / Canvas (Next.js) (pull_request) Successful in 14s
CI / Shellcheck (E2E scripts) (pull_request) Successful in 8s
CI / Python Lint & Test (pull_request) Successful in 9s
Harness Replays / Harness Replays (pull_request) Successful in 12s
E2E Staging Canvas (Playwright) / Canvas tabs E2E (pull_request) Successful in 11s
Runtime PR-Built Compatibility / PR-built wheel + import smoke (pull_request) Successful in 9s
CI / Canvas Deploy Reminder (pull_request) Has been skipped
E2E API Smoke Test / E2E API Smoke Test (pull_request) Failing after 5m51s
Handlers Postgres Integration / Handlers Postgres Integration (pull_request) Failing after 5m56s
CI / Platform (Go) (pull_request) Failing after 11m26s
CI / all-required (pull_request) Failing after 4s
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.
This commit is contained in:
parent
9eb33a9d3c
commit
b76cc3ce61
243
workspace-server/internal/handlers/socket_test.go
Normal file
243
workspace-server/internal/handlers/socket_test.go
Normal file
@ -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[:]
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user