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