geo/server/internal/shared/ipregion/ip_region.go

package ipregion

import (
	"embed"
	"errors"
	"fmt"
	"net"
	"os"
	"strings"
	"sync"

	"go.uber.org/zap"

	"github.com/lionsoul2014/ip2region/binding/golang/xdb"
)

//go:embed ipregiondata/ip2region_v4.xdb ipregiondata/ip2region_v6.xdb
var embeddedIPRegionFS embed.FS

const (
	embeddedV4XDBPath = "ipregiondata/ip2region_v4.xdb"
	embeddedV6XDBPath = "ipregiondata/ip2region_v6.xdb"
)

type Resolver struct {
	searcher searcher
	logger   *zap.Logger
}

type searcher interface {
	Search(ip any) (string, error)
	Close()
}

func NewResolver(v4XDBPath, v6XDBPath string, logger *zap.Logger) (*Resolver, error) {
	searcher, err := newSearcher(v4XDBPath, v6XDBPath)
	if err != nil {
		return nil, err
	}
	return &Resolver{searcher: searcher, logger: logger}, nil
}

func newSearcher(v4XDBPath, v6XDBPath string) (*bufferSearcher, error) {
	v4Content, err := loadContent(strings.TrimSpace(v4XDBPath), embeddedV4XDBPath, xdb.IPv4)
	if err != nil {
		return nil, fmt.Errorf("load ipv4 xdb: %w", err)
	}
	v6Content, err := loadContent(strings.TrimSpace(v6XDBPath), embeddedV6XDBPath, xdb.IPv6)
	if err != nil {
		return nil, fmt.Errorf("load ipv6 xdb: %w", err)
	}

	searcher := &bufferSearcher{}
	if len(v4Content) > 0 {
		searcher.v4, err = xdb.NewWithBuffer(xdb.IPv4, v4Content)
		if err != nil {
			return nil, fmt.Errorf("create ipv4 searcher: %w", err)
		}
	}
	if len(v6Content) > 0 {
		searcher.v6, err = xdb.NewWithBuffer(xdb.IPv6, v6Content)
		if err != nil {
			return nil, fmt.Errorf("create ipv6 searcher: %w", err)
		}
	}
	return searcher, nil
}

func loadContent(externalPath, embeddedPath string, version *xdb.Version) ([]byte, error) {
	if externalPath != "" {
		content, err := loadContentFromFile(externalPath, version)
		if err == nil {
			return content, nil
		}
		if !errors.Is(err, os.ErrNotExist) {
			return nil, err
		}
	}

	content, err := xdb.LoadContentFromFS(embeddedIPRegionFS, embeddedPath)
	if err != nil {
		return nil, err
	}
	if err := verifyVersion(content, version); err != nil {
		return nil, fmt.Errorf("verify embedded xdb %s: %w", embeddedPath, err)
	}
	return content, nil
}

func loadContentFromFile(path string, version *xdb.Version) ([]byte, error) {
	content, err := xdb.LoadContentFromFile(path)
	if err != nil {
		return nil, err
	}
	if err := verifyVersion(content, version); err != nil {
		return nil, fmt.Errorf("verify external xdb %s: %w", path, err)
	}
	return content, nil
}

func verifyVersion(content []byte, expected *xdb.Version) error {
	if len(content) < xdb.HeaderInfoLength {
		return fmt.Errorf("xdb content too short: got %d bytes", len(content))
	}
	header, err := xdb.LoadHeaderFromBuff(content)
	if err != nil {
		return fmt.Errorf("load header: %w", err)
	}
	version, err := xdb.VersionFromHeader(header)
	if err != nil {
		return err
	}
	if version.Id != expected.Id {
		return fmt.Errorf("ip version mismatch: got %s, want %s", version.Name, expected.Name)
	}
	return nil
}

func (r *Resolver) Lookup(rawIP string) string {
	ipText := strings.TrimSpace(rawIP)
	if ipText == "" {
		return ""
	}
	if host, _, err := net.SplitHostPort(ipText); err == nil {
		ipText = host
	}

	ip := net.ParseIP(ipText)
	if ip == nil {
		return ""
	}
	if ip.IsLoopback() {
		return "本机"
	}
	if ip.IsPrivate() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() {
		return "内网IP"
	}
	if r == nil || r.searcher == nil {
		return ""
	}

	region, err := r.searcher.Search(ip.String())
	if err != nil {
		if r.logger != nil {
			r.logger.Warn("ip2region lookup failed",
				zap.String("ip", ip.String()),
				zap.Error(err),
			)
		}
		return ""
	}
	return strings.TrimSpace(region)
}

func (r *Resolver) Close() {
	if r != nil && r.searcher != nil {
		r.searcher.Close()
	}
}

type bufferSearcher struct {
	v4   *xdb.Searcher
	v4Mu sync.Mutex
	v6   *xdb.Searcher
	v6Mu sync.Mutex
}

func (s *bufferSearcher) Search(ip any) (string, error) {
	ipBytes, err := parseIPBytes(ip)
	if err != nil {
		return "", err
	}
	switch len(ipBytes) {
	case xdb.IPv4.Bytes:
		if s == nil || s.v4 == nil {
			return "", nil
		}
		s.v4Mu.Lock()
		defer s.v4Mu.Unlock()
		return s.v4.Search(ipBytes)
	case xdb.IPv6.Bytes:
		if s == nil || s.v6 == nil {
			return "", nil
		}
		s.v6Mu.Lock()
		defer s.v6Mu.Unlock()
		return s.v6.Search(ipBytes)
	default:
		return "", fmt.Errorf("invalid byte ip address with len=%d", len(ipBytes))
	}
}

func (s *bufferSearcher) Close() {
	if s == nil {
		return
	}
	if s.v4 != nil {
		s.v4.Close()
	}
	if s.v6 != nil {
		s.v6.Close()
	}
}

func parseIPBytes(ip any) ([]byte, error) {
	switch v := ip.(type) {
	case string:
		return xdb.ParseIP(v)
	case []byte:
		return v, nil
	default:
		return nil, fmt.Errorf("invalid ip value type %T", v)
	}
}