package utility import ( "context" "errors" "fmt" "net/http" "net/http/httptest" "strings" "sync" "sync/atomic" "testing" "time" "github.com/gogf/gf/v2/encoding/gcharset" ) func waitForIPLookup(t *testing.T, lookup *asyncIPLookup, key string) { t.Helper() deadline := time.Now().Add(2 * time.Second) for time.Now().Before(deadline) { lookup.mutex.Lock() pending := lookup.pending[key] lookup.mutex.Unlock() if !pending { return } time.Sleep(time.Millisecond) } t.Fatal("background lookup did not finish") } func TestAsyncIPLookupDoesNotWaitAndDeduplicates(t *testing.T) { var calls atomic.Int32 started := make(chan struct{}) release := make(chan struct{}) lookup := newAsyncIPLookup(func(ctx context.Context, key string) (string, error) { calls.Add(1) close(started) select { case <-release: return "Shandong Zibo", nil case <-ctx.Done(): return "", ctx.Err() } }, time.Second, time.Hour, time.Minute, 4) defer lookup.cache.Close(context.Background()) defer close(release) if value := lookup.Get("same-ip"); value != "" { t.Fatalf("unexpected initial value: %q", value) } select { case <-started: case <-time.After(time.Second): t.Fatal("background lookup did not start") } var requests sync.WaitGroup for index := 0; index < 50; index++ { requests.Add(1) go func() { defer requests.Done() lookup.Get("same-ip") }() } requests.Wait() if calls.Load() != 1 { t.Fatalf("duplicate external queries: %d", calls.Load()) } release <- struct{}{} waitForIPLookup(t, lookup, "same-ip") if value := lookup.Get("same-ip"); value != "Shandong Zibo" { t.Fatalf("cached value missing: %q", value) } if calls.Load() != 1 { t.Fatal("cache hit triggered a new query") } } func TestAsyncIPLookupFailureKeepsStaleValueAndBacksOff(t *testing.T) { var calls atomic.Int32 lookup := newAsyncIPLookup(func(ctx context.Context, key string) (string, error) { calls.Add(1) return "", errors.New("provider unavailable") }, time.Second, time.Hour, time.Minute, 1) defer lookup.cache.Close(context.Background()) if err := lookup.cache.Set(context.Background(), "ip", ipLookupEntry{ Value: "previous location", RefreshAt: time.Now().Add(-time.Minute), }, time.Hour); err != nil { t.Fatal(err) } if value := lookup.Get("ip"); value != "previous location" { t.Fatalf("stale value was not returned immediately: %q", value) } waitForIPLookup(t, lookup, "ip") for index := 0; index < 20; index++ { if lookup.Get("ip") != "previous location" { t.Fatal("failed refresh discarded the last successful value") } } if calls.Load() != 1 { t.Fatalf("failed provider was repeatedly queried: %d", calls.Load()) } } func TestAsyncIPLookupTimeoutAndRecovery(t *testing.T) { var calls atomic.Int32 lookup := newAsyncIPLookup(func(ctx context.Context, key string) (string, error) { if calls.Add(1) == 1 { <-ctx.Done() return "", ctx.Err() } return "recovered", nil }, 30*time.Millisecond, time.Hour, time.Minute, 1) defer lookup.cache.Close(context.Background()) lookup.Get("ip") waitForIPLookup(t, lookup, "ip") lookup.Get("ip") if calls.Load() != 1 { t.Fatal("timeout did not activate the retry cooldown") } if err := lookup.cache.Set(context.Background(), "ip", ipLookupEntry{}, time.Hour); err != nil { t.Fatal(err) } lookup.Get("ip") waitForIPLookup(t, lookup, "ip") if lookup.Get("ip") != "recovered" || calls.Load() != 2 { t.Fatal("lookup did not recover after the retry cooldown") } } func TestAsyncIPLookupLimitsConcurrentQueries(t *testing.T) { var calls atomic.Int32 release := make(chan struct{}) lookup := newAsyncIPLookup(func(ctx context.Context, key string) (string, error) { calls.Add(1) select { case <-release: return key, nil case <-ctx.Done(): return "", ctx.Err() } }, time.Second, time.Hour, time.Minute, 2) defer lookup.cache.Close(context.Background()) for index := 0; index < 50; index++ { lookup.Get(fmt.Sprintf("ip-%d", index)) } lookup.mutex.Lock() pending := len(lookup.pending) lookup.mutex.Unlock() close(release) waitForIPLookup(t, lookup, "ip-0") waitForIPLookup(t, lookup, "ip-1") if pending != 2 || calls.Load() != 2 { t.Fatalf("query concurrency was not bounded: pending=%d calls=%d", pending, calls.Load()) } } func TestGetCityByIPSkipsPrivateAndInvalidAddresses(t *testing.T) { for _, ip := range []string{"127.0.0.1", "[::1]", "::1", "10.1.2.3", "172.23.0.6", "192.168.100.29", "fe80::1", "0.0.0.0"} { if got := GetCityByIp(ip); got != "内网IP" { t.Errorf("%s: got %q", ip, got) } } for _, ip := range []string{"", "invalid", "8.8.8.8&json=false"} { if got := GetCityByIp(ip); got != "" { t.Errorf("invalid address %q returned %q", ip, got) } } } func TestParseCityByIPResponse(t *testing.T) { for _, payload := range []string{"not JSON", `{"pro":"Unknown"}`, `{"code":1,"pro":"Unknown"}`} { if _, err := parseCityByIPResponse([]byte(payload)); err == nil { t.Errorf("invalid response accepted: %s", payload) } } encoded, err := gcharset.Convert("GBK", "UTF-8", `{"code":0,"pro":"山东省","city":"淄博市"}`) if err != nil { t.Fatal(err) } got, err := parseCityByIPResponse([]byte(encoded)) if err != nil || got != "山东省 淄博市" { t.Fatalf("GBK location parsing failed: %q, %v", got, err) } } func TestIPLookupResponseBodyTimeout(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { writer.WriteHeader(http.StatusOK) writer.(http.Flusher).Flush() <-request.Context().Done() })) defer server.Close() ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) defer cancel() started := time.Now() _, err := readIPLookupResponse(ctx, server.URL) if err == nil || time.Since(started) > time.Second { t.Fatalf("stalled response body was not interrupted: %v", err) } } func TestIPLookupRejectsErrorStatusAndOversizedBody(t *testing.T) { for _, scenario := range []string{"status", "body"} { t.Run(scenario, func(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { if scenario == "status" { writer.WriteHeader(http.StatusServiceUnavailable) return } _, _ = writer.Write([]byte(strings.Repeat("x", 64*1024+1))) })) defer server.Close() if _, err := readIPLookupResponse(context.Background(), server.URL); err == nil { t.Fatal("invalid provider response accepted") } }) } }