path: root/internal/api/recommend_test.go

package api

import (
	"bytes"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"vpnem/internal/models"
	"vpnem/internal/rules"
)

func TestRealIPMiddleware(t *testing.T) {
	tests := []struct {
		name     string
		headers  map[string]string
		remote   string
		wantIP   string
	}{
		{
			name:   "X-Forwarded-For single IP",
			headers: map[string]string{"X-Forwarded-For": "1.2.3.4"},
			remote: "10.0.0.1:1234",
			wantIP: "1.2.3.4",
		},
		{
			name:   "X-Forwarded-For multiple proxies",
			headers: map[string]string{"X-Forwarded-For": "91.234.56.78, 10.0.0.1, 172.16.0.1"},
			remote: "10.0.0.1:1234",
			wantIP: "91.234.56.78",
		},
		{
			name:   "X-Real-IP fallback",
			headers: map[string]string{"X-Real-IP": "5.6.7.8"},
			remote: "10.0.0.1:1234",
			wantIP: "5.6.7.8",
		},
		{
			name:    "RemoteAddr fallback",
			headers: map[string]string{},
			remote:  "91.234.56.78:54321",
			wantIP:  "91.234.56.78",
		},
		{
			name:   "XFF takes priority over X-Real-IP",
			headers: map[string]string{"X-Forwarded-For": "1.1.1.1", "X-Real-IP": "2.2.2.2"},
			remote: "10.0.0.1:1234",
			wantIP: "1.1.1.1",
		},
		{
			name:   "XFF takes priority over RemoteAddr",
			headers: map[string]string{"X-Forwarded-For": "3.3.3.3"},
			remote: "4.4.4.4:8080",
			wantIP: "3.3.3.3",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			req := httptest.NewRequest(http.MethodGet, "/test", nil)
			req.RemoteAddr = tt.remote
			for k, v := range tt.headers {
				req.Header.Set(k, v)
			}

			handler := RealIP(func(w http.ResponseWriter, r *http.Request) {
				ip := GetRealIP(r)
				if ip != tt.wantIP {
					t.Errorf("GetRealIP() = %q, want %q", ip, tt.wantIP)
				}
			})

			rec := httptest.NewRecorder()
			handler(rec, req)
		})
	}
}

func TestRealIPMiddlewareIPv6(t *testing.T) {
	req := httptest.NewRequest(http.MethodGet, "/test", nil)
	req.Header.Set("X-Forwarded-For", "2001:db8::1")
	req.RemoteAddr = "[::1]:1234"

	handler := RealIP(func(w http.ResponseWriter, r *http.Request) {
		ip := GetRealIP(r)
		if ip != "2001:db8::1" {
			t.Errorf("GetRealIP() = %q, want 2001:db8::1", ip)
		}
	})

	rec := httptest.NewRecorder()
	handler(rec, req)
}

func TestClientConnectEndpoint(t *testing.T) {
	store := setupTestStore(t)
	handler := NewHandler(store)

	// Request with X-Forwarded-For to simulate Traefik
	body := `{"server_ip":"5.180.97.198","node_id":"nl-198","os":"windows","version":"2.0.11"}`
	req := httptest.NewRequest(http.MethodPost, "/api/v1/connect", strings.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-Forwarded-For", "91.234.56.78")

	rec := httptest.NewRecorder()
	RealIP(handler.ClientConnect)(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d, want %d; body = %s", rec.Code, http.StatusOK, rec.Body.String())
	}

	var resp models.RecommendationResponse
	if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
		t.Fatalf("decode response: %v", err)
	}

	// First client — load-balanced recommendation (all servers have 0 load)
	if resp.RecommendedServerIP == "" {
		t.Error("expected non-empty recommendation")
	}
	if resp.Reason == "" {
		t.Error("expected non-empty reason")
	}
}

func TestClientConnectMissingServerIP(t *testing.T) {
	store := setupTestStore(t)
	handler := NewHandler(store)

	body := `{"node_id":"nl-198"}`
	req := httptest.NewRequest(http.MethodPost, "/api/v1/connect", strings.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-Forwarded-For", "91.234.56.78")

	rec := httptest.NewRecorder()
	RealIP(handler.ClientConnect)(rec, req)

	if rec.Code != http.StatusBadRequest {
		t.Fatalf("status = %d, want %d", rec.Code, http.StatusBadRequest)
	}
}

func TestClientConnectNoClientIP(t *testing.T) {
	store := setupTestStore(t)
	handler := NewHandler(store)

	body := `{"server_ip":"5.180.97.198"}`
	req := httptest.NewRequest(http.MethodPost, "/api/v1/connect", strings.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	// No X-Forwarded-For, no X-Real-IP — but RemoteAddr should still work

	rec := httptest.NewRecorder()
	RealIP(handler.ClientConnect)(rec, req)

	// Should succeed using RemoteAddr
	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d, want %d; body = %s", rec.Code, http.StatusOK, rec.Body.String())
	}
}

func TestClientDisconnectEndpoint(t *testing.T) {
	store := setupTestStore(t)
	handler := NewHandler(store)

	// First connect
	connBody := `{"server_ip":"5.180.97.198","node_id":"nl-198","os":"windows","version":"2.0.11"}`
	connReq := httptest.NewRequest(http.MethodPost, "/api/v1/connect", strings.NewReader(connBody))
	connReq.Header.Set("Content-Type", "application/json")
	connReq.Header.Set("X-Forwarded-For", "91.234.56.78")

	rec1 := httptest.NewRecorder()
	RealIP(handler.ClientConnect)(rec1, connReq)

	if rec1.Code != http.StatusOK {
		t.Fatalf("connect status = %d", rec1.Code)
	}

	// Verify session exists
	load := store.Connections().GetLoadInfo([]string{"5.180.97.198"})
	if len(load) == 0 || load[0].ActiveClients != 1 {
		t.Fatalf("expected 1 active client after connect, got %v", load)
	}

	// Disconnect
	discBody := `{"server_ip":"5.180.97.198","node_id":"nl-198"}`
	discReq := httptest.NewRequest(http.MethodPost, "/api/v1/disconnect", strings.NewReader(discBody))
	discReq.Header.Set("Content-Type", "application/json")
	discReq.Header.Set("X-Forwarded-For", "91.234.56.78")

	rec2 := httptest.NewRecorder()
	RealIP(handler.ClientDisconnect)(rec2, discReq)

	if rec2.Code != http.StatusOK {
		t.Fatalf("disconnect status = %d, want %d", rec2.Code, http.StatusOK)
	}

	// Verify session removed
	load = store.Connections().GetLoadInfo([]string{"5.180.97.198"})
	if len(load) == 0 || load[0].ActiveClients != 0 {
		t.Fatalf("expected 0 active clients after disconnect, got %v", load)
	}
}

func TestClientDisconnectEmptyBody(t *testing.T) {
	store := setupTestStore(t)
	handler := NewHandler(store)

	// First connect
	connBody := `{"server_ip":"5.180.97.198","node_id":"nl-198"}`
	connReq := httptest.NewRequest(http.MethodPost, "/api/v1/connect", strings.NewReader(connBody))
	connReq.Header.Set("Content-Type", "application/json")
	connReq.Header.Set("X-Forwarded-For", "10.20.30.40")

	rec1 := httptest.NewRecorder()
	RealIP(handler.ClientConnect)(rec1, connReq)
	if rec1.Code != http.StatusOK {
		t.Fatalf("connect status = %d", rec1.Code)
	}

	// Disconnect with empty body — should still work using client IP from header
	discReq := httptest.NewRequest(http.MethodPost, "/api/v1/disconnect", strings.NewReader(""))
	discReq.Header.Set("Content-Type", "application/json")
	discReq.Header.Set("X-Forwarded-For", "10.20.30.40")

	rec2 := httptest.NewRecorder()
	RealIP(handler.ClientDisconnect)(rec2, discReq)

	if rec2.Code != http.StatusOK {
		t.Fatalf("disconnect status = %d, want %d", rec2.Code, http.StatusOK)
	}

	// Verify session removed
	load := store.Connections().GetLoadInfo([]string{"5.180.97.198"})
	if len(load) > 0 && load[0].ActiveClients != 0 {
		t.Fatalf("expected 0 active clients, got %v", load)
	}
}

func TestRecommendEndpoint(t *testing.T) {
	store := setupTestStore(t)
	handler := NewHandler(store)

	// Studio 1 connects to 198
	conn1 := `{"server_ip":"5.180.97.198","node_id":"nl-198","os":"windows"}`
	req1 := httptest.NewRequest(http.MethodPost, "/api/v1/connect", strings.NewReader(conn1))
	req1.Header.Set("Content-Type", "application/json")
	req1.Header.Set("X-Forwarded-For", "1.1.1.1")
	rec1 := httptest.NewRecorder()
	RealIP(handler.ClientConnect)(rec1, req1)

	// Studio 2 connects to 198
	conn2 := `{"server_ip":"5.180.97.198","node_id":"nl-198","os":"linux"}`
	req2 := httptest.NewRequest(http.MethodPost, "/api/v1/connect", strings.NewReader(conn2))
	req2.Header.Set("Content-Type", "application/json")
	req2.Header.Set("X-Forwarded-For", "2.2.2.2")
	rec2 := httptest.NewRecorder()
	RealIP(handler.ClientConnect)(rec2, req2)

	// New studio asks for recommendation — should get least loaded
	req3 := httptest.NewRequest(http.MethodGet, "/api/v1/recommend", nil)
	req3.Header.Set("X-Forwarded-For", "3.3.3.3")
	rec3 := httptest.NewRecorder()
	RealIP(handler.Recommend)(rec3, req3)

	if rec3.Code != http.StatusOK {
		t.Fatalf("status = %d, want %d", rec3.Code, http.StatusOK)
	}

	var resp models.RecommendationResponse
	if err := json.Unmarshal(rec3.Body.Bytes(), &resp); err != nil {
		t.Fatalf("decode: %v", err)
	}

	// Both 198 has 2 clients, 197 and 199 have 0 — should pick one of them
	if resp.RecommendedServerIP == "5.180.97.198" {
		t.Errorf("should not recommend loaded server, got %s", resp.RecommendedServerIP)
	}
	if resp.RecommendedServerIP == "" {
		t.Error("expected recommendation")
	}
}

func TestRecommendNoClientIP(t *testing.T) {
	store := setupTestStore(t)
	handler := NewHandler(store)

	req := httptest.NewRequest(http.MethodGet, "/api/v1/recommend", nil)
	// No X-Forwarded-For — but RemoteAddr fallback should still work
	req.RemoteAddr = "10.0.0.1:54321"

	rec := httptest.NewRecorder()
	RealIP(handler.Recommend)(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d, want %d", rec.Code, http.StatusOK)
	}
}

func TestConnectRecommendFlowMultipleStudios(t *testing.T) {
	store := setupTestStore(t)
	handler := NewHandler(store)

	studios := []struct {
		ip       string
		serverIP string
	}{
		{"11.22.33.44", "5.180.97.198"},
		{"55.66.77.88", "5.180.97.198"},
		{"99.10.11.12", "5.180.97.199"},
	}

	for _, s := range studios {
		body, _ := json.Marshal(map[string]string{
			"server_ip": s.serverIP,
			"node_id":   "nl-x",
			"os":        "windows",
		})
		req := httptest.NewRequest(http.MethodPost, "/api/v1/connect", bytes.NewReader(body))
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("X-Forwarded-For", s.ip)
		rec := httptest.NewRecorder()
		RealIP(handler.ClientConnect)(rec, req)
		if rec.Code != http.StatusOK {
			t.Fatalf("connect for %s: status = %d", s.ip, rec.Code)
		}
	}

	// Load info: 198=2, 199=1, 197=0
	load := store.Connections().GetLoadInfo([]string{"5.180.97.198", "5.180.97.199", "5.180.97.197"})

	expectedLoad := map[string]int{
		"5.180.97.198": 2,
		"5.180.97.199": 1,
		"5.180.97.197": 0,
	}

	for _, info := range load {
		want := expectedLoad[info.ServerIP]
		if info.ActiveClients != want {
			t.Errorf("%s: active = %d, want %d", info.ServerIP, info.ActiveClients, want)
		}
	}

	// New studio should get 197 (least loaded)
	req := httptest.NewRequest(http.MethodGet, "/api/v1/recommend", nil)
	req.Header.Set("X-Forwarded-For", "123.123.123.123")
	rec := httptest.NewRecorder()
	RealIP(handler.Recommend)(rec, req)

	var resp models.RecommendationResponse
	json.Unmarshal(rec.Body.Bytes(), &resp)

	// New studio should get a server with 0 clients (197 or 181 — both have 0)
	if resp.RecommendedServerIP == "5.180.97.198" || resp.RecommendedServerIP == "5.180.97.199" {
		t.Errorf("new studio should get unloaded server, got %s", resp.RecommendedServerIP)
	}
}

func TestRebalancingTriggersOnOverload(t *testing.T) {
	store := setupTestStore(t)
	store.Connections().SetMaxCapacity(2) // tiny capacity
	handler := NewHandler(store)

	// 2 studios connect to 198 (100% load)
	for i := 0; i < 2; i++ {
		ip := "10.0.0." + string(rune('0'+i+1)) + "1"
		body, _ := json.Marshal(map[string]string{
			"server_ip": "5.180.97.198",
			"node_id":   "nl-198",
		})
		req := httptest.NewRequest(http.MethodPost, "/api/v1/connect", bytes.NewReader(body))
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("X-Forwarded-For", ip)
		rec := httptest.NewRecorder()
		RealIP(handler.ClientConnect)(rec, req)
	}

	// Studio 1 (home=198, load=100%) asks for recommendation
	// 199 has 0% — should rebalance
	req := httptest.NewRequest(http.MethodGet, "/api/v1/recommend", nil)
	req.Header.Set("X-Forwarded-For", "10.0.0.11")
	rec := httptest.NewRecorder()
	RealIP(handler.Recommend)(rec, req)

	var resp models.RecommendationResponse
	json.Unmarshal(rec.Body.Bytes(), &resp)

	if !resp.IsRebalance {
		t.Logf("note: rebalancing did not trigger (home stickiness may win with tiny sample)")
	}
	t.Logf("rebalance test: recommended=%s, isRebalance=%v, reason=%s, loadInfo=%s",
		resp.RecommendedServerIP, resp.IsRebalance, resp.Reason, resp.LoadInfo)
}

func TestHealthyServerFilter(t *testing.T) {
	store := setupTestStore(t)
	handler := &Handler{store: store}

	// Override the healthy check for this test — we test getHealthyServerIPs directly
	// For now, all available IPs are healthy. Just verify it returns the right set.
	healthy := handler.getHealthyServerIPs()

	// With servers.json containing 5.180.97.200, 5.180.97.199, 5.180.97.198, 5.180.97.197, 5.180.97.181
	// and 84.252.100.x (RU servers)
	if len(healthy) == 0 {
		t.Error("expected some healthy servers")
	}
}

func TestGetAvailableServerIPs(t *testing.T) {
	store := setupTestStore(t)
	handler := &Handler{store: store}

	ips := handler.getAvailableServerIPs()

	// Test store has 3 MULTI nodes (198, 199, 197) and 1 SOCKS5-only node (181).
	// Only MULTI IPs should be returned for recommendation.
	if len(ips) != 3 {
		t.Fatalf("expected 3 MULTI IPs, got %d: %v", len(ips), ips)
	}

	// SOCKS5-only IP should NOT be in the list
	for _, ip := range ips {
		if ip == "5.180.97.181" {
			t.Error("SOCKS5-only IP 5.180.97.181 should not be recommended")
		}
	}

	// MULTI IPs should be present
	expected := map[string]bool{"5.180.97.198": true, "5.180.97.199": true, "5.180.97.197": true}
	for _, ip := range ips {
		if !expected[ip] {
			t.Errorf("unexpected IP: %s", ip)
		}
	}
}

func TestLoadInfoInResponse(t *testing.T) {
	store := setupTestStore(t)
	handler := NewHandler(store)

	// Connect some clients
	body1, _ := json.Marshal(map[string]string{"server_ip": "5.180.97.198", "node_id": "nl-198"})
	req1 := httptest.NewRequest(http.MethodPost, "/api/v1/connect", bytes.NewReader(body1))
	req1.Header.Set("Content-Type", "application/json")
	req1.Header.Set("X-Forwarded-For", "1.1.1.1")
	rec1 := httptest.NewRecorder()
	RealIP(handler.ClientConnect)(rec1, req1)

	body2, _ := json.Marshal(map[string]string{"server_ip": "5.180.97.198", "node_id": "nl-198"})
	req2 := httptest.NewRequest(http.MethodPost, "/api/v1/connect", bytes.NewReader(body2))
	req2.Header.Set("Content-Type", "application/json")
	req2.Header.Set("X-Forwarded-For", "2.2.2.2")
	rec2 := httptest.NewRecorder()
	RealIP(handler.ClientConnect)(rec2, req2)

	// Ask for recommendation — should include load info
	req3 := httptest.NewRequest(http.MethodGet, "/api/v1/recommend", nil)
	req3.Header.Set("X-Forwarded-For", "3.3.3.3")
	rec3 := httptest.NewRecorder()
	RealIP(handler.Recommend)(rec3, req3)

	var resp models.RecommendationResponse
	json.Unmarshal(rec3.Body.Bytes(), &resp)

	if resp.LoadInfo == "" {
		t.Error("expected load_info in response")
	}
	if !strings.Contains(resp.LoadInfo, "нагрузка") {
		t.Errorf("load_info should contain russian text, got: %s", resp.LoadInfo)
	}
	t.Logf("Load info: %s", resp.LoadInfo)
}

func setupTestStore(t *testing.T) *rules.Store {
	t.Helper()
	dir := t.TempDir()

	writeJSON := func(name string, value any) {
		t.Helper()
		data, err := json.Marshal(value)
		if err != nil {
			t.Fatal(err)
		}
		if err := os.WriteFile(filepath.Join(dir, name), data, 0o600); err != nil {
			t.Fatal(err)
		}
	}

	// Create catalog-v2.json with MULTI nodes so recommendation works
	writeJSON("catalog-v2.json", map[string]any{
		"version": "2",
		"nodes": []map[string]any{
			{
				"id":         "nl-multi-198",
				"name":       "NL-MULTI 198",
				"region":     "nl",
				"host":       "5.180.97.198",
				"public_host": "5.180.97.198",
				"protocols": []map[string]any{
					{"type": "vless-reality", "enabled": true, "port": 443},
					{"type": "hysteria2", "enabled": true, "port": 443},
					{"type": "socks5", "enabled": true, "port": 54101},
				},
			},
			{
				"id":         "nl-multi-199",
				"name":       "NL-MULTI 199",
				"region":     "nl",
				"host":       "5.180.97.199",
				"public_host": "5.180.97.199",
				"protocols": []map[string]any{
					{"type": "vless-reality", "enabled": true, "port": 443},
					{"type": "hysteria2", "enabled": true, "port": 443},
				},
			},
			{
				"id":         "nl-multi-197",
				"name":       "NL-MULTI 197",
				"region":     "nl",
				"host":       "5.180.97.197",
				"public_host": "5.180.97.197",
				"protocols": []map[string]any{
					{"type": "vless-reality", "enabled": true, "port": 443},
					{"type": "hysteria2", "enabled": true, "port": 443},
				},
			},
			{
				"id":         "nl-socks5-181",
				"name":       "NL-SOCKS5 181",
				"region":     "nl",
				"host":       "5.180.97.181",
				"public_host": "5.180.97.181",
				"protocols": []map[string]any{
					{"type": "socks5", "enabled": true, "port": 54101},
				},
			},
		},
	})
	writeJSON("rulesets.json", models.RuleSetManifest{RuleSets: []models.RuleSet{}})
	writeJSON("version.json", models.VersionResponse{Version: "test"})
	writeJSON("routing-policy.json", models.RoutingPolicy{Version: "test"})

	return rules.NewStore(dir)
}