From d7a63f0da39fdc7e624d18a8ae54338b1f90ba61 Mon Sep 17 00:00:00 2001 From: Denozordec Date: Mon, 30 Mar 2026 01:47:08 +0700 Subject: [PATCH] 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. --- README.md | 1 + cmd/gateway/main.go | 17 +++- config.example.yaml | 8 ++ docs/AGGREGATE.md | 2 +- docs/GEOIP.md | 42 ++++++++++ go.mod | 1 + internal/aggregate/geo_enrich.go | 41 ++++++++++ internal/aggregate/handlers.go | 10 ++- internal/aggregate/handlers_test.go | 4 +- internal/aggregate/types.go | 14 +++- internal/config/config.go | 11 +++ internal/geoip/config_open.go | 70 +++++++++++++++++ internal/geoip/download.go | 113 +++++++++++++++++++++++++++ internal/geoip/geoip.go | 117 ++++++++++++++++++++++++++++ internal/server/gateway.go | 10 ++- 15 files changed, 449 insertions(+), 12 deletions(-) create mode 100644 docs/GEOIP.md create mode 100644 internal/aggregate/geo_enrich.go create mode 100644 internal/geoip/config_open.go create mode 100644 internal/geoip/download.go create mode 100644 internal/geoip/geoip.go diff --git a/README.md b/README.md index 518bda8..e1d8026 100644 --- a/README.md +++ b/README.md @@ -56,6 +56,7 @@ docker compose logs -f gateway | **[docs/GATEWAY_RUN.md](docs/GATEWAY_RUN.md)** | Полная инструкция: конфиг, pull/registry, Docker CLI, Compose, CI/CD, неполадки | | **[docs/API.md](docs/API.md)** | Контракт Telemt Control API (`/v1/…`) | | **[docs/AGGREGATE.md](docs/AGGREGATE.md)** | Агрегирующие эндпоинты шлюза (`/api/agg/…`) | +| **[docs/GEOIP.md](docs/GEOIP.md)** | GeoLite2 City (страна/город) и опционально ASN (номер AS, организация) для IP в `unique-ips` | ## Сборка и тесты без Docker diff --git a/cmd/gateway/main.go b/cmd/gateway/main.go index 6cc5886..d945a37 100644 --- a/cmd/gateway/main.go +++ b/cmd/gateway/main.go @@ -3,6 +3,7 @@ package main import ( "context" "log/slog" + "net" "net/http" "os" "os/signal" @@ -10,6 +11,7 @@ import ( "time" "github.com/telemt/telemt-api/internal/config" + "github.com/telemt/telemt-api/internal/geoip" "github.com/telemt/telemt-api/internal/server" ) @@ -31,7 +33,20 @@ func main() { os.Exit(1) } - gw, err := server.NewGateway(parsed, log) + geoDown := &http.Client{ + Transport: &http.Transport{ + DialContext: (&net.Dialer{Timeout: 30 * time.Second}).DialContext, + ResponseHeaderTimeout: 5 * time.Minute, + }, + Timeout: 15 * time.Minute, + } + geo, err := geoip.FromConfig(parsed.Config.GeoIP, geoDown, log) + if err != nil { + log.Error("geoip init failed", "err", err) + os.Exit(1) + } + + gw, err := server.NewGateway(parsed, log, geo) if err != nil { log.Error("gateway init failed", "err", err) os.Exit(1) diff --git a/config.example.yaml b/config.example.yaml index 1ea7bd1..ac071a0 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -29,6 +29,14 @@ trusted_proxies: [] # - gt1 # - gt2 +# Геолокация IP в /api/agg/unique-ips. См. docs/GEOIP.md (City + опционально ASN; Country-only не нужен) +# geoip: +# enabled: true +# database_path: /var/lib/telemt-gateway/GeoLite2-City.mmdb +# download_url: "https://github.com/P3TERX/GeoLite.mmdb/raw/download/GeoLite2-City.mmdb" +# asn_database_path: /var/lib/telemt-gateway/GeoLite2-ASN.mmdb +# asn_download_url: "https://github.com/P3TERX/GeoLite.mmdb/raw/download/GeoLite2-ASN.mmdb" + servers: - alias: main_srv base_url: http://127.0.0.1:9091 diff --git a/docs/AGGREGATE.md b/docs/AGGREGATE.md index 64be5bd..071816d 100644 --- a/docs/AGGREGATE.md +++ b/docs/AGGREGATE.md @@ -12,7 +12,7 @@ | --- | --- | --- | | GET | `/api/agg/summary` | Сводка по флоту, список опросов upstream, `fleet_total_megabytes` / `fleet_total_connections`. Два топа (размер задаётся `top_n`): **`top_users`** — самые «прожорливые» по суммарному трафику (MiB) по всем серверам; **`top_users_by_unique_ips`** — по максимальному `active_unique_ips` среди серверов для пользователя (как в Telemt, снимок). | | GET | `/api/agg/traffic` | Трафик по каждому пользователю в разрезе серверов: `servers..total_megabytes`. | -| GET | `/api/agg/unique-ips` | Уникальные IP по пользователю: на каких серверах IP есть в active/recent списках снимка. | +| GET | `/api/agg/unique-ips` | Уникальные IP по пользователю: на каких серверах IP есть в active/recent списках снимка. При **`geoip.enabled`** в конфиге — из City: `country_code`, `country_name`, `city_name`; при наличии ASN-БД — `asn`, `as_organization` (см. [GEOIP.md](GEOIP.md)); отключить гео для запроса: `?geo=false`. | | GET | `/api/agg/users` | Объединённый список пользователей с `by_server` и суммарным `total_megabytes`. | Все методы — **GET**; действует тот же whitelist, что и для остального API шлюза. diff --git a/docs/GEOIP.md b/docs/GEOIP.md new file mode 100644 index 0000000..d9655b6 --- /dev/null +++ b/docs/GEOIP.md @@ -0,0 +1,42 @@ +# GeoIP / GeoLite2 в шлюзе + +При включённом **`geoip`** шлюз открывает один или два файла **MaxMind DB** (`.mmdb`) и обогащает ответ **`GET /api/agg/unique-ips`** (если не передан `?geo=false`). + +## Нужны ли все три базы (City, Country, ASN)? + +| База | Нужна ли отдельно | Зачем | +| --- | --- | --- | +| **GeoLite2-City** | Для геолокации — **да** (или только она) | Страна, город, англ. названия. | +| **GeoLite2-Country** | **Нет**, если есть City | В City уже есть страна; Country даёт только страну и **дублирует** часть данных City. | +| **GeoLite2-ASN** | **Опционально** | Номер автономной системы (ASN) и имя организации (провайдер/сеть), **не** заменяет город. | + +Итого: минимум — **одна база City** для страны/города. Для «кто оператор сети» добавляют **ASN** (City + ASN вместе дают полную картину). **Three MMDB не обязательны** — достаточно **1–2 файлов** (City + при желании ASN). + +## Поля в JSON + +- Из **City**: `country_code`, `country_name`, `city_name` +- Из **ASN**: `asn`, `as_organization` + +## Конфигурация + +Нужен **хотя бы один** путь: `database_path` (City) и/или `asn_database_path` (ASN). + +Пример с репозитория [P3TERX/GeoLite.mmdb](https://github.com/P3TERX/GeoLite.mmdb) (сырые `.mmdb` по HTTPS): + +```yaml +geoip: + enabled: true + database_path: /var/lib/telemt-gateway/GeoLite2-City.mmdb + download_url: "https://github.com/P3TERX/GeoLite.mmdb/raw/download/GeoLite2-City.mmdb" + asn_database_path: /var/lib/telemt-gateway/GeoLite2-ASN.mmdb + asn_download_url: "https://github.com/P3TERX/GeoLite.mmdb/raw/download/GeoLite2-ASN.mmdb" +``` + +- Если локального файла нет, а указан соответствующий `*_download_url`, при старте выполняется загрузка. +- URL с суффиксом **`.gz`** обрабатываются как gzip; иначе тело ответа пишется в файл как есть (подходит для GitHub raw `.mmdb`). + +## Лицензия + +GeoLite2 распространяется MaxMind на условиях [лицензии](https://dev.maxmind.com/geoip/geolite2-free-data-feed). Сторонние зеркала и репозитории используйте на свой риск; для продакшена часто нужен свой ключ MaxMind и обновление БД. + +Библиотека чтения: [`github.com/oschwald/geoip2-golang`](https://github.com/oschwald/geoip2-golang). diff --git a/go.mod b/go.mod index 06d772a..727332b 100644 --- a/go.mod +++ b/go.mod @@ -3,6 +3,7 @@ module github.com/telemt/telemt-api go 1.22 require ( + github.com/oschwald/geoip2-golang v1.11.0 github.com/prometheus/client_golang v1.20.5 gopkg.in/yaml.v3 v3.0.1 ) diff --git a/internal/aggregate/geo_enrich.go b/internal/aggregate/geo_enrich.go new file mode 100644 index 0000000..e82b3ce --- /dev/null +++ b/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 + } + } + } +} diff --git a/internal/aggregate/handlers.go b/internal/aggregate/handlers.go index a4d2f20..7267e97 100644 --- a/internal/aggregate/handlers.go +++ b/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) } diff --git a/internal/aggregate/handlers_test.go b/internal/aggregate/handlers_test.go index c7bd625..27743f2 100644 --- a/internal/aggregate/handlers_test.go +++ b/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) diff --git a/internal/aggregate/types.go b/internal/aggregate/types.go index 5a9ee31..cb19fd4 100644 --- a/internal/aggregate/types.go +++ b/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. diff --git a/internal/config/config.go b/internal/config/config.go index 5452e33..4d83614 100644 --- a/internal/config/config.go +++ b/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). diff --git a/internal/geoip/config_open.go b/internal/geoip/config_open.go new file mode 100644 index 0000000..dd842e9 --- /dev/null +++ b/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) +} diff --git a/internal/geoip/download.go b/internal/geoip/download.go new file mode 100644 index 0000000..cf62260 --- /dev/null +++ b/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 +} diff --git a/internal/geoip/geoip.go b/internal/geoip/geoip.go new file mode 100644 index 0000000..cc6d33e --- /dev/null +++ b/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 +} diff --git a/internal/server/gateway.go b/internal/server/gateway.go index 2e8d8b2..cab7919 100644 --- a/internal/server/gateway.go +++ b/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 }