iotlcm/admin-svr/server/utility/ip_lookup_test.go

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", "127.0.0.1", "127.0.0.1", "127.0.0.1", "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")
			}
		})
	}
}