3x-ui-agent/internal/web/controller/agent_seen_ip_test.go

package controller

import (
	"bytes"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"testing"

	"github.com/gin-gonic/gin"

	"github.com/mhsanaei/3x-ui/v3/internal/database"
	"github.com/mhsanaei/3x-ui/v3/internal/database/model"
)

// Public-address capture: every authenticated agent request records its
// source address on the agent row — CF-Connecting-IP first (agents come
// through Cloudflare), then the first X-Forwarded-For hop, then the
// transport RemoteAddr. Writes are change-only so the ~10s traffic
// cadence does not become a DB write per request.
func TestAgentSeenIPCapture(t *testing.T) {
	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() })
	db := database.GetDB()

	node := &model.Node{Name: "seen-ip-node", Address: ""}
	if err := db.Create(node).Error; err != nil {
		t.Fatalf("seed node: %v", err)
	}
	const token = "seen-ip-test-token"
	agent := &model.AgentNode{
		TokenHash:   hashAgentToken(token),
		TokenPrefix: "seen-ip-",
		Name:        "seen-ip-node",
		NodeID:      &node.Id,
	}
	if err := db.Create(agent).Error; err != nil {
		t.Fatalf("seed agent: %v", err)
	}

	engine := gin.New()
	NewAgentController(engine.Group("/panel/api/agent"))

	heartbeat := func(headers map[string]string, remoteAddr 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)
		for k, v := range headers {
			req.Header.Set(k, v)
		}
		if remoteAddr != "" {
			req.RemoteAddr = remoteAddr
		}
		rec := httptest.NewRecorder()
		engine.ServeHTTP(rec, req)
		return rec.Code
	}
	load := func() model.AgentNode {
		var row model.AgentNode
		if err := db.First(&row, agent.Id).Error; err != nil {
			t.Fatalf("reload agent: %v", err)
		}
		return row
	}

	// CF-Connecting-IP wins and is stored.
	if code := heartbeat(map[string]string{"CF-Connecting-IP": "111.91.29.150"}, ""); code != http.StatusOK {
		t.Fatalf("heartbeat with CF header: status %d", code)
	}
	row := load()
	if row.SeenIPv4 != "111.91.29.150" || row.SeenIPAt == 0 {
		t.Fatalf("CF source not stored: %+v", row)
	}

	// Change-only: same address again must not rewrite (SeenIPAt frozen).
	frozen := row.SeenIPAt
	if err := db.Model(&model.AgentNode{}).Where("id = ?", agent.Id).
		Update("seen_ip_at", frozen).Error; err != nil {
		t.Fatalf("pin seen_ip_at: %v", err)
	}
	if code := heartbeat(map[string]string{"CF-Connecting-IP": "111.91.29.150"}, ""); code != http.StatusOK {
		t.Fatalf("repeat heartbeat: status %d", code)
	}
	if row = load(); row.SeenIPAt != frozen {
		t.Fatalf("unchanged address rewrote the row: seen_ip_at %d -> %d", frozen, row.SeenIPAt)
	}

	// X-Forwarded-For first hop is next in precedence.
	if code := heartbeat(map[string]string{"X-Forwarded-For": "203.0.113.9, 70.41.3.18"}, ""); code != http.StatusOK {
		t.Fatalf("heartbeat with XFF: status %d", code)
	}
	if row = load(); row.SeenIPv4 != "203.0.113.9" {
		t.Fatalf("XFF first hop not stored: %+v", row)
	}

	// CF beats XFF when both are present.
	if code := heartbeat(map[string]string{
		"CF-Connecting-IP": "111.91.29.150",
		"X-Forwarded-For":  "203.0.113.9",
	}, ""); code != http.StatusOK {
		t.Fatalf("heartbeat with both headers: status %d", code)
	}
	if row = load(); row.SeenIPv4 != "111.91.29.150" {
		t.Fatalf("CF header did not win over XFF: %+v", row)
	}

	// A garbage CF value falls through to XFF.
	if code := heartbeat(map[string]string{
		"CF-Connecting-IP": "not-an-ip",
		"X-Forwarded-For":  "198.51.100.7",
	}, ""); code != http.StatusOK {
		t.Fatalf("heartbeat with garbage CF header: status %d", code)
	}
	if row = load(); row.SeenIPv4 != "198.51.100.7" {
		t.Fatalf("garbage CF value did not fall through to XFF: %+v", row)
	}

	// No headers: the transport RemoteAddr is the fallback.
	if code := heartbeat(nil, "198.51.100.23:5555"); code != http.StatusOK {
		t.Fatalf("heartbeat without headers: status %d", code)
	}
	if row = load(); row.SeenIPv4 != "198.51.100.23" {
		t.Fatalf("RemoteAddr fallback not stored: %+v", row)
	}

	// IPv6 sources land in the v6 column.
	if code := heartbeat(map[string]string{"CF-Connecting-IP": "2602:f32b:1:310f::1"}, ""); code != http.StatusOK {
		t.Fatalf("heartbeat with v6 CF header: status %d", code)
	}
	if row = load(); row.SeenIPv6 != "2602:f32b:1:310f::1" {
		t.Fatalf("v6 source not stored: %+v", row)
	}

	// Unauthenticated requests capture nothing.
	if err := db.Model(&model.AgentNode{}).Where("id = ?", agent.Id).
		Updates(map[string]any{"seen_ipv4": "", "seen_ipv6": ""}).Error; err != nil {
		t.Fatalf("reset seen: %v", err)
	}
	req := httptest.NewRequest(http.MethodPost, "/panel/api/agent/heartbeat", bytes.NewBufferString(`{}`))
	req.Header.Set("Authorization", "Bearer wrong-token")
	req.Header.Set("CF-Connecting-IP", "111.91.29.150")
	rec := httptest.NewRecorder()
	engine.ServeHTTP(rec, req)
	if rec.Code != http.StatusUnauthorized {
		t.Fatalf("bad token heartbeat: status %d, want 401", rec.Code)
	}
	if row = load(); row.SeenIPv4 != "" || row.SeenIPv6 != "" {
		t.Fatalf("unauthenticated request captured a source: %+v", row)
	}

	// Register is an authenticated agent request too: it captures.
	regReq := httptest.NewRequest(http.MethodPost, "/panel/api/agent/register",
		bytes.NewBufferString(`{"token":"`+token+`","name":"seen-ip-node"}`))
	regReq.Header.Set("Content-Type", "application/json")
	regReq.Header.Set("CF-Connecting-IP", "111.91.29.150")
	regRec := httptest.NewRecorder()
	engine.ServeHTTP(regRec, regReq)
	if regRec.Code != http.StatusOK {
		t.Fatalf("register: status %d", regRec.Code)
	}
	if row = load(); row.SeenIPv4 != "111.91.29.150" {
		t.Fatalf("register did not capture source: %+v", row)
	}
}