hfb_sys/backend/internal/middleware/site_domain_fallback_test.go

package middleware

import (
	"context"
	"errors"
	"net/http"
	"net/http/httptest"
	"testing"

	"hfb_sys/backend/internal/sitecode"

	"github.com/gin-gonic/gin"
)

type fakeSiteDomainResolver struct {
	byDomain map[string]string
	err      error
	calls    int
	hosts    []string
}

func (f *fakeSiteDomainResolver) ResolveSiteCodeByDomain(_ context.Context, domain string) (string, error) {
	f.calls++
	f.hosts = append(f.hosts, domain)
	if f.err != nil {
		return "", f.err
	}
	return f.byDomain[domain], nil
}

// newSiteFallbackRouter 构建一个把 sitecode.FromRequest 结果回显的路由，便于断言兜底结果。
func newSiteFallbackRouter(resolver SiteDomainResolver) *gin.Engine {
	gin.SetMode(gin.TestMode)
	engine := gin.New()
	engine.Use(SiteDomainFallback(resolver, nil))
	engine.GET("/probe", func(c *gin.Context) {
		c.String(http.StatusOK, sitecode.FromRequest(c))
	})
	return engine
}

func doSiteFallbackRequest(t *testing.T, engine *gin.Engine, host, target string, headers map[string]string) string {
	t.Helper()
	req := httptest.NewRequest(http.MethodGet, target, nil)
	req.Host = host
	for key, value := range headers {
		req.Header.Set(key, value)
	}
	recorder := httptest.NewRecorder()
	engine.ServeHTTP(recorder, req)
	if recorder.Code != http.StatusOK {
		t.Fatalf("status = %d, want 200", recorder.Code)
	}
	return recorder.Body.String()
}

func TestSiteDomainFallbackResolvesByHost(t *testing.T) {
	resolver := &fakeSiteDomainResolver{byDomain: map[string]string{"tysh.hao858.com": "gz_dist"}}
	engine := newSiteFallbackRouter(resolver)

	if got := doSiteFallbackRequest(t, engine, "tysh.hao858.com", "/probe", nil); got != "gz_dist" {
		t.Fatalf("站点域名应解析为 gz_dist, got %q", got)
	}
	// 带端口的 Host 也要能解析
	if got := doSiteFallbackRequest(t, engine, "tysh.hao858.com:443", "/probe", nil); got != "gz_dist" {
		t.Fatalf("带端口站点域名应解析为 gz_dist, got %q", got)
	}
	// 主站域名无匹配：保持为空
	if got := doSiteFallbackRequest(t, engine, "hao858.com", "/probe", nil); got != "" {
		t.Fatalf("主站域名不应被识别为分站, got %q", got)
	}
}

func TestSiteDomainFallbackPrefersExplicitSiteCode(t *testing.T) {
	resolver := &fakeSiteDomainResolver{byDomain: map[string]string{"tysh.hao858.com": "gz_dist"}}
	engine := newSiteFallbackRouter(resolver)

	// 显式头优先：即使域名能解析出 gz_dist，也必须尊重前端给的站点
	cases := []struct {
		name    string
		target  string
		headers map[string]string
	}{
		{name: "权威头", target: "/probe", headers: map[string]string{sitecode.HeaderName: "other_site"}},
		{name: "旧头", target: "/probe", headers: map[string]string{sitecode.LegacyHeaderName: "legacy_site"}},
		{name: "query", target: "/probe?site_code=query_site", headers: nil},
		{name: "旧query", target: "/probe?distribution_site_code=legacy_query", headers: nil},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			got := doSiteFallbackRequest(t, engine, "tysh.hao858.com", tc.target, tc.headers)
			if got == "gz_dist" || got == "" {
				t.Fatalf("显式站点应优先于域名兜底, got %q", got)
			}
		})
	}
	if resolver.calls != 0 {
		t.Fatalf("显式站点不应触发域名解析, calls = %d", resolver.calls)
	}
}

func TestSiteDomainFallbackSkipsLocalAndIPHosts(t *testing.T) {
	resolver := &fakeSiteDomainResolver{byDomain: map[string]string{}}
	engine := newSiteFallbackRouter(resolver)

	for _, host := range []string{"localhost:8080", "127.0.0.1:18088", "10.0.0.5:8080"} {
		if got := doSiteFallbackRequest(t, engine, host, "/probe", nil); got != "" {
			t.Fatalf("本地/内网地址 %s 不应解析站点, got %q", host, got)
		}
	}
	// 本地地址直接跳过；10.0.0.5 不是 IsLocalHost 白名单，会查一次但无匹配
	if resolver.calls > 1 {
		t.Fatalf("本地地址不应反复查询, calls = %d", resolver.calls)
	}
}

func TestSiteDomainFallbackFallsBackToForwardedHost(t *testing.T) {
	resolver := &fakeSiteDomainResolver{byDomain: map[string]string{"tysh.hao858.com": "gz_dist"}}
	engine := newSiteFallbackRouter(resolver)

	got := doSiteFallbackRequest(t, engine, "hao858.com", "/probe", map[string]string{
		"X-Forwarded-Host": "tysh.hao858.com",
	})
	if got != "gz_dist" {
		t.Fatalf("应回退到 X-Forwarded-Host 解析站点, got %q", got)
	}
}

func TestSiteDomainFallbackCachesAndFailsOpen(t *testing.T) {
	resolver := &fakeSiteDomainResolver{byDomain: map[string]string{"tysh.hao858.com": "gz_dist"}}
	engine := newSiteFallbackRouter(resolver)

	for i := 0; i < 3; i++ {
		if got := doSiteFallbackRequest(t, engine, "tysh.hao858.com", "/probe", nil); got != "gz_dist" {
			t.Fatalf("第 %d 次解析 = %q, want gz_dist", i+1, got)
		}
	}
	if resolver.calls != 1 {
		t.Fatalf("同域名应命中缓存只查一次, calls = %d", resolver.calls)
	}

	// 解析失败时不阻断请求，退化为未指定站点
	failing := &fakeSiteDomainResolver{err: errors.New("db down")}
	failEngine := newSiteFallbackRouter(failing)
	for i := 0; i < 3; i++ {
		if got := doSiteFallbackRequest(t, failEngine, "tysh.hao858.com", "/probe", nil); got != "" {
			t.Fatalf("解析失败应退化为空站点, got %q", got)
		}
	}
	if failing.calls != 1 {
		t.Fatalf("解析失败也应缓存避免刷日志, calls = %d", failing.calls)
	}
}