Volcano-Engine-TTS-UI/middleware/metricsip_test.go

package middleware

import (
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

// allowListFrom 把逗号分隔的 CIDR 串解析成白名单,供测试直接注入。
func allowListFrom(t *testing.T, spec string) []*net.IPNet {
	t.Helper()
	var out []*net.IPNet
	for _, part := range strings.Split(spec, ",") {
		entry := strings.TrimSpace(part)
		if entry == "" {
			continue
		}
		_, n, err := net.ParseCIDR(entry)
		if err != nil {
			t.Fatalf("测试用例里的 CIDR 非法 %q: %v", entry, err)
		}
		out = append(out, n)
	}
	return out
}

// reqFrom 构造带指定来源地址的请求。
func reqFrom(remoteAddr string) *http.Request {
	r := httptest.NewRequest("GET", "/metrics", nil)
	r.RemoteAddr = remoteAddr
	return r
}

// TestIPAllowed 验证 CIDR 白名单判定:
// 命中放行、未命中拒绝,解析失败的 IP 一律拒绝(宁可拒绝也不误放行)。
func TestIPAllowed(t *testing.T) {
	// 直接构造白名单,不依赖环境变量
	metricsAllowList = allowListFrom(t, "127.0.0.1/32,10.0.0.0/8,172.16.0.0/12,::1/128")
	metricsAllowConfigured = true
	t.Cleanup(func() {
		metricsAllowList = nil
		metricsAllowConfigured = false
	})

	cases := []struct {
		name string
		ip   string
		want bool
	}{
		{"本机 IPv4 命中 /32", "127.0.0.1", true},
		{"10.x 命中 /8", "10.1.2.3", true},
		{"172.16.x 命中 /12", "172.16.5.9", true},
		{"172.31.x 仍在 /12 内", "172.31.255.254", true},
		{"IPv6 回环命中 /128", "::1", true},
		{"公网 IP 不在白名单", "8.8.8.8", false},
		{"172.32.x 超出 /12 范围", "172.32.0.1", false},
		{"192.168.x 未配置", "192.168.1.1", false},
		{"空 IP 拒绝", "", false},
		{"非法 IP 拒绝", "not-an-ip", false},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			if got := ipAllowed(c.ip); got != c.want {
				t.Errorf("ipAllowed(%q) = %v, want %v", c.ip, got, c.want)
			}
		})
	}
}

// TestIPAllowed_EmptyListRejectsAll 白名单为空时全部拒绝(未配置 = 不开放)。
func TestIPAllowed_EmptyListRejectsAll(t *testing.T) {
	metricsAllowList = nil
	metricsAllowConfigured = false
	t.Cleanup(func() {
		metricsAllowList = nil
		metricsAllowConfigured = false
	})

	if ipAllowed("127.0.0.1") {
		t.Error("白名单为空时应拒绝所有 IP")
	}
}

// TestMetricsIPAllowList_Scope 白名单内放行、白名单外 404。
// 用 404 而不是 403,是为了不向扫描者确认"该端点存在,只是你没权限"。
func TestMetricsIPAllowList_Scope(t *testing.T) {
	metricsAllowList = allowListFrom(t, "10.0.0.0/8")
	t.Cleanup(func() { metricsAllowList = nil })

	h := MetricsIPAllowList(okHandler())

	// 命中白名单 → 放行到 next
	w := httptest.NewRecorder()
	h.ServeHTTP(w, reqFrom("10.1.1.1:1234"))
	if w.Code != http.StatusOK {
		t.Errorf("白名单内来源: code=%d, want 200", w.Code)
	}

	// 白名单外(RFC 5737 文档地址)→ 404,且不应触达 next
	w2 := httptest.NewRecorder()
	h.ServeHTTP(w2, reqFrom("192.0.2.1:1234"))
	if w2.Code != http.StatusNotFound {
		t.Errorf("白名单外来源: code=%d, want 404", w2.Code)
	}
	if w2.Body.String() == "ok" {
		t.Error("白名单外来源不应触达被保护的 handler")
	}
}

// TestMetricsIPAllowList_XForwardedForSpoof 伪造 X-Forwarded-For 不能绕过白名单。
//
// GetClientIP 的启发式模式(trustedProxyHops==0,默认)从 XFF 链**尾部**取第一个
// **公网** IP,刻意跳过私网跳 —— 这样攻击者无法用 "X-Forwarded-For: <内网IP>"
// 把自己伪装成白名单来源。本测试锁死这个安全属性。
func TestMetricsIPAllowList_XForwardedForSpoof(t *testing.T) {
	metricsAllowList = allowListFrom(t, "10.0.0.0/8")
	t.Cleanup(func() { metricsAllowList = nil })

	h := MetricsIPAllowList(okHandler())

	// 直连是私网(像反代),XFF 里塞一个内网 IP 想混进白名单
	r := reqFrom("127.0.0.1:1234")
	r.Header.Set("X-Forwarded-For", "10.9.9.9")
	w := httptest.NewRecorder()
	h.ServeHTTP(w, r)

	// 启发式模式会跳过私网跳、落到直连地址 127.0.0.1(不在 10.0.0.0/8)→ 404。
	// 关键断言:伪造的内网 XFF **没有**让它通过。
	if w.Code == http.StatusOK {
		t.Errorf("伪造私网 XFF 不应绕过白名单: code=%d", w.Code)
	}

	// 反向对照:XFF 填公网 IP 时,GetClientIP 会采用它,同样不在白名单 → 404
	r2 := reqFrom("127.0.0.1:1234")
	r2.Header.Set("X-Forwarded-For", "8.8.8.8")
	w2 := httptest.NewRecorder()
	h.ServeHTTP(w2, r2)
	if w2.Code != http.StatusNotFound {
		t.Errorf("公网来源不在白名单: code=%d, want 404", w2.Code)
	}
}