tamabox/internal/security/geoip/geoip_test.go

// Copyright 2026 tamakyi. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package geoip

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"net/url"
	"sync"
	"sync/atomic"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// rewriteTransport 把真实域名的请求改投到本地测试服务器。
//
// 用改投而不是给生产代码加一个 baseURL 字段：生产代码里不该为了测试多出
// 「可以被指向别处」的开关，那类开关一旦被配置错就是把攻击者 IP 发到未知主机。
type rewriteTransport struct {
	target *url.URL
}

func (t *rewriteTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	next := req.Clone(req.Context())
	next.URL.Scheme = t.target.Scheme
	next.URL.Host = t.target.Host
	return http.DefaultTransport.RoundTrip(next)
}

// newTestClient 返回一个把所有出站请求改投到 srv 的 client。
func newTestClient(t *testing.T, srv *httptest.Server) *http.Client {
	t.Helper()
	u, err := url.Parse(srv.URL)
	require.NoError(t, err)
	return &http.Client{Transport: &rewriteTransport{target: u}}
}

// ipAPIServer 起一个假的 ip-api 批量接口。
//
// respond 收到本批请求的地址列表，返回对应的 ipAPIResult 切片；
// status 非 0 时直接返回该 HTTP 状态码（用于模拟网关故障）。
func ipAPIServer(t *testing.T, respond func(addrs []string) ([]ipAPIResult, int)) (*httptest.Server, *int32) {
	t.Helper()
	var calls int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		if req.URL.Path != "/batch" {
			http.NotFound(w, req)
			return
		}
		if req.Method != http.MethodPost {
			http.Error(w, "method", http.StatusMethodNotAllowed)
			return
		}
		atomic.AddInt32(&calls, 1)
		var addrs []string
		if err := json.NewDecoder(req.Body).Decode(&addrs); err != nil {
			http.Error(w, "bad body", http.StatusBadRequest)
			return
		}
		results, code := respond(addrs)
		if code != 0 {
			http.Error(w, "upstream failure", code)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(results)
	}))
	t.Cleanup(srv.Close)
	return srv, &calls
}

// okIP 构造一条成功的 ip-api 结果。
func okIP(addr, country, code, city string, lat, lon float64, isp, as string) ipAPIResult {
	return ipAPIResult{
		Query: addr, Status: "success", Country: country, CountryCode: code,
		City: city, Lat: lat, Lon: lon, ISP: isp, AS: as,
	}
}

// TestLookupSuccess 三态里的「查到了」与「服务明确说没有」。
func TestLookupSuccess(t *testing.T) {
	srv, calls := ipAPIServer(t, func(addrs []string) ([]ipAPIResult, int) {
		return []ipAPIResult{
			okIP(addrs[0], "德国", "DE", "Berlin", 52.52, 13.4, "Hetzner", "AS24940 Hetzner Online GmbH"),
			{Query: addrs[1], Status: "fail", Message: "reserved range"},
		}, 0
	})

	r := New(Options{Provider: ProviderIPAPI, HTTPClient: newTestClient(t, srv)})
	res := r.Lookup(context.Background(), []string{"5.9.1.1", "203.0.113.7"})

	require.Equal(t, int32(1), atomic.LoadInt32(calls))
	require.Equal(t, 1, res.Requests)
	require.Equal(t, 0, res.Failures)

	g, ok := res.Geo["5.9.1.1"]
	require.True(t, ok)
	assert.Equal(t, "德国", g.Country)
	assert.Equal(t, "DE", g.CountryCode)
	assert.Equal(t, "Berlin", g.City)
	assert.Equal(t, "AS24940", g.ASN, "AS 字段必须截成纯 ASN，否则国家统计里全是长串")
	assert.True(t, g.HasCoords)
	assert.Equal(t, ProviderIPAPI, g.Source)

	// status != success 才是「服务明确说没有」，才可以被缓存成查不到
	assert.True(t, res.Unresolved["203.0.113.7"])
	assert.NotContains(t, res.Geo, "203.0.113.7")
}

// TestLookupNetworkFailure 三态里的「没查成」：绝不能进 Unresolved。
//
// 这条断言是本包最重要的回归保护：一旦失败也进 Unresolved，上层就会把它
// 写进缓存（3 天），此后每份报告的定位覆盖率都会莫名其妙变成 0。
func TestLookupNetworkFailure(t *testing.T) {
	srv, calls := ipAPIServer(t, func([]string) ([]ipAPIResult, int) {
		return nil, http.StatusBadGateway
	})

	r := New(Options{Provider: ProviderIPAPI, HTTPClient: newTestClient(t, srv)})
	res := r.Lookup(context.Background(), []string{"8.8.4.4", "1.0.0.1"})

	assert.Equal(t, int32(ipAPIAttempts), atomic.LoadInt32(calls), "必须重试，但只重试一次")
	assert.Equal(t, 1, res.Requests)
	assert.Equal(t, 1, res.Failures)
	assert.Empty(t, res.Geo)
	assert.Empty(t, res.Unresolved, "网络失败不能记成「这个 IP 没有位置」")
}

// TestLookupRetrySucceeds 第一次抖动、第二次正常：不该丢数据。
func TestLookupRetrySucceeds(t *testing.T) {
	var attempt int32
	srv, _ := ipAPIServer(t, func(addrs []string) ([]ipAPIResult, int) {
		if atomic.AddInt32(&attempt, 1) == 1 {
			return nil, http.StatusGatewayTimeout
		}
		return []ipAPIResult{okIP(addrs[0], "日本", "JP", "Tokyo", 35.68, 139.69, "", "")}, 0
	})

	r := New(Options{Provider: ProviderIPAPI, HTTPClient: newTestClient(t, srv)})
	res := r.Lookup(context.Background(), []string{"133.242.0.1"})

	require.Equal(t, 0, res.Failures)
	assert.Equal(t, "日本", res.Geo["133.242.0.1"].Country)
}

// TestLookupBatches 超过 100 个地址必须分批，而不是一次塞进去被拒。
func TestLookupBatches(t *testing.T) {
	srv, calls := ipAPIServer(t, func(addrs []string) ([]ipAPIResult, int) {
		assert.LessOrEqual(t, len(addrs), ipAPIBatchSize, "单批不能超过 ip-api 的硬上限")
		out := make([]ipAPIResult, 0, len(addrs))
		for _, a := range addrs {
			out = append(out, okIP(a, "美国", "US", "Ashburn", 39.0, -77.5, "", ""))
		}
		return out, 0
	})

	addrs := make([]string, 0, 150)
	for i := 0; i < 150; i++ {
		addrs = append(addrs, fmt.Sprintf("198.51.%d.%d", i/256, i%256))
	}

	r := New(Options{Provider: ProviderIPAPI, HTTPClient: newTestClient(t, srv)})
	res := r.Lookup(context.Background(), addrs)

	assert.Equal(t, int32(2), atomic.LoadInt32(calls))
	assert.Equal(t, 2, res.Requests)
	assert.Len(t, res.Geo, 150)
}

// TestLookupSkipsPrivate 私网 / 回环 / CGNAT 一律不外发：既不泄露内网结构，
// 也省下配额（这类地址查了也是空的）。
func TestLookupSkipsPrivate(t *testing.T) {
	srv, calls := ipAPIServer(t, func([]string) ([]ipAPIResult, int) {
		return nil, http.StatusInternalServerError
	})

	r := New(Options{Provider: ProviderIPAPI, HTTPClient: newTestClient(t, srv)})
	res := r.Lookup(context.Background(), []string{
		"10.0.0.1", "192.168.1.1", "172.16.0.9", "127.0.0.1",
		"100.64.3.4", "169.254.1.1", "::1", "fe80::1", "fd00::1",
	})

	assert.Equal(t, int32(0), atomic.LoadInt32(calls), "不可查询的地址不该产生任何请求")
	assert.Equal(t, 0, res.Requests)
	assert.Empty(t, res.Geo)
	assert.Empty(t, res.Unresolved)
}

// TestLookupNone provider = none：一个字节都不外发。
func TestLookupNone(t *testing.T) {
	srv, calls := ipAPIServer(t, func([]string) ([]ipAPIResult, int) {
		return nil, http.StatusInternalServerError
	})
	r := New(Options{Provider: ProviderNone, HTTPClient: newTestClient(t, srv)})

	assert.False(t, r.Enabled())
	res := r.Lookup(context.Background(), []string{"1.2.3.4"})
	assert.Equal(t, int32(0), atomic.LoadInt32(calls))
	assert.Empty(t, res.Geo)
}

// TestLookupEmptyKeys 空输入不该发请求。
func TestLookupEmptyKeys(t *testing.T) {
	srv, calls := ipAPIServer(t, func([]string) ([]ipAPIResult, int) {
		return nil, http.StatusInternalServerError
	})
	r := New(Options{Provider: ProviderIPAPI, HTTPClient: newTestClient(t, srv)})
	assert.Empty(t, r.Lookup(context.Background(), nil).Geo)
	assert.Equal(t, int32(0), atomic.LoadInt32(calls))
}

// TestLookupIPv6Prefix IPv6 的 "/64" 缓存键要还原成基址去查。
func TestLookupIPv6Prefix(t *testing.T) {
	srv, _ := ipAPIServer(t, func(addrs []string) ([]ipAPIResult, int) {
		require.Len(t, addrs, 1)
		assert.Equal(t, "2001:4860:4860::", addrs[0], "带 /64 后缀的键必须剥掉后缀再查")
		return []ipAPIResult{okIP(addrs[0], "美国", "US", "Mountain View", 37.4, -122.0, "Google", "AS15169 Google LLC")}, 0
	})

	r := New(Options{Provider: ProviderIPAPI, HTTPClient: newTestClient(t, srv)})
	res := r.Lookup(context.Background(), []string{"2001:4860:4860::/64"})

	g, ok := res.Geo["2001:4860:4860::/64"]
	require.True(t, ok, "结果键必须仍是原来的缓存键，不能变成查询用的地址")
	assert.Equal(t, "US", g.CountryCode)
	assert.Equal(t, "AS15169", g.ASN)
}

// tencentServer 起一个假的腾讯 IP 定位接口。
func tencentServer(t *testing.T, respond func(ip string) (tencentResult, int)) (*httptest.Server, *int32) {
	t.Helper()
	var calls int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		if req.URL.Path != "/ws/location/v1/ip" {
			http.NotFound(w, req)
			return
		}
		atomic.AddInt32(&calls, 1)
		tr, code := respond(req.URL.Query().Get("ip"))
		if code != 0 {
			http.Error(w, "upstream failure", code)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(tr)
	}))
	t.Cleanup(srv.Close)
	return srv, &calls
}

// combinedServer 同时提供两条通道的测试服务器（auto 模式要两条都通）。
func combinedServer(t *testing.T, ipapi func([]string) ([]ipAPIResult, int), tencent func(ip string) (tencentResult, int)) (*httptest.Server, *int32, *int32) {
	t.Helper()
	var ipCalls, txCalls int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		switch req.URL.Path {
		case "/batch":
			atomic.AddInt32(&ipCalls, 1)
			var addrs []string
			if err := json.NewDecoder(req.Body).Decode(&addrs); err != nil {
				http.Error(w, "bad body", http.StatusBadRequest)
				return
			}
			results, code := ipapi(addrs)
			if code != 0 {
				http.Error(w, "upstream failure", code)
				return
			}
			w.Header().Set("Content-Type", "application/json")
			_ = json.NewEncoder(w).Encode(results)
		case "/ws/location/v1/ip":
			atomic.AddInt32(&txCalls, 1)
			tr, code := tencent(req.URL.Query().Get("ip"))
			if code != 0 {
				http.Error(w, "upstream failure", code)
				return
			}
			w.Header().Set("Content-Type", "application/json")
			_ = json.NewEncoder(w).Encode(tr)
		default:
			http.NotFound(w, req)
		}
	}))
	t.Cleanup(srv.Close)
	return srv, &ipCalls, &txCalls
}

func tencentOK(nation, province, city string, lat, lon float64) tencentResult {
	var tr tencentResult
	tr.Status = 0
	tr.Result.Location.Lat, tr.Result.Location.Lng = lat, lon
	tr.Result.AdInfo.Nation = nation
	tr.Result.AdInfo.Province = province
	tr.Result.AdInfo.City = city
	return tr
}

// TestLookupAutoOnlyRefinesChina auto 模式只对国内 IP 调腾讯：
// 海外 IP 腾讯给不出省市，查了也是白花配额。
func TestLookupAutoOnlyRefinesChina(t *testing.T) {
	srv, _, txCalls := combinedServer(t,
		func(addrs []string) ([]ipAPIResult, int) {
			return []ipAPIResult{
				okIP(addrs[0], "中国", "CN", "", 0, 0, "", ""),
				okIP(addrs[1], "美国", "US", "Ashburn", 39.0, -77.5, "", ""),
			}, 0
		},
		func(ip string) (tencentResult, int) {
			return tencentOK("中国", "北京市", "北京市", 39.9087, 116.3975), 0
		})

	r := New(Options{Provider: ProviderAuto, TencentKey: "k", HTTPClient: newTestClient(t, srv)})
	res := r.Lookup(context.Background(), []string{"114.114.114.114", "45.32.0.1"})

	assert.Equal(t, int32(1), atomic.LoadInt32(txCalls), "只该精修 1 个国内 IP")

	cn := res.Geo["114.114.114.114"]
	assert.Equal(t, "北京市", cn.Region)
	assert.Equal(t, "北京市", cn.City)
	assert.Equal(t, "CN", cn.CountryCode)
	assert.InDelta(t, 39.9087, cn.Lat, 1e-9)
	assert.Equal(t, ProviderTencent, cn.Source)

	us := res.Geo["45.32.0.1"]
	assert.Equal(t, "US", us.CountryCode)
	assert.Equal(t, "Ashburn", us.City, "海外 IP 应保留 ip-api 的结果")
	assert.Equal(t, ProviderIPAPI, us.Source)
}

// TestLookupTencentRefineCap 精修数量必须有上限，否则周报会跑十几分钟。
//
// 顺带守住一个容易踩的坑：精修必须按调用方给的顺序取前 N 个。报告是按命中
// 次数从高到低传进来的，若这里按 map 遍历顺序取，每次运行精修到的 IP 都不一样，
// 配额会花在长尾地址上，「主要从哪个城市来」这个问题就答不准了。
func TestLookupTencentRefineCap(t *testing.T) {
	var mu sync.Mutex
	var refined []string
	srv, _, _ := combinedServer(t,
		func(addrs []string) ([]ipAPIResult, int) {
			out := make([]ipAPIResult, 0, len(addrs))
			for _, a := range addrs {
				out = append(out, okIP(a, "中国", "CN", "", 0, 0, "", ""))
			}
			return out, 0
		},
		func(ip string) (tencentResult, int) {
			mu.Lock()
			refined = append(refined, ip)
			mu.Unlock()
			return tencentOK("中国", "广东省", "深圳市", 22.5, 114.0), 0
		})

	addrs := []string{"1.1.1.1", "2.2.2.2", "3.3.3.3", "4.4.4.4", "5.5.5.5"}
	r := New(Options{
		Provider: ProviderAuto, TencentKey: "k", TencentRefineCap: 2,
		HTTPClient: newTestClient(t, srv),
	})
	res := r.Lookup(context.Background(), addrs)

	require.Len(t, refined, 2)
	assert.ElementsMatch(t, []string{"1.1.1.1", "2.2.2.2"}, refined,
		"必须精修前 2 个（最活跃的），而不是随机的 2 个")

	// 精修到的用腾讯结果
	assert.Equal(t, ProviderTencent, res.Geo["1.1.1.1"].Source)
	// 没精修到的仍保留 ip-api 的国家
	assert.Equal(t, "CN", res.Geo["5.5.5.5"].CountryCode)
	assert.Equal(t, ProviderIPAPI, res.Geo["5.5.5.5"].Source)
}

// TestLookupTencentFailureNotCached 腾讯配额用尽 / key 无效属于「没查成」，
// 不能把整批 IP 记成「查不到」。ip-api 已经给出的结果也不能被抹掉。
func TestLookupTencentFailureNotCached(t *testing.T) {
	srv, _, _ := combinedServer(t,
		func(addrs []string) ([]ipAPIResult, int) {
			return []ipAPIResult{okIP(addrs[0], "中国", "CN", "", 0, 0, "", "")}, 0
		},
		func(string) (tencentResult, int) {
			var tr tencentResult
			tr.Status = 311
			tr.Message = "key格式错误"
			return tr, 0
		})

	r := New(Options{Provider: ProviderAuto, TencentKey: "bad", HTTPClient: newTestClient(t, srv)})
	res := r.Lookup(context.Background(), []string{"114.114.114.114"})

	assert.Equal(t, 1, res.Failures)
	assert.Empty(t, res.Unresolved, "配额/key 问题不是「这个 IP 没数据」")
	assert.Equal(t, "CN", res.Geo["114.114.114.114"].CountryCode, "ip-api 的结果不能被腾讯的失败抹掉")
}

// TestLookupTencentNoData 状态成功但没有任何行政区划 = 真的没数据。
// 但若 ip-api 已经定位到了，也不能反过来把这条记成「查不到」。
func TestLookupTencentNoData(t *testing.T) {
	srv, _, _ := combinedServer(t,
		func(addrs []string) ([]ipAPIResult, int) {
			return []ipAPIResult{okIP(addrs[0], "中国", "CN", "", 0, 0, "", "")}, 0
		},
		func(string) (tencentResult, int) { return tencentOK("中国", "", "", 0, 0), 0 })

	r := New(Options{Provider: ProviderAuto, TencentKey: "k", HTTPClient: newTestClient(t, srv)})
	res := r.Lookup(context.Background(), []string{"114.114.114.114"})

	assert.Equal(t, 0, res.Failures)
	assert.NotContains(t, res.Unresolved, "114.114.114.114", "已定位的 IP 不该被降级成查不到")
	assert.Equal(t, "CN", res.Geo["114.114.114.114"].CountryCode)
}

// TestTencentOnly 只用腾讯（provider = tencent）：不经 ip-api，全部走腾讯。
func TestTencentOnly(t *testing.T) {
	srv, _, txCalls := combinedServer(t,
		func([]string) ([]ipAPIResult, int) {
			t.Error("provider = tencent 时不该请求 ip-api")
			return nil, 0
		},
		func(string) (tencentResult, int) {
			return tencentOK("中国", "浙江省", "杭州市", 30.27, 120.15), 0
		})

	r := New(Options{Provider: ProviderTencent, TencentKey: "k", HTTPClient: newTestClient(t, srv)})
	res := r.Lookup(context.Background(), []string{"115.236.1.1"})

	assert.Equal(t, int32(1), atomic.LoadInt32(txCalls))
	assert.Equal(t, "杭州市", res.Geo["115.236.1.1"].City)
	assert.Equal(t, "CN", res.Geo["115.236.1.1"].CountryCode)
}

// TestNewFallbacks 配置错误不能表现为「静默不工作」。
func TestNewFallbacks(t *testing.T) {
	// 只选 tencent 却没 key → 降级 ipapi
	r := New(Options{Provider: ProviderTencent})
	assert.Equal(t, ProviderIPAPI, r.provider)
	assert.True(t, r.Enabled())

	// 空值 → auto
	assert.Equal(t, ProviderAuto, New(Options{}).provider)

	// 不认识的值 → auto
	assert.Equal(t, ProviderAuto, New(Options{Provider: "maxmind"}).provider)

	// none 不工作
	assert.False(t, New(Options{Provider: ProviderNone}).Enabled())

	// refineCap 缺省
	assert.Equal(t, defaultTencentRefineCap, New(Options{}).refineCap)
}

// TestDescribe 描述会写进报告，读者靠它判断坐标可信度。
func TestDescribe(t *testing.T) {
	assert.Contains(t, New(Options{Provider: ProviderAuto, TencentKey: "k"}).Describe(), "腾讯")
	assert.Contains(t, New(Options{Provider: ProviderAuto}).Describe(), "未配置腾讯 key")
	assert.Contains(t, New(Options{Provider: ProviderIPAPI}).Describe(), "ip-api")
	assert.Contains(t, New(Options{Provider: ProviderNone}).Describe(), "未启用")
}

// TestQueryAddr 缓存键还原。
func TestQueryAddr(t *testing.T) {
	cases := []struct {
		key  string
		want string
		ok   bool
	}{
		{"8.8.8.8", "8.8.8.8", true},
		{"2001:4860:4860::/64", "2001:4860:4860::", true},
		{"2001:4860:4860::8888", "2001:4860:4860::8888", true},
		{"10.1.2.3", "", false},
		{"::ffff:8.8.8.8", "8.8.8.8", true},
		{"not-an-ip", "", false},
		{"", "", false},
	}
	for _, c := range cases {
		got, ok := queryAddr(c.key)
		assert.Equal(t, c.ok, ok, "key = %q", c.key)
		assert.Equal(t, c.want, got, "key = %q", c.key)
	}
}

// TestNormalizeASN ASN 只留编号。
func TestNormalizeASN(t *testing.T) {
	assert.Equal(t, "AS15169", normalizeASN("AS15169 Google LLC"))
	assert.Equal(t, "AS4134", normalizeASN("  AS4134  "))
	assert.Equal(t, "AS4134", normalizeASN("AS4134"))
	assert.Equal(t, "", normalizeASN(""))
}

// TestTencentCountryCode 不能把海外 IP 算成中国来源。
func TestTencentCountryCode(t *testing.T) {
	assert.Equal(t, "CN", tencentCountryCode("中国"))
	assert.Equal(t, "", tencentCountryCode("中国香港"))
	assert.Equal(t, "", tencentCountryCode("美国"))
	assert.Equal(t, "", tencentCountryCode(""))
}