207 lines
7.0 KiB
Go
207 lines
7.0 KiB
Go
package controller
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"path/filepath"
|
|
"strconv"
|
|
"testing"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
|
|
"github.com/mhsanaei/3x-ui/v3/internal/database"
|
|
"github.com/mhsanaei/3x-ui/v3/internal/database/model"
|
|
)
|
|
|
|
// Item 19: agent tokens are kept recoverably so the node edit page can
|
|
// re-display the full install command. Rows created before plaintext
|
|
// storage existed are backfilled from the first request that passes the
|
|
// token hash check — and only from such a request.
|
|
|
|
func setupBackfillTest(t *testing.T) (*gin.Engine, *model.Node) {
|
|
t.Helper()
|
|
gin.SetMode(gin.TestMode)
|
|
dbDir := t.TempDir()
|
|
t.Setenv("XUI_DB_FOLDER", dbDir)
|
|
if err := database.InitDB(filepath.Join(dbDir, "x-ui.db")); err != nil {
|
|
t.Fatalf("InitDB: %v", err)
|
|
}
|
|
t.Cleanup(func() { _ = database.CloseDB() })
|
|
|
|
node := &model.Node{Name: "backfill-node", Address: "10.0.9.2"}
|
|
if err := database.GetDB().Create(node).Error; err != nil {
|
|
t.Fatalf("seed node: %v", err)
|
|
}
|
|
engine := gin.New()
|
|
ctrl := NewAgentController(engine.Group("/panel/api/agent"))
|
|
ctrl.RegisterAdminRoutes(engine.Group("/panel/api/agent"))
|
|
return engine, node
|
|
}
|
|
|
|
func seedLegacyAgent(t *testing.T, nodeID int, token, plain string) *model.AgentNode {
|
|
t.Helper()
|
|
agent := &model.AgentNode{
|
|
TokenHash: hashAgentToken(token),
|
|
TokenPrefix: token[:12],
|
|
TokenPlain: plain,
|
|
Name: "backfill-node",
|
|
NodeID: &nodeID,
|
|
}
|
|
if err := database.GetDB().Create(agent).Error; err != nil {
|
|
t.Fatalf("seed agent: %v", err)
|
|
}
|
|
return agent
|
|
}
|
|
|
|
func reloadAgent(t *testing.T, id int) model.AgentNode {
|
|
t.Helper()
|
|
var row model.AgentNode
|
|
if err := database.GetDB().First(&row, id).Error; err != nil {
|
|
t.Fatalf("reload agent: %v", err)
|
|
}
|
|
return row
|
|
}
|
|
|
|
func heartbeatWith(engine *gin.Engine, token string) int {
|
|
req := httptest.NewRequest(http.MethodPost, "/panel/api/agent/heartbeat", bytes.NewBufferString(`{}`))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Authorization", "Bearer "+token)
|
|
rec := httptest.NewRecorder()
|
|
engine.ServeHTTP(rec, req)
|
|
return rec.Code
|
|
}
|
|
|
|
func TestIssueAgentTokenStoresPlaintext(t *testing.T) {
|
|
_, node := setupBackfillTest(t)
|
|
token, id, err := issueAgentToken(&node.Id, "backfill-node")
|
|
if err != nil {
|
|
t.Fatalf("issueAgentToken: %v", err)
|
|
}
|
|
row := reloadAgent(t, id)
|
|
if row.TokenPlain != token {
|
|
t.Fatalf("issued token not stored recoverably: TokenPlain=%q", row.TokenPlain)
|
|
}
|
|
if row.TokenHash != hashAgentToken(token) {
|
|
t.Fatal("issued token hash mismatch")
|
|
}
|
|
}
|
|
|
|
func TestAgentTokenBackfillOnAuthenticatedRequest(t *testing.T) {
|
|
engine, node := setupBackfillTest(t)
|
|
const token = "legacy-token-backfill-00000000000000000001"
|
|
agent := seedLegacyAgent(t, node.Id, token, "")
|
|
|
|
if code := heartbeatWith(engine, token); code != http.StatusOK {
|
|
t.Fatalf("heartbeat with valid token: status %d", code)
|
|
}
|
|
if row := reloadAgent(t, agent.Id); row.TokenPlain != token {
|
|
t.Fatalf("token not backfilled after authenticated request: %q", row.TokenPlain)
|
|
}
|
|
|
|
// The install endpoint now serves the complete command material.
|
|
req := httptest.NewRequest(http.MethodGet, "/panel/api/agent/install/"+strconv.Itoa(node.Id), nil)
|
|
rec := httptest.NewRecorder()
|
|
engine.ServeHTTP(rec, req)
|
|
if rec.Code != http.StatusOK {
|
|
t.Fatalf("install endpoint: status %d", rec.Code)
|
|
}
|
|
var resp struct {
|
|
Success bool `json:"success"`
|
|
Obj struct {
|
|
Name string `json:"name"`
|
|
Token string `json:"token"`
|
|
HasToken bool `json:"hasToken"`
|
|
} `json:"obj"`
|
|
}
|
|
if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
|
|
t.Fatalf("install response not JSON: %v", err)
|
|
}
|
|
if !resp.Success || !resp.Obj.HasToken || resp.Obj.Token != token {
|
|
t.Fatalf("install endpoint did not return the backfilled token: %+v", resp.Obj)
|
|
}
|
|
}
|
|
|
|
// Regression for the production failure of the first V1.0.0 deployment:
|
|
// rows that existed before item 19 got token_plain from AutoMigrate as
|
|
// SQL NULL, not ”, and the original guard `token_plain = ”` never
|
|
// matches NULL — the backfill silently never fired on any legacy node.
|
|
// Seed the row exactly as production has it (raw SQL, explicit NULL)
|
|
// instead of through GORM, which would write ”.
|
|
func TestAgentTokenBackfillFromNullLegacyRow(t *testing.T) {
|
|
engine, node := setupBackfillTest(t)
|
|
const token = "legacy-token-backfill-00000000000000000004"
|
|
db := database.GetDB()
|
|
if err := db.Exec(
|
|
"INSERT INTO agent_nodes (token_hash, token_prefix, token_plain, name, node_id, last_seen, created_at, updated_at) VALUES (?, ?, NULL, ?, ?, 0, 0, 0)",
|
|
hashAgentToken(token), token[:12], "backfill-node", node.Id,
|
|
).Error; err != nil {
|
|
t.Fatalf("seed NULL legacy row: %v", err)
|
|
}
|
|
var seeded model.AgentNode
|
|
if err := db.Where("token_hash = ?", hashAgentToken(token)).First(&seeded).Error; err != nil {
|
|
t.Fatalf("load seeded row: %v", err)
|
|
}
|
|
var nullRows int
|
|
if err := db.Raw("SELECT COUNT(*) FROM agent_nodes WHERE id = ? AND token_plain IS NULL", seeded.Id).Scan(&nullRows).Error; err != nil {
|
|
t.Fatalf("NULL check: %v", err)
|
|
}
|
|
if nullRows != 1 {
|
|
t.Fatalf("seeded row is not SQL NULL as intended (null rows = %d)", nullRows)
|
|
}
|
|
|
|
if code := heartbeatWith(engine, token); code != http.StatusOK {
|
|
t.Fatalf("heartbeat with valid token: status %d", code)
|
|
}
|
|
if row := reloadAgent(t, seeded.Id); row.TokenPlain != token {
|
|
t.Fatalf("token not backfilled from NULL legacy row: %q", row.TokenPlain)
|
|
}
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/panel/api/agent/install/"+strconv.Itoa(node.Id), nil)
|
|
rec := httptest.NewRecorder()
|
|
engine.ServeHTTP(rec, req)
|
|
if rec.Code != http.StatusOK {
|
|
t.Fatalf("install endpoint: status %d", rec.Code)
|
|
}
|
|
var resp struct {
|
|
Success bool `json:"success"`
|
|
Obj struct {
|
|
Token string `json:"token"`
|
|
HasToken bool `json:"hasToken"`
|
|
} `json:"obj"`
|
|
}
|
|
if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
|
|
t.Fatalf("install response not JSON: %v", err)
|
|
}
|
|
if !resp.Success || !resp.Obj.HasToken || resp.Obj.Token != token {
|
|
t.Fatalf("install endpoint did not return the backfilled token: %+v", resp.Obj)
|
|
}
|
|
}
|
|
|
|
func TestAgentTokenNoBackfillOnFailedAuth(t *testing.T) {
|
|
engine, node := setupBackfillTest(t)
|
|
const token = "legacy-token-backfill-00000000000000000002"
|
|
agent := seedLegacyAgent(t, node.Id, token, "")
|
|
|
|
if code := heartbeatWith(engine, "definitely-the-wrong-token"); code != http.StatusUnauthorized {
|
|
t.Fatalf("heartbeat with wrong token: status %d, want 401", code)
|
|
}
|
|
if row := reloadAgent(t, agent.Id); row.TokenPlain != "" {
|
|
t.Fatalf("failed authentication must not backfill, got %q", row.TokenPlain)
|
|
}
|
|
}
|
|
|
|
func TestAgentTokenBackfillNeverOverwrites(t *testing.T) {
|
|
engine, node := setupBackfillTest(t)
|
|
const token = "legacy-token-backfill-00000000000000000003"
|
|
agent := seedLegacyAgent(t, node.Id, token, "previously-stored-value")
|
|
|
|
if code := heartbeatWith(engine, token); code != http.StatusOK {
|
|
t.Fatalf("heartbeat with valid token: status %d", code)
|
|
}
|
|
if row := reloadAgent(t, agent.Id); row.TokenPlain != "previously-stored-value" {
|
|
t.Fatalf("existing plaintext was overwritten by a request: %q", row.TokenPlain)
|
|
}
|
|
}
|