431 lines
16 KiB
Go
431 lines
16 KiB
Go
package server
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"github.com/golang-jwt/jwt/v5"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
"vdi-cloud/platform/internal/model"
|
|
"vdi-cloud/platform/internal/store"
|
|
"vdi-cloud/platform/pkg/proxmox"
|
|
)
|
|
|
|
type batchPVE struct {
|
|
mu sync.Mutex
|
|
state, outcome string
|
|
locked, deny bool
|
|
writes int
|
|
}
|
|
|
|
func fakeBatchPVE(t *testing.T, f *fixture) *batchPVE {
|
|
p := &batchPVE{state: "stopped"}
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
path := strings.TrimPrefix(r.URL.Path, "/api2/json")
|
|
var data any
|
|
switch {
|
|
case strings.Contains(path, "/tasks/"):
|
|
if p.outcome == "poll-failure" {
|
|
w.WriteHeader(503)
|
|
return
|
|
}
|
|
status, exit := "stopped", "OK"
|
|
if p.outcome == "running" {
|
|
status = "running"
|
|
}
|
|
if p.outcome == "failed" {
|
|
exit = "TASK ERROR"
|
|
}
|
|
data = map[string]any{"status": status, "exitstatus": exit}
|
|
case path == "/access/permissions":
|
|
n := 1
|
|
if p.deny {
|
|
n = 0
|
|
}
|
|
data = map[string]any{r.URL.Query().Get("path"): map[string]int{"VM.PowerMgmt": n, "VM.Snapshot": n, "VM.Config.CPU": n, "VM.Config.Memory": n}}
|
|
case strings.HasSuffix(path, "/status/current"):
|
|
data = map[string]string{"status": p.state}
|
|
case strings.HasSuffix(path, "/config"):
|
|
c := map[string]any{"digest": "d1", "cores": 2, "memory": 2048, "scsi0": "local:vm-9000,size=4G"}
|
|
if p.locked {
|
|
c["lock"] = "backup"
|
|
}
|
|
data = c
|
|
case strings.HasSuffix(path, "/pending"):
|
|
data = []any{}
|
|
case r.Method == "POST" && strings.Contains(path, "/status/"):
|
|
p.writes++
|
|
if p.outcome == "reject" {
|
|
w.WriteHeader(403)
|
|
return
|
|
}
|
|
if p.outcome == "ambiguous" {
|
|
w.WriteHeader(500)
|
|
return
|
|
}
|
|
p.state = "running"
|
|
if strings.HasSuffix(path, "/shutdown") {
|
|
p.state = "stopped"
|
|
}
|
|
data = "UPID:batch-power"
|
|
if p.outcome == "missing-upid" {
|
|
data = ""
|
|
}
|
|
default:
|
|
t.Error("unexpected batch PVE call", r.Method, path)
|
|
http.NotFound(w, r)
|
|
return
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"data": data})
|
|
}))
|
|
t.Cleanup(srv.Close)
|
|
f.server.PVE = &proxmox.Client{BaseURL: srv.URL, HTTP: srv.Client()}
|
|
if _, err := f.db.Exec("UPDATE desktops SET provider='proxmox',node='node-a',vmid=9000,power_state='stopped' WHERE id='desk_demo_windows'"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return p
|
|
}
|
|
func batchRequest(action string, ids ...string) map[string]any {
|
|
return map[string]any{"action": action, "desktopIds": ids}
|
|
}
|
|
func checkedBatch(t *testing.T, f *fixture, token string, body map[string]any) map[string]any {
|
|
t.Helper()
|
|
code, out := f.request("POST", "/desktop-batches/check", token, body, "")
|
|
expect(t, 200, code, out)
|
|
b := copySnapshotConfig(body)
|
|
b["previewToken"] = out["previewToken"]
|
|
return b
|
|
}
|
|
func queueBatchTest(t *testing.T, f *fixture, token string, body map[string]any, key string) map[string]any {
|
|
t.Helper()
|
|
code, out := f.request("POST", "/desktop-batches", token, body, key)
|
|
expect(t, 202, code, out)
|
|
return out
|
|
}
|
|
func batchTask(t *testing.T, f *fixture, out map[string]any) model.Task {
|
|
t.Helper()
|
|
for _, value := range out["items"].([]any) {
|
|
item := value.(map[string]any)
|
|
if id, _ := item["taskId"].(string); id != "" {
|
|
return readSnapshotTask(t, f, id)
|
|
}
|
|
}
|
|
t.Fatal("no batch task", out)
|
|
return model.Task{}
|
|
}
|
|
func TestPowerBatchLifecycleAndPartialResults(t *testing.T) {
|
|
for _, action := range []string{"start", "shutdown", "reboot"} {
|
|
t.Run(action, func(t *testing.T) {
|
|
f := newFixtureWorker(t, false)
|
|
p := fakeBatchPVE(t, f)
|
|
if action != "start" {
|
|
p.state = "running"
|
|
}
|
|
operator := f.login("operator")
|
|
body := checkedBatch(t, f, operator, batchRequest(action, "desk_demo_windows", "not-managed"))
|
|
out := queueBatchTest(t, f, operator, body, "batch-once")
|
|
batch := out["batch"].(map[string]any)
|
|
if batch["submittedCount"] != float64(1) || batch["skippedCount"] != float64(1) || batch["activeCount"] != float64(1) || batch["successCount"] != float64(0) {
|
|
t.Fatal(batch)
|
|
}
|
|
task := batchTask(t, f, out)
|
|
f.server.execute(context.Background(), task)
|
|
task = readSnapshotTask(t, f, task.ID)
|
|
if task.Status != "success" || p.writes != 1 || task.ProviderRef == "" {
|
|
t.Fatal(task, p.writes)
|
|
}
|
|
code, view := f.request("GET", "/desktop-batches/"+batch["id"].(string), f.login("viewer"), nil, "")
|
|
expect(t, 200, code, view)
|
|
counts := view["batch"].(map[string]any)
|
|
if counts["successCount"] != float64(1) || counts["activeCount"] != float64(0) {
|
|
t.Fatal(counts)
|
|
}
|
|
code, view = f.request("GET", "/desktop-batches", operator, nil, "")
|
|
expect(t, 200, code, view)
|
|
if view["total"] != float64(1) {
|
|
t.Fatal(view)
|
|
}
|
|
// A retry after completion, even without a preview token, returns the same batch.
|
|
delete(body, "previewToken")
|
|
again := queueBatchTest(t, f, operator, body, "batch-once")
|
|
if again["batch"].(map[string]any)["id"] != batch["id"] {
|
|
t.Fatal("duplicate batch")
|
|
}
|
|
body["action"] = "start"
|
|
if action == "start" {
|
|
body["action"] = "shutdown"
|
|
}
|
|
code, view = f.request("POST", "/desktop-batches", operator, body, "batch-once")
|
|
expect(t, 409, code, view)
|
|
var submitted, completed int
|
|
_ = f.db.Get(&submitted, "SELECT count(*) FROM audit_logs WHERE result='submitted'")
|
|
_ = f.db.Get(&completed, "SELECT count(*) FROM audit_logs WHERE result='success' AND action=$1", action)
|
|
if submitted != 2 || completed != 1 {
|
|
t.Fatal("audit missing", submitted, completed)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
func TestPowerBatchPreviewValidationAndRBAC(t *testing.T) {
|
|
f := newFixtureWorker(t, false)
|
|
fakeBatchPVE(t, f)
|
|
admin := f.login("admin")
|
|
viewer := f.login("viewer")
|
|
operator := f.login("operator")
|
|
for _, path := range []string{"/desktop-batches/check", "/desktop-batches"} {
|
|
code, out := f.request("POST", path, viewer, batchRequest("start", "desk_demo_windows"), "no")
|
|
expect(t, 403, code, out)
|
|
}
|
|
for _, body := range []map[string]any{batchRequest("reset", "desk_demo_windows"), batchRequest("start"), batchRequest("start", "desk_demo_windows", "desk_demo_windows"), batchRequest("start", " desk_demo_windows"), {"action": "start", "desktopIds": []string{"desk_demo_windows"}, "force": true}} {
|
|
code, out := f.request("POST", "/desktop-batches/check", admin, body, "")
|
|
expect(t, 400, code, out)
|
|
}
|
|
ids := []string{}
|
|
for i := 0; i < 51; i++ {
|
|
ids = append(ids, store.ID("d_"))
|
|
}
|
|
code, out := f.request("POST", "/desktop-batches/check", admin, batchRequest("start", ids...), "")
|
|
expect(t, 400, code, out)
|
|
body := checkedBatch(t, f, admin, batchRequest("start", "desk_demo_windows"))
|
|
preview := body["previewToken"].(string)
|
|
code, out = f.request("GET", "/auth/me", preview, nil, "")
|
|
expect(t, 401, code, out)
|
|
code, out = f.request("POST", "/desktop-batches", admin, body, "")
|
|
expect(t, 400, code, out)
|
|
code, out = f.request("POST", "/desktop-batches", operator, body, "wrong-user")
|
|
expect(t, 409, code, out)
|
|
body["action"] = "shutdown"
|
|
code, out = f.request("POST", "/desktop-batches", admin, body, "changed")
|
|
expect(t, 409, code, out)
|
|
body["action"] = "start"
|
|
body["previewToken"] = admin
|
|
code, out = f.request("POST", "/desktop-batches", admin, body, "access-not-preview")
|
|
expect(t, 409, code, out)
|
|
cl := &powerBatchClaims{}
|
|
_, err := jwt.ParseWithClaims(preview, cl, func(*jwt.Token) (any, error) { return []byte(f.server.Config.JWTSecret), nil })
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
cl.ExpiresAt = jwt.NewNumericDate(time.Now().Add(-time.Minute))
|
|
expired, _ := jwt.NewWithClaims(jwt.SigningMethodHS256, cl).SignedString([]byte(f.server.Config.JWTSecret))
|
|
body["previewToken"] = expired
|
|
code, out = f.request("POST", "/desktop-batches", admin, body, "expired")
|
|
expect(t, 409, code, out)
|
|
body["previewToken"] = preview
|
|
secondSession := f.login("admin")
|
|
code, out = f.request("POST", "/desktop-batches", secondSession, body, "wrong-session")
|
|
expect(t, 409, code, out)
|
|
var n int
|
|
_ = f.db.Get(&n, "SELECT count(*) FROM power_batches")
|
|
if n != 0 {
|
|
t.Fatal("invalid request persisted batch")
|
|
}
|
|
}
|
|
func TestPowerBatchPreflightBlockersAndRecheck(t *testing.T) {
|
|
f := newFixtureWorker(t, false)
|
|
p := fakeBatchPVE(t, f)
|
|
admin := f.login("admin")
|
|
for _, block := range []string{"running", "lock", "deny", "active-session", "maintenance", "unknown"} {
|
|
t.Run(block, func(t *testing.T) {
|
|
p.state = "stopped"
|
|
p.locked = false
|
|
p.deny = false
|
|
switch block {
|
|
case "running":
|
|
p.state = "running"
|
|
case "lock":
|
|
p.locked = true
|
|
case "deny":
|
|
p.deny = true
|
|
case "active-session":
|
|
_, _ = f.db.Exec("UPDATE terminals SET session_desktop_id='desk_demo_windows',last_seen=now() WHERE id='term_demo_001'")
|
|
case "maintenance":
|
|
_, _ = f.db.Exec("INSERT INTO tasks(id,kind,action,target_id,target_name,actor,status,idempotency_key) VALUES('block','desktop','start','desk_demo_windows','blocked','admin','queued','block')")
|
|
case "unknown":
|
|
_, _ = f.db.Exec("INSERT INTO tasks(id,kind,action,target_id,target_name,actor,status,idempotency_key) VALUES('block','desktop','start','desk_demo_windows','blocked','admin','timeout','block')")
|
|
}
|
|
code, out := f.request("POST", "/desktop-batches/check", admin, batchRequest("start", "desk_demo_windows"), "")
|
|
expect(t, 200, code, out)
|
|
if out["eligibleCount"] != float64(0) {
|
|
t.Fatal(block, out)
|
|
}
|
|
body := batchRequest("start", "desk_demo_windows")
|
|
body["previewToken"] = out["previewToken"]
|
|
code, out = f.request("POST", "/desktop-batches", admin, body, "blocked-"+block)
|
|
expect(t, 409, code, out)
|
|
_, _ = f.db.Exec("DELETE FROM tasks WHERE id='block'")
|
|
_, _ = f.db.Exec("UPDATE terminals SET session_desktop_id=NULL WHERE id='term_demo_001'")
|
|
})
|
|
}
|
|
p.state = "stopped"
|
|
p.locked = false
|
|
p.deny = false
|
|
// A desktop that was excluded from the preview cannot become a newly queued target.
|
|
body := checkedBatch(t, f, admin, batchRequest("start", "desk_demo_windows", "desk_demo_ubuntu"))
|
|
_, _ = f.db.Exec("UPDATE desktops SET power_state='stopped' WHERE id='desk_demo_ubuntu'")
|
|
out := queueBatchTest(t, f, admin, body, "preview-fence")
|
|
if out["batch"].(map[string]any)["submittedCount"] != float64(1) {
|
|
t.Fatal("added unapproved target", out)
|
|
}
|
|
task := batchTask(t, f, out)
|
|
p.state = "running"
|
|
f.server.execute(context.Background(), task)
|
|
if readSnapshotTask(t, f, task.ID).Status != "failed" || p.writes != 0 {
|
|
t.Fatal("worker did not recheck before send")
|
|
}
|
|
// Queueing rechecks newly active sessions and records a skipped row, never sends.
|
|
p.state = "stopped"
|
|
body = checkedBatch(t, f, admin, batchRequest("start", "desk_demo_windows"))
|
|
_, _ = f.db.Exec("UPDATE terminals SET session_desktop_id='desk_demo_windows',last_seen=now() WHERE id='term_demo_001'")
|
|
out = queueBatchTest(t, f, admin, body, "session-after-preview")
|
|
if out["batch"].(map[string]any)["submittedCount"] != float64(0) {
|
|
t.Fatal(out)
|
|
}
|
|
}
|
|
func TestPowerBatchUnknownAndMaintenanceProtection(t *testing.T) {
|
|
for _, outcome := range []string{"reject", "ambiguous", "missing-upid", "poll-failure", "failed"} {
|
|
t.Run(outcome, func(t *testing.T) {
|
|
f := newFixtureWorker(t, false)
|
|
p := fakeBatchPVE(t, f)
|
|
admin := f.login("admin")
|
|
body := checkedBatch(t, f, admin, batchRequest("start", "desk_demo_windows"))
|
|
out := queueBatchTest(t, f, admin, body, "uncertain")
|
|
task := batchTask(t, f, out)
|
|
p.outcome = outcome
|
|
f.server.execute(context.Background(), task)
|
|
task = readSnapshotTask(t, f, task.ID)
|
|
want := "timeout"
|
|
if outcome == "reject" {
|
|
want = "failed"
|
|
}
|
|
if task.Status != want {
|
|
t.Fatal(task)
|
|
}
|
|
if want == "failed" {
|
|
return
|
|
}
|
|
code, view := f.request("POST", "/desktops/desk_demo_windows/actions", admin, map[string]string{"action": "start"}, "blocked-single")
|
|
expect(t, 409, code, view)
|
|
_, _ = f.db.Exec("UPDATE desktops SET power_state='running' WHERE id='desk_demo_windows'")
|
|
_, _ = f.db.Exec("UPDATE terminals SET token_hash=$1,last_seen=now() WHERE id='term_demo_001'", store.Hash("batch-device"))
|
|
code, view = f.request("POST", "/terminal/connect-ticket", "batch-device", map[string]string{"desktopId": "desk_demo_windows"}, "")
|
|
expect(t, 409, code, view)
|
|
if view["code"] != "hardware_unconfirmed" {
|
|
t.Fatal(view)
|
|
}
|
|
code, view = f.request("PUT", "/desktops/desk_demo_windows/connection", admin, map[string]any{"endpoint": "192.0.2.90", "width": 1920, "height": 1080, "fps": 30, "initialized": true}, "")
|
|
expect(t, 409, code, view)
|
|
_, _ = f.db.Exec("UPDATE tasks SET status='running' WHERE id=$1", task.ID)
|
|
f.server.execute(context.Background(), task)
|
|
task = readSnapshotTask(t, f, task.ID)
|
|
if p.writes != 1 {
|
|
t.Fatal("power replayed", p.writes)
|
|
}
|
|
code, view = f.request("POST", "/tasks/"+task.ID+"/batch-power-reconcile", admin, nil, "")
|
|
expect(t, 409, code, view)
|
|
p.outcome = ""
|
|
f.db.SetMaxOpenConns(1)
|
|
code, view = f.request("POST", "/tasks/"+task.ID+"/batch-power-reconcile", admin, nil, "")
|
|
if outcome == "poll-failure" || outcome == "failed" {
|
|
expect(t, 200, code, view)
|
|
} else {
|
|
expect(t, 409, code, view)
|
|
}
|
|
if p.writes != 1 {
|
|
t.Fatal("read-only reconcile wrote PVE")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
func TestPowerBatchConcurrentReplayAndExclusivity(t *testing.T) {
|
|
f := newFixtureWorker(t, false)
|
|
fakeBatchPVE(t, f)
|
|
admin := f.login("admin")
|
|
body := checkedBatch(t, f, admin, batchRequest("start", "desk_demo_windows"))
|
|
var wg sync.WaitGroup
|
|
results := make(chan string, 5)
|
|
for i := 0; i < 5; i++ {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
out := queueBatchTest(t, f, admin, body, "concurrent")
|
|
results <- out["batch"].(map[string]any)["id"].(string)
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
close(results)
|
|
first := ""
|
|
for id := range results {
|
|
if first == "" {
|
|
first = id
|
|
}
|
|
if id != first {
|
|
t.Fatal("duplicate batch")
|
|
}
|
|
}
|
|
code, out := f.request("POST", "/desktops/desk_demo_windows/actions", admin, map[string]string{"action": "start"}, "single-conflict")
|
|
expect(t, 409, code, out)
|
|
code, out = f.request("POST", "/desktops/desk_demo_windows/hardware", admin, map[string]any{"kind": "compute", "cores": 4, "sockets": 1, "memoryMiB": 4096, "digest": "d1"}, "hardware-conflict")
|
|
expect(t, 409, code, out)
|
|
var n int
|
|
_ = f.db.Get(&n, "SELECT count(*) FROM tasks WHERE kind='batch-desktop'")
|
|
if n != 1 {
|
|
t.Fatal(n)
|
|
}
|
|
}
|
|
func TestPowerBatchExpiryAndOrdinaryLane(t *testing.T) {
|
|
f := newFixtureWorker(t, false)
|
|
p := fakeBatchPVE(t, f)
|
|
admin := f.login("admin")
|
|
body := checkedBatch(t, f, admin, batchRequest("start", "desk_demo_windows"))
|
|
out := queueBatchTest(t, f, admin, body, "prepared-expire")
|
|
task := batchTask(t, f, out)
|
|
f.server.finishBatchPower(task, "", "expired", "")
|
|
claimed, err := f.server.claimBatchPower(context.Background(), task)
|
|
if err != nil || claimed {
|
|
t.Fatal("expired task claimed", err)
|
|
}
|
|
out = queueBatchTest(t, f, admin, body, "submitting-expire")
|
|
task = batchTask(t, f, out)
|
|
claimed, err = f.server.claimBatchPower(context.Background(), task)
|
|
if err != nil || !claimed {
|
|
t.Fatal(err)
|
|
}
|
|
f.server.finishBatchPower(task, "", "expired", "")
|
|
if readSnapshotTask(t, f, task.ID).Status != "timeout" {
|
|
t.Fatal("submitted expiry not protected")
|
|
}
|
|
// Clear only the isolated test job, then verify a queued batch older than 2 minutes remains eligible.
|
|
_, _ = f.db.Exec("UPDATE power_batch_items SET stage='failed' WHERE task_id=$1", task.ID)
|
|
_, _ = f.db.Exec("UPDATE tasks SET status='failed' WHERE id=$1", task.ID)
|
|
out = queueBatchTest(t, f, admin, body, "long")
|
|
task = batchTask(t, f, out)
|
|
_, _ = f.db.Exec("UPDATE tasks SET created_at=now()-interval '3 minutes' WHERE id=$1", task.ID)
|
|
p.outcome = "running"
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
done := make(chan struct{})
|
|
go func() { f.server.RunWorker(ctx); close(done) }()
|
|
defer func() { cancel(); <-done }()
|
|
deadline := time.Now().Add(6 * time.Second)
|
|
for readSnapshotTask(t, f, task.ID).ProviderRef == "" {
|
|
if time.Now().After(deadline) {
|
|
t.Fatal("batch failed to start")
|
|
}
|
|
time.Sleep(25 * time.Millisecond)
|
|
}
|
|
code, view := f.request("POST", "/desktops/desk_demo_kylin/actions", admin, map[string]string{"action": "start"}, "ordinary")
|
|
expect(t, 202, code, view)
|
|
for readSnapshotTask(t, f, view["id"].(string)).Status != "success" {
|
|
if time.Now().After(deadline) {
|
|
t.Fatal("batch blocked normal task lane")
|
|
}
|
|
time.Sleep(25 * time.Millisecond)
|
|
}
|
|
}
|