3x-ui-agent/internal/web/controller/agent_token_backfill_test.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)
	}
}

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)
	}
}