core-lead/molecule-core
1
0
Fork 0
forked from molecule-ai/molecule-core
Code Pull Requests Activity

Merge pull request #719 from Molecule-AI/fix/issue-697-validate-token-removed-workspace

Browse Source 
fix(wsauth): add removed-workspace JOIN to ValidateToken (#697)
This commit is contained in:
molecule-ai[bot] 2026-04-17 12:50:52 +00:00 committed by GitHub
commit c53bf6eebd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 42 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
platform/internal/handlers/a2a_proxy_test.go
View File

@ -742,7 +742,7 @@ func TestValidateCallerToken_InvalidToken(t *testing.T) {
mock.ExpectQuery(`SELECT COUNT\(\*\) FROM workspace_auth_tokens`).
WithArgs("ws-authed").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1))
mock.ExpectQuery(`SELECT id, workspace_id FROM workspace_auth_tokens`).
mock.ExpectQuery(`SELECT t\.id, t\.workspace_id.*FROM workspace_auth_tokens t.*JOIN workspaces`).
WillReturnError(sql.ErrNoRows)
w := httptest.NewRecorder()
@ -766,7 +766,7 @@ func TestValidateCallerToken_ValidToken(t *testing.T) {
mock.ExpectQuery(`SELECT COUNT\(\*\) FROM workspace_auth_tokens`).
WithArgs("ws-authed").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1))
mock.ExpectQuery(`SELECT id, workspace_id FROM workspace_auth_tokens`).
mock.ExpectQuery(`SELECT t\.id, t\.workspace_id.*FROM workspace_auth_tokens t.*JOIN workspaces`).
WithArgs(sqlmock.AnyArg()).
WillReturnRows(sqlmock.NewRows([]string{"id", "workspace_id"}).AddRow("t1", "ws-authed"))
mock.ExpectExec(`UPDATE workspace_auth_tokens SET last_used_at`).
@ -794,7 +794,7 @@ func TestValidateCallerToken_WrongWorkspaceBindingRejected(t *testing.T) {
mock.ExpectQuery(`SELECT COUNT\(\*\) FROM workspace_auth_tokens`).
WithArgs("ws-b-attacker").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1))
mock.ExpectQuery(`SELECT id, workspace_id FROM workspace_auth_tokens`).
mock.ExpectQuery(`SELECT t\.id, t\.workspace_id.*FROM workspace_auth_tokens t.*JOIN workspaces`).
WillReturnRows(sqlmock.NewRows([]string{"id", "workspace_id"}).AddRow("t-a", "ws-a-owner"))
w := httptest.NewRecorder()

2
platform/internal/handlers/admin_test_token_test.go
View File

@ -117,7 +117,7 @@ func TestAdminTestToken_HappyPath_TokenValidates(t *testing.T) {
// doesn't capture live args; the important invariant is that the issued
// token passes ValidateToken given a matching hash row exists.)
_ = capturedHash
mock.ExpectQuery("SELECT id, workspace_id\\s+FROM workspace_auth_tokens").
mock.ExpectQuery("SELECT t\\.id, t\\.workspace_id.*FROM workspace_auth_tokens t.*JOIN workspaces").
WithArgs(sqlmock.AnyArg()).
WillReturnRows(sqlmock.NewRows([]string{"id", "workspace_id"}).AddRow("tok-1", "ws-1"))
mock.ExpectExec("UPDATE workspace_auth_tokens SET last_used_at").

4
platform/internal/handlers/secrets_test.go
View File

@ -614,7 +614,7 @@ func TestSecretsValues_WrongToken(t *testing.T) {
WithArgs(testWsID).
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1))
// ValidateToken lookup returns nothing
mock.ExpectQuery(`SELECT id, workspace_id FROM workspace_auth_tokens`).
mock.ExpectQuery(`SELECT t\.id, t\.workspace_id.*FROM workspace_auth_tokens t.*JOIN workspaces`).
WillReturnError(sql.ErrNoRows)
w := httptest.NewRecorder()
@ -633,7 +633,7 @@ func TestSecretsValues_ValidTokenReturnsDecryptedMerge(t *testing.T) {
mock.ExpectQuery(`SELECT COUNT\(\*\) FROM workspace_auth_tokens`).
WithArgs(testWsID).
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1))
mock.ExpectQuery(`SELECT id, workspace_id FROM workspace_auth_tokens`).
mock.ExpectQuery(`SELECT t\.id, t\.workspace_id.*FROM workspace_auth_tokens t.*JOIN workspaces`).
WithArgs(sqlmock.AnyArg()).
WillReturnRows(sqlmock.NewRows([]string{"id", "workspace_id"}).AddRow("tok-1", testWsID))
mock.ExpectExec(`UPDATE workspace_auth_tokens SET last_used_at`).

2
platform/internal/handlers/workspace_test.go
View File

@ -752,7 +752,7 @@ func TestWorkspaceState_ValidTokenReturnsStatus(t *testing.T) {
mock.ExpectQuery(`SELECT COUNT\(\*\) FROM workspace_auth_tokens`).
WithArgs(stateWsID).
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1))
mock.ExpectQuery(`SELECT id, workspace_id FROM workspace_auth_tokens`).
mock.ExpectQuery(`SELECT t\.id, t\.workspace_id.*FROM workspace_auth_tokens t.*JOIN workspaces`).
WithArgs(sqlmock.AnyArg()).
WillReturnRows(sqlmock.NewRows([]string{"id", "workspace_id"}).AddRow("t1", stateWsID))
mock.ExpectExec(`UPDATE workspace_auth_tokens SET last_used_at`).

2
platform/internal/middleware/wsauth_middleware_test.go
View File

@ -23,7 +23,7 @@ import (
const hasAnyLiveTokenGlobalQuery = "SELECT COUNT.*FROM workspace_auth_tokens"
// validateTokenQuery is matched for ValidateToken (SELECT).
const validateTokenSelectQuery = "SELECT id, workspace_id.*FROM workspace_auth_tokens.*token_hash"
const validateTokenSelectQuery = "SELECT t\\.id, t\\.workspace_id.*FROM workspace_auth_tokens t.*JOIN workspaces"
// validateAnyTokenQuery is matched for ValidateAnyToken (SELECT).
// The JOIN on workspaces filters removed-workspace tokens (#682 defense-in-depth).

14
platform/internal/wsauth/tokens.go
View File

@ -73,6 +73,11 @@ func IssueToken(ctx context.Context, db *sql.DB, workspaceID string) (string, er
// The expectedWorkspaceID binding is required because a token is only
// valid for the workspace it was issued to. A compromised token from
// workspace A must never authenticate workspace B.
//
// Defense-in-depth (#697): the JOIN on workspaces filters out tokens that
// belong to removed workspaces so that a deleted workspace's tokens cannot
// be replayed against its former sub-routes even before the token row is
// explicitly revoked. Mirrors the same guard added to ValidateAnyToken (#696).
func ValidateToken(ctx context.Context, db *sql.DB, expectedWorkspaceID, plaintext string) error {
if plaintext == "" || expectedWorkspaceID == "" {
return ErrInvalidToken
@ -81,9 +86,12 @@ func ValidateToken(ctx context.Context, db *sql.DB, expectedWorkspaceID, plainte
var tokenID, workspaceID string
err := db.QueryRowContext(ctx, `
SELECT id, workspace_id
FROM workspace_auth_tokens
WHERE token_hash = $1 AND revoked_at IS NULL
SELECT t.id, t.workspace_id
FROM workspace_auth_tokens t
JOIN workspaces w ON w.id = t.workspace_id
WHERE t.token_hash = $1
AND t.revoked_at IS NULL
AND w.status != 'removed'
`, hash[:]).Scan(&tokenID, &workspaceID)
if err != nil {
// Includes sql.ErrNoRows — collapse to a single public-facing error

27
platform/internal/wsauth/tokens_test.go
View File

@ -76,8 +76,8 @@ func TestValidateToken_HappyPath(t *testing.T) {
t.Fatalf("IssueToken: %v", err)
}
// Validate: lookup by hash returns matching workspace.
mock.ExpectQuery(`SELECT id, workspace_id FROM workspace_auth_tokens`).
// Validate: lookup by hash with removed-workspace JOIN returns matching row.
mock.ExpectQuery(`SELECT t\.id, t\.workspace_id.*FROM workspace_auth_tokens t.*JOIN workspaces`).
WithArgs(sqlmock.AnyArg()).
WillReturnRows(sqlmock.NewRows([]string{"id", "workspace_id"}).AddRow("tok-id-1", "ws-xyz"))
// Best-effort last_used_at update.
@ -94,7 +94,7 @@ func TestValidateToken_WrongWorkspaceRejected(t *testing.T) {
db, mock := setupMock(t)
// Token belongs to ws-owner; caller claims to be ws-attacker.
mock.ExpectQuery(`SELECT id, workspace_id FROM workspace_auth_tokens`).
mock.ExpectQuery(`SELECT t\.id, t\.workspace_id.*FROM workspace_auth_tokens t.*JOIN workspaces`).
WillReturnRows(sqlmock.NewRows([]string{"id", "workspace_id"}).AddRow("tok-id-2", "ws-owner"))
err := ValidateToken(context.Background(), db, "ws-attacker", "some-token")
@ -115,7 +115,7 @@ func TestValidateToken_RejectsEmptyInputs(t *testing.T) {
func TestValidateToken_UnknownTokenRejected(t *testing.T) {
db, mock := setupMock(t)
mock.ExpectQuery(`SELECT id, workspace_id FROM workspace_auth_tokens`).
mock.ExpectQuery(`SELECT t\.id, t\.workspace_id.*FROM workspace_auth_tokens t.*JOIN workspaces`).
WillReturnError(sql.ErrNoRows)
if err := ValidateToken(context.Background(), db, "ws-a", "not-a-real-token"); err != ErrInvalidToken {
@ -123,6 +123,25 @@ func TestValidateToken_UnknownTokenRejected(t *testing.T) {
}
}
// TestValidateToken_RemovedWorkspaceRejected — defense-in-depth (#697):
// a token belonging to a workspace with status='removed' must be rejected
// even when the token row itself is still live (revoked_at IS NULL).
// The JOIN on workspaces with AND w.status != 'removed' filters the row
// out at the DB layer, returning ErrNoRows which collapses to ErrInvalidToken.
func TestValidateToken_RemovedWorkspaceRejected(t *testing.T) {
db, mock := setupMock(t)
// JOIN with w.status != 'removed' causes no rows — same path as ErrNoRows.
mock.ExpectQuery(`SELECT t\.id, t\.workspace_id.*FROM workspace_auth_tokens t.*JOIN workspaces`).
WithArgs(sqlmock.AnyArg()).
WillReturnRows(sqlmock.NewRows([]string{"id", "workspace_id"})) // empty: workspace removed
err := ValidateToken(context.Background(), db, "ws-removed", "token-for-removed-workspace")
if err != ErrInvalidToken {
t.Errorf("removed workspace token: expected ErrInvalidToken, got %v", err)
}
}
// ------------------------------------------------------------
// HasAnyLiveToken
// ------------------------------------------------------------