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.

internal/aggregate/geo_enrich.go

@@ -0,0 +1,41 @@
+package aggregate
+
+import (
+	"github.com/telemt/telemt-api/internal/geoip"
+)
+
+// EnrichUniqueIPsGeo fills country/city on each IP when lookup succeeds.
+func EnrichUniqueIPsGeo(rows []UniqueIPsRow, g *geoip.Service) {
+	if g == nil || len(rows) == 0 {
+		return
+	}
+	for i := range rows {
+		for j := range rows[i].IPs {
+			ip := rows[i].IPs[j].IP
+			res, ok := g.Lookup(ip)
+			if !ok {
+				continue
+			}
+			if res.CountryCode != "" {
+				s := res.CountryCode
+				rows[i].IPs[j].CountryCode = &s
+			}
+			if res.CountryName != "" {
+				s := res.CountryName
+				rows[i].IPs[j].CountryName = &s
+			}
+			if res.CityName != "" {
+				s := res.CityName
+				rows[i].IPs[j].CityName = &s
+			}
+			if res.ASN != 0 {
+				a := res.ASN
+				rows[i].IPs[j].ASN = &a
+			}
+			if res.ASOrg != "" {
+				o := res.ASOrg
+				rows[i].IPs[j].ASOrganization = &o
+			}
+		}
+	}
+}

internal/aggregate/handlers.go

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/telemt/telemt-api/internal/config"
+	"github.com/telemt/telemt-api/internal/geoip"
 )
 
 const pathPrefix = "/api/agg"
@@ -17,11 +18,13 @@ const pathPrefix = "/api/agg"
 type Handler struct {
 	Parsed *config.Parsed
 	Client *http.Client
+	Geo    *geoip.Service
 }
 
 // NewHandler builds an aggregate handler; client must use a non-nil Transport (e.g. gateway shared transport).
-func NewHandler(p *config.Parsed, client *http.Client) *Handler {
-	return &Handler{Parsed: p, Client: client}
+// Geo may be nil (no GeoLite2 lookups).
+func NewHandler(p *config.Parsed, client *http.Client, geo *geoip.Service) *Handler {
+	return &Handler{Parsed: p, Client: client, Geo: geo}
 }
 
 func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
@@ -154,6 +157,9 @@ func (h *Handler) handleUniqueIPs(w http.ResponseWriter, r *http.Request) {
 	defer cancel()
 	results := FetchStatsUsers(ctx, h.Client, h.Parsed, aliases)
 	data := BuildUniqueIPs(results)
+	if h.Geo != nil && !strings.EqualFold(r.URL.Query().Get("geo"), "false") {
+		EnrichUniqueIPsGeo(data, h.Geo)
+	}
 	writeOK(w, data)
 }
 

internal/aggregate/handlers_test.go

@@ -35,7 +35,7 @@ func TestHandlerResolveAndFetch(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	h := NewHandler(parsed, up.Client())
+	h := NewHandler(parsed, up.Client(), nil)
 	req := httptest.NewRequest(http.MethodGet, "/api/agg/summary?aliases=test", nil)
 	rec := httptest.NewRecorder()
 	h.ServeHTTP(rec, req)
@@ -65,7 +65,7 @@ func TestHandlerMethodNotAllowed(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	h := NewHandler(parsed, http.DefaultClient)
+	h := NewHandler(parsed, http.DefaultClient, nil)
 	req := httptest.NewRequest(http.MethodPost, "/api/agg/summary", nil)
 	rec := httptest.NewRecorder()
 	h.ServeHTTP(rec, req)

internal/aggregate/types.go

@@ -62,10 +62,16 @@ type UniqueIPsRow struct {
 
 // IPAssignments per-IP server visibility (snapshot-based).
 type IPAssignments struct {
-	IP                string   `json:"ip"`
-	ActiveOnServers   []string `json:"active_on_servers"`
-	RecentOnServers   []string `json:"recent_on_servers"`
-	PrimaryServer     *string  `json:"primary_server,omitempty"`
+	IP              string   `json:"ip"`
+	ActiveOnServers []string `json:"active_on_servers"`
+	RecentOnServers []string `json:"recent_on_servers"`
+	PrimaryServer   *string  `json:"primary_server,omitempty"`
+	// GeoIP when gateway geoip.enabled and lookup succeeds:
+	CountryCode *string `json:"country_code,omitempty"`
+	CountryName *string `json:"country_name,omitempty"`
+	CityName    *string `json:"city_name,omitempty"`
+	ASN         *uint64 `json:"asn,omitempty"`                 // GeoLite2-ASN
+	ASOrganization *string `json:"as_organization,omitempty"` // ASN org name
 }
 
 // UsersRow merged user view with optional per-server detail and links.

internal/config/config.go

@@ -21,6 +21,17 @@ type Config struct {
 	TrustedProxies []string         `yaml:"trusted_proxies"`
 	Servers        []Server         `yaml:"servers"`
 	Aggregate      *AggregateConfig `yaml:"aggregate"`
+	GeoIP          *GeoIPConfig     `yaml:"geoip"`
+}
+
+// GeoIPConfig enables GeoLite2 lookups for /api/agg/unique-ips (optional).
+// Нужен хотя бы один из: database_path (City) или asn_database_path (ASN). Country-only DB не используется — страна в City.
+type GeoIPConfig struct {
+	Enabled          bool   `yaml:"enabled"`
+	DatabasePath     string `yaml:"database_path"`      // GeoLite2-City.mmdb
+	DownloadURL      string `yaml:"download_url"`       // если city-файла нет; .mmdb или .gz
+	AsnDatabasePath  string `yaml:"asn_database_path"`  // GeoLite2-ASN.mmdb
+	AsnDownloadURL   string `yaml:"asn_download_url"`   // если asn-файла нет
 }
 
 // AggregateConfig controls default scope of /api/agg/* (optional).

internal/geoip/config_open.go

@@ -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)
+}

internal/geoip/download.go

@@ -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
+}

internal/geoip/geoip.go

@@ -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
+}

internal/server/gateway.go

@@ -15,6 +15,7 @@ import (
 
 	"github.com/telemt/telemt-api/internal/aggregate"
 	"github.com/telemt/telemt-api/internal/config"
+	"github.com/telemt/telemt-api/internal/geoip"
 	"github.com/telemt/telemt-api/internal/proxy"
 )
 
@@ -23,13 +24,14 @@ type Gateway struct {
 	parsed        *config.Parsed
 	proxies       map[string]*httputil.ReverseProxy
 	agg           *aggregate.Handler
+	geo           *geoip.Service
 	log           *slog.Logger
 	transport     *http.Transport
 	promHandler   http.Handler
 }
 
 // NewGateway builds handlers and reverse proxies from parsed config.
-func NewGateway(p *config.Parsed, log *slog.Logger) (*Gateway, error) {
+func NewGateway(p *config.Parsed, log *slog.Logger, geo *geoip.Service) (*Gateway, error) {
 	t := proxy.DirectTransport()
 	t.MaxIdleConns = 64
 	t.IdleConnTimeout = 90 * time.Second
@@ -40,6 +42,7 @@ func NewGateway(p *config.Parsed, log *slog.Logger) (*Gateway, error) {
 	g := &Gateway{
 		parsed:      p,
 		proxies:     make(map[string]*httputil.ReverseProxy),
+		geo:         geo,
 		log:         log,
 		transport:   t,
 		promHandler: promhttp.Handler(),
@@ -65,7 +68,7 @@ func NewGateway(p *config.Parsed, log *slog.Logger) (*Gateway, error) {
 		}
 		g.proxies[s.Alias] = rp
 	}
-	g.agg = aggregate.NewHandler(p, &http.Client{Transport: t})
+	g.agg = aggregate.NewHandler(p, &http.Client{Transport: t}, geo)
 	return g, nil
 }
 
@@ -222,5 +225,8 @@ func (s *statusWriter) WriteHeader(code int) {
 // Shutdown idle connections on the shared transport.
 func (g *Gateway) Shutdown(ctx context.Context) error {
 	g.transport.CloseIdleConnections()
+	if g.geo != nil {
+		_ = g.geo.Close()
+	}
 	return nil
 }