[优化] go 引用调整

2026-10-07 08:06:37 +08:00 · 2026-06-06 10:26:20 +08:00
parent ee1110b752
commit 3cfefb4367
552 changed files with 1642 additions and 2185 deletions
@@ -0,0 +1,171 @@
+package geoip
+
+import (
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"os"
+	"path/filepath"
+	"sync"
+
+	"github.com/oschwald/maxminddb-golang"
+)
+
+var GeoIpUrl = "https://raw.githubusercontent.com/Loyalsoldier/geoip/release/GeoLite2-Country.mmdb"
+var GeoIpFilePath = "./data/GeoLite2-Country.mmdb"
+
+type GeoIpRecord struct {
+	Country struct {
+		ISOCode string            `maxminddb:"iso_code"`
+		Names   map[string]string `maxminddb:"names"`
+	} `maxminddb:"country"`
+}
+
+type MaxMindGeoIPService struct {
+	maxMindDBReader *maxminddb.Reader
+	dbFilePath      string
+	mu              sync.RWMutex
+}
+
+func (s *MaxMindGeoIPService) Name() string {
+	return "MaxMind"
+}
+
+func NewMaxMindGeoIPService() (*MaxMindGeoIPService, error) {
+	return NewMaxMindGeoIPServiceWithConfig(GeoIpFilePath, GeoIpUrl)
+}
+
+func NewMaxMindGeoIPServiceWithConfig(dbFilePath string, downloadURL string) (*MaxMindGeoIPService, error) {
+	if dbFilePath == "" {
+		dbFilePath = GeoIpFilePath
+	}
+	if downloadURL == "" {
+		downloadURL = GeoIpUrl
+	}
+	service := &MaxMindGeoIPService{
+		dbFilePath: dbFilePath,
+	}
+
+	if err := os.MkdirAll(filepath.Dir(service.dbFilePath), os.ModePerm); err != nil {
+		return nil, fmt.Errorf("failed to create data directory for MaxMind database: %w", err)
+	}
+
+	if _, err := os.Stat(service.dbFilePath); os.IsNotExist(err) {
+		if err := DownloadMaxMindDatabase(service.dbFilePath, downloadURL); err != nil {
+			return nil, fmt.Errorf("failed to download initial MaxMind database: %w", err)
+		}
+	}
+
+	if err := service.initialize(); err != nil {
+		return nil, fmt.Errorf("failed to initialize MaxMind database: %w", err)
+	}
+
+	return service, nil
+}
+
+func (s *MaxMindGeoIPService) initialize() error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if s.maxMindDBReader != nil {
+		_ = s.maxMindDBReader.Close()
+		s.maxMindDBReader = nil
+	}
+
+	reader, err := maxminddb.Open(s.dbFilePath)
+	if err != nil {
+		return fmt.Errorf("error opening MaxMind database at %s: %w", s.dbFilePath, err)
+	}
+	s.maxMindDBReader = reader
+	return nil
+}
+
+func (s *MaxMindGeoIPService) GetGeoInfo(ip net.IP) (*GeoInfo, error) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	if s.maxMindDBReader == nil {
+		return nil, fmt.Errorf("MaxMind database is not initialized or failed to open")
+	}
+	if ip == nil {
+		return nil, fmt.Errorf("IP address cannot be nil")
+	}
+
+	var record GeoIpRecord
+	if err := s.maxMindDBReader.Lookup(ip, &record); err != nil {
+		return nil, fmt.Errorf("error looking up IP %s in MaxMind database: %w", ip.String(), err)
+	}
+
+	geoInfo := &GeoInfo{
+		ISOCode: record.Country.ISOCode,
+		Name:    record.Country.Names["en"],
+	}
+	if geoInfo.Name == "" && geoInfo.ISOCode != "" {
+		geoInfo.Name = geoInfo.ISOCode
+	}
+
+	return geoInfo, nil
+}
+
+func (s *MaxMindGeoIPService) UpdateDatabase() error {
+	if err := DownloadMaxMindDatabase(s.dbFilePath, GeoIpUrl); err != nil {
+		return err
+	}
+	return s.initialize()
+}
+
+func DownloadMaxMindDatabase(dbFilePath string, downloadURL string) error {
+	if dbFilePath == "" {
+		dbFilePath = GeoIpFilePath
+	}
+	if downloadURL == "" {
+		downloadURL = GeoIpUrl
+	}
+	resp, err := http.Get(downloadURL)
+	if err != nil {
+		return fmt.Errorf("failed to initiate MaxMind database download: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return fmt.Errorf("failed to download MaxMind database: HTTP status %s", resp.Status)
+	}
+
+	if err := os.MkdirAll(filepath.Dir(dbFilePath), os.ModePerm); err != nil {
+		return fmt.Errorf("failed to create data directory for MaxMind database update: %w", err)
+	}
+
+	tempPath := dbFilePath + ".download"
+	out, err := os.Create(tempPath)
+	if err != nil {
+		return fmt.Errorf("failed to create MaxMind database file at %s: %w", tempPath, err)
+	}
+	defer func() {
+		_ = out.Close()
+	}()
+
+	if _, err = io.Copy(out, resp.Body); err != nil {
+		return fmt.Errorf("failed to write MaxMind database file: %w", err)
+	}
+	if err = out.Close(); err != nil {
+		return fmt.Errorf("failed to close MaxMind database file: %w", err)
+	}
+	if err = os.Rename(tempPath, dbFilePath); err != nil {
+		return fmt.Errorf("failed to move MaxMind database file into place: %w", err)
+	}
+	return nil
+}
+
+func (s *MaxMindGeoIPService) Close() error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.maxMindDBReader != nil {
+		err := s.maxMindDBReader.Close()
+		s.maxMindDBReader = nil
+		if err != nil {
+			return fmt.Errorf("error closing MaxMind database: %w", err)
+		}
+	}
+	return nil
+}