summaryrefslogtreecommitdiff
path: root/internal/api/recommend_test.go
diff options
context:
space:
mode:
Diffstat (limited to 'internal/api/recommend_test.go')
-rw-r--r--internal/api/recommend_test.go549
1 files changed, 549 insertions, 0 deletions
diff --git a/internal/api/recommend_test.go b/internal/api/recommend_test.go
new file mode 100644
index 0000000..8449db0
--- /dev/null
+++ b/internal/api/recommend_test.go
@@ -0,0 +1,549 @@
+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)
+}
generated by cgit v1.2.3 (git 2.39.1) at 2026-05-31 22:13:34 +0000