Add GeoIP support for unique IP aggregation

- Introduced GeoIP configuration options in config.example.yaml to enable geolocation lookups for the /api/agg/unique-ips endpoint. - Updated the aggregate handler to include optional GeoIP data in responses, enriching unique IP information with country and city details, as well as ASN data if available. - Enhanced documentation in AGGREGATE.md and README.md to reflect the new GeoIP functionality and its usage. - Added a dependency on the geoip2-golang library in go.mod for GeoIP lookups. - Modified tests to accommodate the new GeoIP integration in the aggregate handler.
2026-03-30 01:47:08 +07:00
parent 4a4f15bd0b
commit d7a63f0da3
15 changed files with 449 additions and 12 deletions
@@ -0,0 +1,70 @@
+package geoip
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"log/slog"
+	"net/http"
+	"os"
+	"strings"
+	"time"
+
+	"github.com/telemt/telemt-api/internal/config"
+)
+
+// FromConfig ensures MMDB files exist (optional download) and opens readers. Returns nil if geoip disabled.
+func FromConfig(cfg *config.GeoIPConfig, httpClient *http.Client, log *slog.Logger) (*Service, error) {
+	if cfg == nil || !cfg.Enabled {
+		return nil, nil
+	}
+	cityPath := strings.TrimSpace(cfg.DatabasePath)
+	asnPath := strings.TrimSpace(cfg.AsnDatabasePath)
+	if cityPath == "" && asnPath == "" {
+		return nil, errors.New("geoip.enabled requires at least one of geoip.database_path or geoip.asn_database_path")
+	}
+
+	if httpClient == nil {
+		httpClient = http.DefaultClient
+	}
+
+	if cityPath != "" {
+		if _, err := os.Stat(cityPath); err != nil {
+			if !os.IsNotExist(err) {
+				return nil, err
+			}
+			if cfg.DownloadURL == "" {
+				return nil, fmt.Errorf("geoip city database %q missing and geoip.download_url is empty", cityPath)
+			}
+			if log != nil {
+				log.Info("geoip downloading city database", "url", cfg.DownloadURL, "path", cityPath)
+			}
+			ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
+			defer cancel()
+			if err := DownloadMMDB(ctx, httpClient, cfg.DownloadURL, cityPath); err != nil {
+				return nil, err
+			}
+		}
+	}
+
+	if asnPath != "" {
+		if _, err := os.Stat(asnPath); err != nil {
+			if !os.IsNotExist(err) {
+				return nil, err
+			}
+			if cfg.AsnDownloadURL == "" {
+				return nil, fmt.Errorf("geoip asn database %q missing and geoip.asn_download_url is empty", asnPath)
+			}
+			if log != nil {
+				log.Info("geoip downloading asn database", "url", cfg.AsnDownloadURL, "path", asnPath)
+			}
+			ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
+			defer cancel()
+			if err := DownloadMMDB(ctx, httpClient, cfg.AsnDownloadURL, asnPath); err != nil {
+				return nil, err
+			}
+		}
+	}
+
+	return OpenDatabases(cityPath, asnPath)
+}
@@ -0,0 +1,113 @@
+package geoip
+
+import (
+	"compress/gzip"
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"path/filepath"
+	"strings"
+)
+
+// DownloadMMDB downloads an MMDB URL to destPath. If the URL ends with `.gz`, the body is gunzipped.
+// Otherwise the response body is written as-is (raw `.mmdb`, e.g. GitHub raw).
+func DownloadMMDB(ctx context.Context, client *http.Client, url, destPath string) error {
+	if strings.HasSuffix(strings.ToLower(strings.TrimSpace(url)), ".gz") {
+		return downloadGzippedMMDB(ctx, client, url, destPath)
+	}
+	return downloadRawMMDB(ctx, client, url, destPath)
+}
+
+func downloadRawMMDB(ctx context.Context, client *http.Client, url, destPath string) error {
+	if client == nil {
+		client = http.DefaultClient
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		return err
+	}
+	resp, err := client.Do(req)
+	if err != nil {
+		return fmt.Errorf("geoip download: %w", err)
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		return fmt.Errorf("geoip download: http %s", resp.Status)
+	}
+	dir := filepath.Dir(destPath)
+	if err := os.MkdirAll(dir, 0o755); err != nil {
+		return fmt.Errorf("geoip mkdir: %w", err)
+	}
+	tmp, err := os.CreateTemp(dir, "geolite-*.mmdb.part")
+	if err != nil {
+		return err
+	}
+	tmpPath := tmp.Name()
+	_, copyErr := io.Copy(tmp, resp.Body)
+	closeErr := tmp.Close()
+	if copyErr != nil {
+		_ = os.Remove(tmpPath)
+		return copyErr
+	}
+	if closeErr != nil {
+		_ = os.Remove(tmpPath)
+		return closeErr
+	}
+	_ = os.Remove(destPath)
+	if err := os.Rename(tmpPath, destPath); err != nil {
+		_ = os.Remove(tmpPath)
+		return fmt.Errorf("geoip rename: %w", err)
+	}
+	return nil
+}
+
+func downloadGzippedMMDB(ctx context.Context, client *http.Client, url, destPath string) error {
+	if client == nil {
+		client = http.DefaultClient
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		return err
+	}
+	resp, err := client.Do(req)
+	if err != nil {
+		return fmt.Errorf("geoip download: %w", err)
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		return fmt.Errorf("geoip download: http %s", resp.Status)
+	}
+	gzr, err := gzip.NewReader(resp.Body)
+	if err != nil {
+		return fmt.Errorf("geoip gzip: %w", err)
+	}
+	defer gzr.Close()
+
+	dir := filepath.Dir(destPath)
+	if err := os.MkdirAll(dir, 0o755); err != nil {
+		return fmt.Errorf("geoip mkdir: %w", err)
+	}
+	tmp, err := os.CreateTemp(dir, "geolite-*.mmdb.part")
+	if err != nil {
+		return err
+	}
+	tmpPath := tmp.Name()
+	_, copyErr := io.Copy(tmp, gzr)
+	closeErr := tmp.Close()
+	if copyErr != nil {
+		_ = os.Remove(tmpPath)
+		return copyErr
+	}
+	if closeErr != nil {
+		_ = os.Remove(tmpPath)
+		return closeErr
+	}
+	_ = os.Remove(destPath)
+	if err := os.Rename(tmpPath, destPath); err != nil {
+		_ = os.Remove(tmpPath)
+		return fmt.Errorf("geoip rename: %w", err)
+	}
+	return nil
+}
@@ -0,0 +1,117 @@
+package geoip
+
+import (
+	"fmt"
+	"net"
+
+	"github.com/oschwald/geoip2-golang/geoip2"
+)
+
+// Service holds optional GeoLite2 City and/or ASN MMDB readers.
+// Country-only DB не нужен: страна уже есть в City.
+type Service struct {
+	city *geoip2.Reader
+	asn  *geoip2.Reader
+}
+
+// Open opens a single MMDB (City or ASN) — для тестов и простых случаев.
+func Open(path string) (*Service, error) {
+	r, err := geoip2.Open(path)
+	if err != nil {
+		return nil, fmt.Errorf("geoip open %q: %w", path, err)
+	}
+	return &Service{city: r}, nil
+}
+
+// OpenDatabases opens city and/or asn paths; at least one must be non-empty.
+func OpenDatabases(cityPath, asnPath string) (*Service, error) {
+	var s Service
+	if cityPath != "" {
+		r, err := geoip2.Open(cityPath)
+		if err != nil {
+			return nil, fmt.Errorf("geoip city open %q: %w", cityPath, err)
+		}
+		s.city = r
+	}
+	if asnPath != "" {
+		r, err := geoip2.Open(asnPath)
+		if err != nil {
+			if s.city != nil {
+				_ = s.city.Close()
+			}
+			return nil, fmt.Errorf("geoip asn open %q: %w", asnPath, err)
+		}
+		s.asn = r
+	}
+	if s.city == nil && s.asn == nil {
+		return nil, fmt.Errorf("geoip: no database paths")
+	}
+	return &s, nil
+}
+
+// Close releases database handles.
+func (s *Service) Close() error {
+	if s == nil {
+		return nil
+	}
+	var first error
+	if s.city != nil {
+		if err := s.city.Close(); err != nil {
+			first = err
+		}
+		s.city = nil
+	}
+	if s.asn != nil {
+		if err := s.asn.Close(); err != nil && first == nil {
+			first = err
+		}
+		s.asn = nil
+	}
+	return first
+}
+
+// Result combines City + ASN fields when available.
+type Result struct {
+	CountryCode string
+	CountryName string
+	CityName    string
+	ASN         uint64
+	ASOrg       string
+}
+
+// Lookup fills geographic and/or ASN data for an IP string (IPv4/IPv6).
+func (s *Service) Lookup(ipStr string) (Result, bool) {
+	var out Result
+	if s == nil {
+		return out, false
+	}
+	ip := net.ParseIP(ipStr)
+	if ip == nil {
+		return out, false
+	}
+	ok := false
+	if s.city != nil {
+		if rec, err := s.city.City(ip); err == nil {
+			out.CountryCode = rec.Country.IsoCode
+			if rec.Country.Names != nil {
+				out.CountryName = rec.Country.Names["en"]
+			}
+			if rec.City.Names != nil {
+				out.CityName = rec.City.Names["en"]
+			}
+			if out.CountryCode != "" || out.CountryName != "" || out.CityName != "" {
+				ok = true
+			}
+		}
+	}
+	if s.asn != nil {
+		if rec, err := s.asn.ASN(ip); err == nil {
+			out.ASN = uint64(rec.AutonomousSystemNumber)
+			out.ASOrg = rec.AutonomousSystemOrganization
+			if out.ASN != 0 || out.ASOrg != "" {
+				ok = true
+			}
+		}
+	}
+	return out, ok
+}