Compare commits
4 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 9a9efbf8c3 | |||
| b5411d2c37 | |||
| 03ad7ab2d8 | |||
| fd545a332b |
@@ -0,0 +1,425 @@
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"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/events"
|
||||
"github.com/Molecule-AI/molecule-monorepo/platform/internal/ws"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// setupBroadcastDB uses QueryMatcherEqual so SQL strings with quoted literals
|
||||
// (e.g. status != 'removed') are compared verbatim, not as regex.
|
||||
func setupBroadcastDB(t *testing.T) sqlmock.Sqlmock {
|
||||
t.Helper()
|
||||
mockDB, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual))
|
||||
if err != nil {
|
||||
t.Fatalf("failed to create sqlmock: %v", err)
|
||||
}
|
||||
prevDB := db.DB
|
||||
db.DB = mockDB
|
||||
t.Cleanup(func() { db.DB = prevDB; mockDB.Close() })
|
||||
return mock
|
||||
}
|
||||
|
||||
// validUUID is a properly formatted test UUID.
|
||||
const broadcastTestUUID = "bbbbbbbb-0001-0001-0001-000000000001"
|
||||
|
||||
// buildBroadcastCtx creates a gin.Context wired for POST /workspaces/:id/broadcast.
|
||||
func buildBroadcastCtx(id, body string) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
req := httptest.NewRequest(http.MethodPost, "/workspaces/"+id+"/broadcast", bytes.NewBufferString(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
c.Request = req.WithContext(context.Background())
|
||||
c.Params = gin.Params{{Key: "id", Value: id}}
|
||||
return c, w
|
||||
}
|
||||
|
||||
// ─── Pure function ────────────────────────────────────────────────────────────
|
||||
|
||||
func TestBroadcastTruncate(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
s string
|
||||
max int
|
||||
want string
|
||||
}{
|
||||
{"empty string", "", 10, ""},
|
||||
{"under limit", "hello", 10, "hello"},
|
||||
{"exactly at limit", "hello", 5, "hello"},
|
||||
{"over limit", "hello world", 5, "hello…"},
|
||||
{"unicode over limit", "こんにちは世界", 5, "こんにちは…"},
|
||||
{"ascii over limit", "abcdefghij", 5, "abcde…"},
|
||||
}
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
got := broadcastTruncate(tc.s, tc.max)
|
||||
if got != tc.want {
|
||||
t.Errorf("broadcastTruncate(%q, %d) = %q; want %q", tc.s, tc.max, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Validation ────────────────────────────────────────────────────────────────
|
||||
|
||||
func TestBroadcast_InvalidWorkspaceID(t *testing.T) {
|
||||
c, w := buildBroadcastCtx("not-a-uuid", `{"message":"hello"}`)
|
||||
handler := NewBroadcastHandler(events.NewBroadcaster(ws.NewHub(nil)))
|
||||
handler.Broadcast(c)
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Errorf("want 400, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestBroadcast_MissingMessage(t *testing.T) {
|
||||
mock := setupBroadcastDB(t)
|
||||
c, w := buildBroadcastCtx(broadcastTestUUID, `{}`)
|
||||
|
||||
handler := NewBroadcastHandler(events.NewBroadcaster(ws.NewHub(nil)))
|
||||
handler.Broadcast(c)
|
||||
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Errorf("want 400, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet mock expectations: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBroadcast_MalformedJSON(t *testing.T) {
|
||||
mock := setupBroadcastDB(t)
|
||||
c, w := buildBroadcastCtx(broadcastTestUUID, `not json`)
|
||||
|
||||
handler := NewBroadcastHandler(events.NewBroadcaster(ws.NewHub(nil)))
|
||||
handler.Broadcast(c)
|
||||
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Errorf("want 400, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet mock expectations: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Auth / Authz ─────────────────────────────────────────────────────────────
|
||||
|
||||
func TestBroadcast_WorkspaceNotFound(t *testing.T) {
|
||||
mock := setupBroadcastDB(t)
|
||||
c, w := buildBroadcastCtx(broadcastTestUUID, `{"message":"hello"}`)
|
||||
|
||||
// Workspace lookup returns no rows.
|
||||
mock.ExpectQuery("SELECT name, broadcast_enabled FROM workspaces WHERE id = $1 AND status != 'removed'").
|
||||
WithArgs(broadcastTestUUID).
|
||||
WillReturnError(sql.ErrNoRows)
|
||||
|
||||
handler := NewBroadcastHandler(events.NewBroadcaster(ws.NewHub(nil)))
|
||||
handler.Broadcast(c)
|
||||
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Errorf("want 404, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet mock expectations: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBroadcast_WorkspaceLookupQueryError(t *testing.T) {
|
||||
mock := setupBroadcastDB(t)
|
||||
c, w := buildBroadcastCtx(broadcastTestUUID, `{"message":"hello"}`)
|
||||
|
||||
mock.ExpectQuery("SELECT name, broadcast_enabled FROM workspaces WHERE id = $1 AND status != 'removed'").
|
||||
WithArgs(broadcastTestUUID).
|
||||
WillReturnError(sql.ErrConnDone)
|
||||
|
||||
handler := NewBroadcastHandler(events.NewBroadcaster(ws.NewHub(nil)))
|
||||
handler.Broadcast(c)
|
||||
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Errorf("want 404, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet mock expectations: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBroadcast_BroadcastDisabled(t *testing.T) {
|
||||
mock := setupBroadcastDB(t)
|
||||
c, w := buildBroadcastCtx(broadcastTestUUID, `{"message":"hello"}`)
|
||||
|
||||
// Workspace found but broadcast_enabled=false.
|
||||
rows := sqlmock.NewRows([]string{"name", "broadcast_enabled"}).
|
||||
AddRow("test-workspace", false)
|
||||
mock.ExpectQuery("SELECT name, broadcast_enabled FROM workspaces WHERE id = $1 AND status != 'removed'").
|
||||
WithArgs(broadcastTestUUID).
|
||||
WillReturnRows(rows)
|
||||
|
||||
handler := NewBroadcastHandler(events.NewBroadcaster(ws.NewHub(nil)))
|
||||
handler.Broadcast(c)
|
||||
|
||||
if w.Code != http.StatusForbidden {
|
||||
t.Errorf("want 403, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet mock expectations: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── DB error paths ───────────────────────────────────────────────────────────
|
||||
|
||||
func TestBroadcast_RecipientQueryError(t *testing.T) {
|
||||
mock := setupBroadcastDB(t)
|
||||
c, w := buildBroadcastCtx(broadcastTestUUID, `{"message":"hello"}`)
|
||||
|
||||
// Workspace lookup succeeds with broadcast_enabled=true.
|
||||
rows := sqlmock.NewRows([]string{"name", "broadcast_enabled"}).
|
||||
AddRow("test-workspace", true)
|
||||
mock.ExpectQuery("SELECT name, broadcast_enabled FROM workspaces WHERE id = $1 AND status != 'removed'").
|
||||
WithArgs(broadcastTestUUID).
|
||||
WillReturnRows(rows)
|
||||
|
||||
// Recipient query fails.
|
||||
mock.ExpectQuery("SELECT id FROM workspaces WHERE status != 'removed' AND id != $1").
|
||||
WithArgs(broadcastTestUUID).
|
||||
WillReturnError(sql.ErrConnDone)
|
||||
|
||||
handler := NewBroadcastHandler(events.NewBroadcaster(ws.NewHub(nil)))
|
||||
handler.Broadcast(c)
|
||||
|
||||
if w.Code != http.StatusInternalServerError {
|
||||
t.Errorf("want 500, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet mock expectations: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBroadcast_RecipientRowsError(t *testing.T) {
|
||||
mock := setupBroadcastDB(t)
|
||||
c, w := buildBroadcastCtx(broadcastTestUUID, `{"message":"hello"}`)
|
||||
|
||||
rows := sqlmock.NewRows([]string{"name", "broadcast_enabled"}).
|
||||
AddRow("test-workspace", true)
|
||||
mock.ExpectQuery("SELECT name, broadcast_enabled FROM workspaces WHERE id = $1 AND status != 'removed'").
|
||||
WithArgs(broadcastTestUUID).
|
||||
WillReturnRows(rows)
|
||||
|
||||
// Recipient query succeeds but rows.Err() fails.
|
||||
badRows := sqlmock.NewRows([]string{"id"}).AddRow("ws-2").RowError(0, sql.ErrConnDone)
|
||||
mock.ExpectQuery("SELECT id FROM workspaces WHERE status != 'removed' AND id != $1").
|
||||
WithArgs(broadcastTestUUID).
|
||||
WillReturnRows(badRows)
|
||||
|
||||
handler := NewBroadcastHandler(events.NewBroadcaster(ws.NewHub(nil)))
|
||||
handler.Broadcast(c)
|
||||
|
||||
if w.Code != http.StatusInternalServerError {
|
||||
t.Errorf("want 500, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet mock expectations: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Success paths ───────────────────────────────────────────────────────────
|
||||
|
||||
func TestBroadcast_Success_OneRecipient(t *testing.T) {
|
||||
mock := setupBroadcastDB(t)
|
||||
c, w := buildBroadcastCtx(broadcastTestUUID, `{"message":"hello world"}`)
|
||||
|
||||
// Workspace lookup.
|
||||
wsRows := sqlmock.NewRows([]string{"name", "broadcast_enabled"}).
|
||||
AddRow("sender-workspace", true)
|
||||
mock.ExpectQuery("SELECT name, broadcast_enabled FROM workspaces WHERE id = $1 AND status != 'removed'").
|
||||
WithArgs(broadcastTestUUID).
|
||||
WillReturnRows(wsRows)
|
||||
|
||||
// Recipient query: one recipient.
|
||||
recipRows := sqlmock.NewRows([]string{"id"}).AddRow("ws-recipient-1")
|
||||
mock.ExpectQuery("SELECT id FROM workspaces WHERE status != 'removed' AND id != $1").
|
||||
WithArgs(broadcastTestUUID).
|
||||
WillReturnRows(recipRows)
|
||||
|
||||
// Activity log insert for recipient.
|
||||
mock.ExpectExec("INSERT INTO activity_logs (workspace_id, activity_type, method, source_id, summary, status) VALUES ($1, 'broadcast_receive', 'broadcast', $2, $3, 'ok')").
|
||||
WithArgs("ws-recipient-1", broadcastTestUUID, sqlmock.AnyArg()).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
// Activity log insert for sender (broadcast_sent).
|
||||
mock.ExpectExec("INSERT INTO activity_logs (workspace_id, activity_type, method, summary, status) VALUES ($1, 'broadcast_sent', 'broadcast', $2, 'ok')").
|
||||
WithArgs(broadcastTestUUID, sqlmock.AnyArg()).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
handler := NewBroadcastHandler(events.NewBroadcaster(ws.NewHub(nil)))
|
||||
handler.Broadcast(c)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Errorf("want 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet mock expectations: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBroadcast_Success_NoRecipients(t *testing.T) {
|
||||
mock := setupBroadcastDB(t)
|
||||
c, w := buildBroadcastCtx(broadcastTestUUID, `{"message":"hello"}`)
|
||||
|
||||
wsRows := sqlmock.NewRows([]string{"name", "broadcast_enabled"}).
|
||||
AddRow("solo-workspace", true)
|
||||
mock.ExpectQuery("SELECT name, broadcast_enabled FROM workspaces WHERE id = $1 AND status != 'removed'").
|
||||
WithArgs(broadcastTestUUID).
|
||||
WillReturnRows(wsRows)
|
||||
|
||||
// No recipients.
|
||||
recipRows := sqlmock.NewRows([]string{"id"})
|
||||
mock.ExpectQuery("SELECT id FROM workspaces WHERE status != 'removed' AND id != $1").
|
||||
WithArgs(broadcastTestUUID).
|
||||
WillReturnRows(recipRows)
|
||||
|
||||
// Activity log insert for sender (broadcast_sent).
|
||||
mock.ExpectExec("INSERT INTO activity_logs (workspace_id, activity_type, method, summary, status) VALUES ($1, 'broadcast_sent', 'broadcast', $2, 'ok')").
|
||||
WithArgs(broadcastTestUUID, sqlmock.AnyArg()).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
handler := NewBroadcastHandler(events.NewBroadcaster(ws.NewHub(nil)))
|
||||
handler.Broadcast(c)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Errorf("want 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet mock expectations: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBroadcast_Success_MultipleRecipients(t *testing.T) {
|
||||
mock := setupBroadcastDB(t)
|
||||
c, w := buildBroadcastCtx(broadcastTestUUID, `{"message":"hello"}`)
|
||||
|
||||
wsRows := sqlmock.NewRows([]string{"name", "broadcast_enabled"}).
|
||||
AddRow("broadcaster", true)
|
||||
mock.ExpectQuery("SELECT name, broadcast_enabled FROM workspaces WHERE id = $1 AND status != 'removed'").
|
||||
WithArgs(broadcastTestUUID).
|
||||
WillReturnRows(wsRows)
|
||||
|
||||
// Three recipients.
|
||||
recipRows := sqlmock.NewRows([]string{"id"}).
|
||||
AddRow("ws-1").AddRow("ws-2").AddRow("ws-3")
|
||||
mock.ExpectQuery("SELECT id FROM workspaces WHERE status != 'removed' AND id != $1").
|
||||
WithArgs(broadcastTestUUID).
|
||||
WillReturnRows(recipRows)
|
||||
|
||||
// Each recipient gets a broadcast_receive log.
|
||||
for _, rid := range []string{"ws-1", "ws-2", "ws-3"} {
|
||||
mock.ExpectExec("INSERT INTO activity_logs (workspace_id, activity_type, method, source_id, summary, status) VALUES ($1, 'broadcast_receive', 'broadcast', $2, $3, 'ok')").
|
||||
WithArgs(rid, broadcastTestUUID, sqlmock.AnyArg()).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
}
|
||||
|
||||
// Sender log.
|
||||
mock.ExpectExec("INSERT INTO activity_logs (workspace_id, activity_type, method, summary, status) VALUES ($1, 'broadcast_sent', 'broadcast', $2, 'ok')").
|
||||
WithArgs(broadcastTestUUID, sqlmock.AnyArg()).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
handler := NewBroadcastHandler(events.NewBroadcaster(ws.NewHub(nil)))
|
||||
handler.Broadcast(c)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Errorf("want 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet mock expectations: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Recipient insert failure (logged, continues) ─────────────────────────────
|
||||
|
||||
func TestBroadcast_RecipientInsertError_ContinuesAndSucceeds(t *testing.T) {
|
||||
mock := setupBroadcastDB(t)
|
||||
c, w := buildBroadcastCtx(broadcastTestUUID, `{"message":"hello"}`)
|
||||
|
||||
wsRows := sqlmock.NewRows([]string{"name", "broadcast_enabled"}).
|
||||
AddRow("broadcaster", true)
|
||||
mock.ExpectQuery("SELECT name, broadcast_enabled FROM workspaces WHERE id = $1 AND status != 'removed'").
|
||||
WithArgs(broadcastTestUUID).
|
||||
WillReturnRows(wsRows)
|
||||
|
||||
// Two recipients.
|
||||
recipRows := sqlmock.NewRows([]string{"id"}).
|
||||
AddRow("ws-1").AddRow("ws-2")
|
||||
mock.ExpectQuery("SELECT id FROM workspaces WHERE status != 'removed' AND id != $1").
|
||||
WithArgs(broadcastTestUUID).
|
||||
WillReturnRows(recipRows)
|
||||
|
||||
// First recipient insert fails (logged, continues).
|
||||
mock.ExpectExec("INSERT INTO activity_logs (workspace_id, activity_type, method, source_id, summary, status) VALUES ($1, 'broadcast_receive', 'broadcast', $2, $3, 'ok')").
|
||||
WithArgs("ws-1", broadcastTestUUID, sqlmock.AnyArg()).
|
||||
WillReturnError(sql.ErrConnDone)
|
||||
|
||||
// Second recipient insert succeeds.
|
||||
mock.ExpectExec("INSERT INTO activity_logs (workspace_id, activity_type, method, source_id, summary, status) VALUES ($1, 'broadcast_receive', 'broadcast', $2, $3, 'ok')").
|
||||
WithArgs("ws-2", broadcastTestUUID, sqlmock.AnyArg()).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
// Sender log.
|
||||
mock.ExpectExec("INSERT INTO activity_logs (workspace_id, activity_type, method, summary, status) VALUES ($1, 'broadcast_sent', 'broadcast', $2, 'ok')").
|
||||
WithArgs(broadcastTestUUID, sqlmock.AnyArg()).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
handler := NewBroadcastHandler(events.NewBroadcaster(ws.NewHub(nil)))
|
||||
handler.Broadcast(c)
|
||||
|
||||
// Handler returns 200 even though one insert failed — it logs and continues.
|
||||
if w.Code != http.StatusOK {
|
||||
t.Errorf("want 200 despite insert error, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet mock expectations: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Sender activity log insert failure (logged, still 200) ───────────────────
|
||||
|
||||
func TestBroadcast_SenderLogInsertError_Still200(t *testing.T) {
|
||||
mock := setupBroadcastDB(t)
|
||||
c, w := buildBroadcastCtx(broadcastTestUUID, `{"message":"hello"}`)
|
||||
|
||||
wsRows := sqlmock.NewRows([]string{"name", "broadcast_enabled"}).
|
||||
AddRow("broadcaster", true)
|
||||
mock.ExpectQuery("SELECT name, broadcast_enabled FROM workspaces WHERE id = $1 AND status != 'removed'").
|
||||
WithArgs(broadcastTestUUID).
|
||||
WillReturnRows(wsRows)
|
||||
|
||||
recipRows := sqlmock.NewRows([]string{"id"}).AddRow("ws-1")
|
||||
mock.ExpectQuery("SELECT id FROM workspaces WHERE status != 'removed' AND id != $1").
|
||||
WithArgs(broadcastTestUUID).
|
||||
WillReturnRows(recipRows)
|
||||
|
||||
mock.ExpectExec("INSERT INTO activity_logs (workspace_id, activity_type, method, source_id, summary, status) VALUES ($1, 'broadcast_receive', 'broadcast', $2, $3, 'ok')").
|
||||
WithArgs("ws-1", broadcastTestUUID, sqlmock.AnyArg()).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
// Sender log fails — but handler still returns 200 (logged only).
|
||||
mock.ExpectExec("INSERT INTO activity_logs (workspace_id, activity_type, method, summary, status) VALUES ($1, 'broadcast_sent', 'broadcast', $2, 'ok')").
|
||||
WithArgs(broadcastTestUUID, sqlmock.AnyArg()).
|
||||
WillReturnError(sql.ErrConnDone)
|
||||
|
||||
handler := NewBroadcastHandler(events.NewBroadcaster(ws.NewHub(nil)))
|
||||
handler.Broadcast(c)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Errorf("want 200 despite sender log error, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("unmet mock expectations: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -178,12 +178,21 @@ func (p *CPProvisioner) Start(ctx context.Context, cfg WorkspaceConfig) (string,
|
||||
// /admin/liveness and other admin-gated platform endpoints (core#831).
|
||||
// p.adminToken is read from os.Getenv("ADMIN_TOKEN") at provisioner creation;
|
||||
// it is also used for CP→platform HTTP auth but those are separate concerns.
|
||||
env := cfg.EnvVars
|
||||
if p.adminToken != "" {
|
||||
env = make(map[string]string, len(cfg.EnvVars)+1)
|
||||
for k, v := range cfg.EnvVars {
|
||||
env[k] = v
|
||||
//
|
||||
// Forensic #145 hardening: tenant workspaces run on EC2 via this path, so
|
||||
// the SCM-write-token denylist (see buildContainerEnv) is enforced here
|
||||
// too. Always build a filtered copy — never pass cfg.EnvVars through
|
||||
// verbatim — so a latent persona-merged GITEA_TOKEN can't reach the
|
||||
// tenant container regardless of whether ADMIN_TOKEN is set.
|
||||
env := make(map[string]string, len(cfg.EnvVars)+1)
|
||||
for k, v := range cfg.EnvVars {
|
||||
if isSCMWriteTokenKey(k) {
|
||||
log.Printf("CPProvisioner.Start: dropped SCM-write credential %q from tenant workspace env (forensic #145 guard)", k)
|
||||
continue
|
||||
}
|
||||
env[k] = v
|
||||
}
|
||||
if p.adminToken != "" {
|
||||
env["ADMIN_TOKEN"] = p.adminToken
|
||||
}
|
||||
// Collect template files and generated configs, with OFFSEC-010 guards:
|
||||
@@ -343,6 +352,7 @@ func collectCPConfigFiles(cfg WorkspaceConfig) (map[string]string, error) {
|
||||
}
|
||||
return files, nil
|
||||
}
|
||||
|
||||
// Stop terminates the workspace's EC2 instance via the control plane.
|
||||
//
|
||||
// Looks up the actual EC2 instance_id from the workspaces table before
|
||||
@@ -497,7 +507,9 @@ func (p *CPProvisioner) IsRunning(ctx context.Context, workspaceID string) (bool
|
||||
// Don't leak the body — upstream errors may echo headers.
|
||||
return true, fmt.Errorf("cp provisioner: status: unexpected %d", resp.StatusCode)
|
||||
}
|
||||
var result struct{ State string `json:"state"` }
|
||||
var result struct {
|
||||
State string `json:"state"`
|
||||
}
|
||||
// Cap body read at 64 KiB for parity with Start — a misconfigured
|
||||
// or compromised CP streaming a huge body could otherwise exhaust
|
||||
// memory in this hot path (called reactively per-request from
|
||||
|
||||
@@ -591,6 +591,28 @@ func ValidateWorkspaceAccess(access, workspacePath string) error {
|
||||
}
|
||||
}
|
||||
|
||||
// scmWriteTokenKeys is the explicit denylist of environment variable names
|
||||
// that carry a Git SCM *write* credential (push / merge / approve). These
|
||||
// must never reach a tenant workspace container — see the forensic #145
|
||||
// rationale in buildContainerEnv. Kept as an exact-match set rather than a
|
||||
// substring/prefix heuristic so the guard is auditable and can't silently
|
||||
// over-strip a legitimately-named var.
|
||||
var scmWriteTokenKeys = map[string]struct{}{
|
||||
"GITEA_TOKEN": {},
|
||||
"GITHUB_TOKEN": {},
|
||||
"GH_TOKEN": {}, // gh CLI honours GH_TOKEN as a GITHUB_TOKEN alias
|
||||
"GITLAB_TOKEN": {},
|
||||
"GL_TOKEN": {}, // glab CLI alias
|
||||
"BITBUCKET_TOKEN": {},
|
||||
}
|
||||
|
||||
// isSCMWriteTokenKey reports whether an env var name is a known Git SCM
|
||||
// write credential that must be stripped from tenant workspace env.
|
||||
func isSCMWriteTokenKey(key string) bool {
|
||||
_, ok := scmWriteTokenKeys[key]
|
||||
return ok
|
||||
}
|
||||
|
||||
// buildContainerEnv assembles the initial environment variables injected
|
||||
// into every workspace container.
|
||||
//
|
||||
@@ -627,6 +649,21 @@ func buildContainerEnv(cfg WorkspaceConfig) []string {
|
||||
env = append(env, fmt.Sprintf("AWARENESS_URL=%s", cfg.AwarenessURL))
|
||||
}
|
||||
for k, v := range cfg.EnvVars {
|
||||
// Forensic #145 hardening: tenant workspace containers run
|
||||
// agent-controlled code and must NEVER receive a Git SCM *write*
|
||||
// credential. Without merge/approve creds in-container the
|
||||
// two-eyes review gate is structurally self-bypass-proof — an
|
||||
// agent that forges an approval has no token to act on it. A
|
||||
// latent path exists (loadPersonaEnvFile merges a per-role
|
||||
// persona `GITEA_TOKEN` into cfg.EnvVars when MOLECULE_PERSONA_ROOT
|
||||
// is set on a tenant host); it is inert today (persona dirs are
|
||||
// operator-host-only) but unguarded. Strip SCM-write tokens here
|
||||
// by construction so the invariant holds regardless of whether
|
||||
// that path ever becomes reachable.
|
||||
if isSCMWriteTokenKey(k) {
|
||||
log.Printf("buildContainerEnv: dropped SCM-write credential %q from workspace env (forensic #145 guard)", k)
|
||||
continue
|
||||
}
|
||||
env = append(env, fmt.Sprintf("%s=%s", k, v))
|
||||
}
|
||||
// Inject ADMIN_TOKEN from the platform server's environment so workspace
|
||||
|
||||
@@ -636,10 +636,15 @@ func TestBuildContainerEnv_AwarenessOnlyWhenBothSet(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestBuildContainerEnv_CustomEnvVarsAppended(t *testing.T) {
|
||||
// NOTE: this test previously asserted GITHUB_TOKEN passed through
|
||||
// verbatim. That assertion encoded the forensic #145 latent leak as
|
||||
// expected behavior. Post-guard, ordinary custom env still flows but
|
||||
// SCM-write credentials are stripped — see
|
||||
// TestBuildContainerEnv_StripsSCMWriteTokens for the negative assertion.
|
||||
cfg := WorkspaceConfig{
|
||||
WorkspaceID: "ws-x",
|
||||
PlatformURL: "http://localhost:8080",
|
||||
EnvVars: map[string]string{"CUSTOM": "value", "GITHUB_TOKEN": "fake-token-for-test"},
|
||||
EnvVars: map[string]string{"CUSTOM": "value", "ANTHROPIC_API_KEY": "sk-not-an-scm-token"},
|
||||
}
|
||||
env := buildContainerEnv(cfg)
|
||||
seen := map[string]string{}
|
||||
@@ -652,8 +657,8 @@ func TestBuildContainerEnv_CustomEnvVarsAppended(t *testing.T) {
|
||||
if seen["CUSTOM"] != "value" {
|
||||
t.Errorf("CUSTOM env missing, got env=%v", env)
|
||||
}
|
||||
if seen["GITHUB_TOKEN"] != "fake-token-for-test" {
|
||||
t.Errorf("GITHUB_TOKEN env missing, got env=%v", env)
|
||||
if seen["ANTHROPIC_API_KEY"] != "sk-not-an-scm-token" {
|
||||
t.Errorf("non-SCM custom env must still pass through, got env=%v", env)
|
||||
}
|
||||
// Built-in defaults still present
|
||||
if seen["MOLECULE_URL"] == "" {
|
||||
@@ -661,6 +666,129 @@ func TestBuildContainerEnv_CustomEnvVarsAppended(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// ---------- forensic #145: SCM-write-token denylist guard ----------
|
||||
|
||||
// TestBuildContainerEnv_StripsSCMWriteTokens is the core negative
|
||||
// assertion: a tenant workspace env constructed via buildContainerEnv MUST
|
||||
// NOT contain any Git SCM *write* credential, regardless of how it got into
|
||||
// cfg.EnvVars. This proves the two-eyes review gate stays structurally
|
||||
// self-bypass-proof — an agent in-container has no merge/approve token to
|
||||
// act on a forged approval. See forensic #145.
|
||||
//
|
||||
// This test FAILS on the pre-guard code (where buildContainerEnv passed
|
||||
// cfg.EnvVars through verbatim) and PASSES once the denylist filter is in
|
||||
// place — i.e. the guard is proven by construction, not by environment
|
||||
// accident.
|
||||
func TestBuildContainerEnv_StripsSCMWriteTokens(t *testing.T) {
|
||||
scmTokens := []string{
|
||||
"GITEA_TOKEN", "GITHUB_TOKEN", "GH_TOKEN",
|
||||
"GITLAB_TOKEN", "GL_TOKEN", "BITBUCKET_TOKEN",
|
||||
}
|
||||
|
||||
t.Run("normal path — SCM tokens explicitly set in EnvVars", func(t *testing.T) {
|
||||
envVars := map[string]string{"CUSTOM": "ok", "ANTHROPIC_API_KEY": "sk-keep"}
|
||||
for _, k := range scmTokens {
|
||||
envVars[k] = "leaked-write-credential-" + k
|
||||
}
|
||||
cfg := WorkspaceConfig{
|
||||
WorkspaceID: "ws-tenant",
|
||||
PlatformURL: "http://localhost:8080",
|
||||
Tier: 2,
|
||||
EnvVars: envVars,
|
||||
}
|
||||
assertNoSCMWriteToken(t, buildContainerEnv(cfg), scmTokens)
|
||||
|
||||
// Sanity: non-SCM custom env is NOT collateral-damaged by the filter.
|
||||
if !envContains(buildContainerEnv(cfg), "CUSTOM=ok") {
|
||||
t.Errorf("filter must not strip non-SCM custom env")
|
||||
}
|
||||
if !envContains(buildContainerEnv(cfg), "ANTHROPIC_API_KEY=sk-keep") {
|
||||
t.Errorf("filter must not strip non-SCM API keys")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("persona-file path — simulates loadPersonaEnvFile merge", func(t *testing.T) {
|
||||
// The latent path: handlers.loadPersonaEnvFile() merges a per-role
|
||||
// persona env file (carrying GITEA_USER, GITEA_TOKEN, …) into the
|
||||
// workspace env map when MOLECULE_PERSONA_ROOT is set on a tenant
|
||||
// host. We can't invoke that cross-package helper here, but its
|
||||
// observable effect is exactly "a GITEA_TOKEN appears in
|
||||
// cfg.EnvVars". Constructing that condition directly proves the
|
||||
// guard holds even if the latent path becomes reachable.
|
||||
cfg := WorkspaceConfig{
|
||||
WorkspaceID: "ws-tenant",
|
||||
PlatformURL: "http://localhost:8080",
|
||||
Tier: 2,
|
||||
EnvVars: map[string]string{
|
||||
// Persona identity fields that are SAFE to keep (read-only
|
||||
// identity, not a write credential):
|
||||
"GITEA_USER": "backend-engineer",
|
||||
"GITEA_USER_EMAIL": "backend-engineer@agents.moleculesai.app",
|
||||
// The credential that must be stripped:
|
||||
"GITEA_TOKEN": "persona-merged-write-pat",
|
||||
"GITEA_TOKEN_SCOPES": "write:repository",
|
||||
},
|
||||
}
|
||||
got := buildContainerEnv(cfg)
|
||||
assertNoSCMWriteToken(t, got, scmTokens)
|
||||
// Non-credential persona identity may still flow through — only the
|
||||
// write token is the denied surface.
|
||||
if !envContains(got, "GITEA_USER=backend-engineer") {
|
||||
t.Errorf("non-credential persona identity (GITEA_USER) should not be stripped")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestCPProvisionerEnv_StripsSCMWriteTokens covers the tenant-EC2 path:
|
||||
// CPProvisioner.Start builds the env map the control plane forwards to the
|
||||
// EC2 workspace container. The same forensic #145 denylist must hold there.
|
||||
func TestCPProvisionerEnv_StripsSCMWriteTokens(t *testing.T) {
|
||||
// isSCMWriteTokenKey is the single source of truth shared by both
|
||||
// buildContainerEnv (local Docker) and CPProvisioner.Start (tenant EC2).
|
||||
// Assert it classifies every known SCM-write var as denied and leaves
|
||||
// ordinary / read-only-identity vars alone.
|
||||
for _, k := range []string{
|
||||
"GITEA_TOKEN", "GITHUB_TOKEN", "GH_TOKEN",
|
||||
"GITLAB_TOKEN", "GL_TOKEN", "BITBUCKET_TOKEN",
|
||||
} {
|
||||
if !isSCMWriteTokenKey(k) {
|
||||
t.Errorf("isSCMWriteTokenKey(%q) = false, want true (SCM-write credential must be denied)", k)
|
||||
}
|
||||
}
|
||||
for _, k := range []string{
|
||||
"GITEA_USER", "GITEA_USER_EMAIL", "ANTHROPIC_API_KEY",
|
||||
"CUSTOM", "PLATFORM_URL", "ADMIN_TOKEN", "",
|
||||
} {
|
||||
if isSCMWriteTokenKey(k) {
|
||||
t.Errorf("isSCMWriteTokenKey(%q) = true, want false (must not over-strip non-SCM env)", k)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func assertNoSCMWriteToken(t *testing.T, env []string, scmTokens []string) {
|
||||
t.Helper()
|
||||
for _, e := range env {
|
||||
key := e
|
||||
if i := strings.IndexByte(e, '='); i >= 0 {
|
||||
key = e[:i]
|
||||
}
|
||||
for _, banned := range scmTokens {
|
||||
if key == banned {
|
||||
t.Errorf("SCM-write credential %q leaked into workspace env (forensic #145 invariant violated): %q", banned, e)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func envContains(env []string, want string) bool {
|
||||
for _, e := range env {
|
||||
if e == want {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// ---------- buildWorkspaceMount — #65 workspace_access ----------
|
||||
|
||||
func TestBuildWorkspaceMount_SelectionMatrix(t *testing.T) {
|
||||
|
||||
Reference in New Issue
Block a user