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

package utility

import (
	"context"
	"encoding/json"
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/gogf/gf/v2/encoding/gcharset"
	"github.com/gogf/gf/v2/errors/gerror"
	"github.com/gogf/gf/v2/os/gcache"
)

type ipLookupEntry struct {
	Value     string
	RefreshAt time.Time
}

type asyncIPLookup struct {
	cache      *gcache.Cache
	mutex      sync.Mutex
	pending    map[string]bool
	workers    chan struct{}
	lookup     func(context.Context, string) (string, error)
	timeout    time.Duration
	ttl        time.Duration
	retryDelay time.Duration
}

var (
	cityLookup     = newAsyncIPLookup(lookupCityByIP, time.Second, 24*time.Hour, time.Minute, 4)
	publicIPLookup = newAsyncIPLookup(func(ctx context.Context, _ string) (string, error) {
		return lookupPublicIP(ctx)
	}, 2*time.Second, time.Hour, time.Minute, 1)
)

func newAsyncIPLookup(lookup func(context.Context, string) (string, error), timeout, ttl, retryDelay time.Duration, workers int) *asyncIPLookup {
	return &asyncIPLookup{
		cache:      gcache.New(4096),
		pending:    make(map[string]bool),
		workers:    make(chan struct{}, workers),
		lookup:     lookup,
		timeout:    timeout,
		ttl:        ttl,
		retryDelay: retryDelay,
	}
}

func (lookup *asyncIPLookup) Get(key string) string {
	lookup.mutex.Lock()
	defer lookup.mutex.Unlock()
	entry := ipLookupEntry{}
	if value, err := lookup.cache.Get(context.Background(), key); err == nil && value != nil {
		entry, _ = value.Val().(ipLookupEntry)
	}
	if time.Now().Before(entry.RefreshAt) || lookup.pending[key] {
		return entry.Value
	}
	select {
	case lookup.workers <- struct{}{}:
		lookup.pending[key] = true
		go lookup.refresh(key, entry)
	default:
	}
	return entry.Value
}

func (lookup *asyncIPLookup) refresh(key string, entry ipLookupEntry) {
	ctx, cancel := context.WithTimeout(context.Background(), lookup.timeout)
	defer cancel()
	value, err := lookup.lookup(ctx, key)
	delay := lookup.retryDelay
	if err == nil && value != "" {
		entry.Value = value
		delay = lookup.ttl
	}
	entry.RefreshAt = time.Now().Add(delay)
	retention := delay
	if entry.Value != "" {
		retention = 2 * lookup.ttl
	}
	lookup.mutex.Lock()
	defer lookup.mutex.Unlock()
	_ = lookup.cache.Set(context.Background(), key, entry, retention)
	delete(lookup.pending, key)
	<-lookup.workers
}

func GetCachedPublicIP() string {
	return publicIPLookup.Get("public-ip")
}

func lookupCityByIP(ctx context.Context, ip string) (string, error) {
	query := url.Values{"json": {"true"}, "ip": {ip}}
	payload, err := readIPLookupResponse(ctx, "https://whois.pconline.com.cn/ipJson.jsp?"+query.Encode())
	if err != nil {
		return "", err
	}
	return parseCityByIPResponse(payload)
}

func parseCityByIPResponse(payload []byte) (string, error) {
	text, err := gcharset.ToUTF8("GBK", string(payload))
	if err != nil {
		return "", gerror.Wrap(err, "decode IP location response failed")
	}
	var response struct {
		Code     *int   `json:"code"`
		Province string `json:"pro"`
		City     string `json:"city"`
	}
	if err := json.Unmarshal([]byte(text), &response); err != nil {
		return "", gerror.Wrap(err, "parse IP location response failed")
	}
	if response.Code == nil || *response.Code != 0 {
		return "", gerror.New("IP location lookup failed")
	}
	if strings.EqualFold(response.Province, response.City) {
		return strings.TrimSpace(response.Province), nil
	}
	return strings.TrimSpace(response.Province + " " + response.City), nil
}

func lookupPublicIP(ctx context.Context) (string, error) {
	payload, err := readIPLookupResponse(ctx, "https://ifconfig.co/ip")
	if err != nil {
		return "", err
	}
	ip := net.ParseIP(strings.TrimSpace(string(payload)))
	if ip == nil || !ip.IsGlobalUnicast() || ip.IsPrivate() {
		return "", gerror.New("invalid public IP response")
	}
	return ip.String(), nil
}

func readIPLookupResponse(ctx context.Context, endpoint string) ([]byte, error) {
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil)
	if err != nil {
		return nil, gerror.Wrap(err, "create IP lookup request failed")
	}
	client := &http.Client{Timeout: 2 * time.Second}
	response, err := client.Do(request)
	if err != nil {
		return nil, gerror.Wrap(err, "IP lookup request failed")
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		return nil, gerror.Newf("IP lookup HTTP status: %d", response.StatusCode)
	}
	const maxBodySize = 64 * 1024
	payload, err := io.ReadAll(io.LimitReader(response.Body, maxBodySize+1))
	if err != nil {
		return nil, gerror.Wrap(err, "read IP lookup response failed")
	}
	if len(payload) > maxBodySize {
		return nil, gerror.New("IP lookup response is too large")
	}
	return payload, nil
}