160 lines
4.9 KiB
Go
160 lines
4.9 KiB
Go
// Package geo wraps the GeoLite2 country resolver and exposes the
|
|
// platform-level geo helpers consumed by `backend/internal/auth` at user
|
|
// registration time and by the user-surface middleware on every
|
|
// authenticated request.
|
|
//
|
|
// The implementation shipped `LookupCountry`, `LanguageForIP` and
|
|
// `SetDeclaredCountryAtRegistration`. The implementation added the
|
|
// `OnUserDeleted` cascade leg. The implementation layers `IncrementCounterAsync`
|
|
// and `ListUserCounters` on top of the same Service plus the
|
|
// background-goroutine machinery (cancellable context and WaitGroup)
|
|
// needed to drain pending counter upserts on shutdown.
|
|
package geo
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"errors"
|
|
"fmt"
|
|
"sync"
|
|
"sync/atomic"
|
|
|
|
"galaxy/geoip"
|
|
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
// Service is the geo-domain entry point. It is safe for concurrent use.
|
|
type Service struct {
|
|
db *sql.DB
|
|
resolver *geoip.Resolver
|
|
|
|
logger *zap.Logger
|
|
|
|
// bgCtx is the lifetime context passed to fire-and-forget goroutines
|
|
// launched by IncrementCounterAsync. It is cancelled by Close so that
|
|
// in-flight counter upserts observe shutdown promptly. The matching
|
|
// WaitGroup tracks live goroutines so Drain (and Close) can wait for
|
|
// them.
|
|
bgCtx context.Context
|
|
bgCancel context.CancelFunc
|
|
wg sync.WaitGroup
|
|
closed atomic.Bool
|
|
}
|
|
|
|
// NewService constructs a Service backed by the GeoLite2 country database
|
|
// at databasePath and the supplied Postgres pool. Closing the returned
|
|
// Service releases the memory-mapped database file; the database pool is
|
|
// owned by the caller.
|
|
//
|
|
// A trimmed-empty databasePath is rejected with a non-nil error so that
|
|
// boot fails fast rather than silently hiding lookups behind a permanent
|
|
// failure path. Callers that explicitly want a no-op Service should
|
|
// inject their own implementation via the auth-level interfaces.
|
|
//
|
|
// The returned Service uses a no-op zap logger by default; callers that
|
|
// want diagnostic output from the asynchronous counter path inject one
|
|
// via SetLogger.
|
|
func NewService(databasePath string, db *sql.DB) (*Service, error) {
|
|
if db == nil {
|
|
return nil, errors.New("geo: db must not be nil")
|
|
}
|
|
resolver, err := geoip.Open(databasePath)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("geo: open resolver: %w", err)
|
|
}
|
|
bgCtx, bgCancel := context.WithCancel(context.Background())
|
|
return &Service{
|
|
db: db,
|
|
resolver: resolver,
|
|
logger: zap.NewNop(),
|
|
bgCtx: bgCtx,
|
|
bgCancel: bgCancel,
|
|
}, nil
|
|
}
|
|
|
|
// SetLogger replaces the diagnostic logger used by the asynchronous
|
|
// counter path. A nil argument resets the logger to a no-op so that
|
|
// production wiring can supply a real logger after construction without
|
|
// the test paths having to thread one through. SetLogger is nil-safe on
|
|
// the Service receiver.
|
|
func (s *Service) SetLogger(logger *zap.Logger) {
|
|
if s == nil {
|
|
return
|
|
}
|
|
if logger == nil {
|
|
logger = zap.NewNop()
|
|
}
|
|
s.logger = logger.Named("geo")
|
|
}
|
|
|
|
// Drain blocks until every fire-and-forget goroutine launched through
|
|
// IncrementCounterAsync has finished, or until ctx is done. It cancels
|
|
// the Service-internal background context so live goroutines observe
|
|
// shutdown and stop waiting on the database. Drain is nil-safe and
|
|
// idempotent: subsequent calls return immediately.
|
|
//
|
|
// Drain does not close the GeoLite2 resolver — Close does. The split
|
|
// lets the boot orchestrator wait for in-flight writes within the
|
|
// shutdown deadline before the resolver and database pool are torn
|
|
// down.
|
|
func (s *Service) Drain(ctx context.Context) {
|
|
if s == nil {
|
|
return
|
|
}
|
|
if s.bgCancel != nil {
|
|
s.bgCancel()
|
|
}
|
|
done := make(chan struct{})
|
|
go func() {
|
|
s.wg.Wait()
|
|
close(done)
|
|
}()
|
|
select {
|
|
case <-done:
|
|
case <-ctx.Done():
|
|
}
|
|
}
|
|
|
|
// Close releases the underlying GeoLite2 database resources. Pending
|
|
// counter goroutines launched through IncrementCounterAsync are
|
|
// signalled to stop via the internal background context but are NOT
|
|
// awaited; callers that need to wait must invoke Drain first. Close is
|
|
// idempotent and nil-safe; subsequent lookups return the empty country
|
|
// / language ("" treated as no data).
|
|
func (s *Service) Close() error {
|
|
if s == nil {
|
|
return nil
|
|
}
|
|
if !s.closed.CompareAndSwap(false, true) {
|
|
return nil
|
|
}
|
|
if s.bgCancel != nil {
|
|
s.bgCancel()
|
|
}
|
|
if s.resolver == nil {
|
|
return nil
|
|
}
|
|
if err := s.resolver.Close(); err != nil {
|
|
return fmt.Errorf("geo: close resolver: %w", err)
|
|
}
|
|
s.resolver = nil
|
|
return nil
|
|
}
|
|
|
|
// LookupCountry resolves an uppercase ISO 3166-1 alpha-2 country code
|
|
// from sourceIP. The lookup is best-effort: the empty string is returned
|
|
// for any invalid address, missing record, or closed resolver. The
|
|
// returned error is always nil; callers that need diagnostic detail
|
|
// should query the geoip resolver directly.
|
|
func (s *Service) LookupCountry(sourceIP string) string {
|
|
if s == nil || s.resolver == nil || sourceIP == "" {
|
|
return ""
|
|
}
|
|
code, err := s.resolver.CountryString(sourceIP)
|
|
if err != nil {
|
|
return ""
|
|
}
|
|
return code
|
|
}
|