ipinfo/internal/db/manager.go

package db

import (
	"fmt"
	"log/slog"
	"net"
	"net/http"
	"sync"
	"time"

	"github.com/oschwald/maxminddb-golang"
)

// GeoIPManager manages the GeoIP databases
type GeoIPManager struct {
	cityDB       *maxminddb.Reader
	asnDB        *maxminddb.Reader
	asnPrefixMap map[uint][]*net.IPNet
	httpClient   *http.Client
	mu           sync.RWMutex
}

// NewGeoIPManager creates a new GeoIPManager
func NewGeoIPManager() (*GeoIPManager, error) {
	manager := &GeoIPManager{
		httpClient: &http.Client{Timeout: 5 * time.Minute},
	}
	if err := manager.Initialize(); err != nil {
		return nil, fmt.Errorf("initializing geoip manager: %w", err)
	}
	return manager, nil
}

// Initialize initializes the GeoIPManager by opening the database files.
func (g *GeoIPManager) Initialize() error {
	g.mu.Lock()
	cityErr := g.openDB(CityDBPath)
	asnErr := g.openDB(ASNDBPath)
	g.mu.Unlock()

	if cityErr != nil || asnErr != nil {
		slog.Info("databases missing or invalid, performing initial update")
		if err := g.UpdateDatabases(); err != nil {
			return fmt.Errorf("initial update failed: %w", err)
		}
	} else {
		g.mu.Lock()
		g.buildASNPrefixMap()
		g.mu.Unlock()
	}

	return nil
}

// Close closes the GeoIP database readers
func (g *GeoIPManager) Close() {
	g.mu.Lock()
	defer g.mu.Unlock()
	if g.cityDB != nil {
		if err := g.cityDB.Close(); err != nil {
			slog.Warn("failed to close citydb", "err", err)
		}
	}
	if g.asnDB != nil {
		if err := g.asnDB.Close(); err != nil {
			slog.Warn("failed to close asndb", "err", err)
		}
	}
}

// openDB opens a MaxMind DB file.
func (g *GeoIPManager) openDB(path string) error {
	db, err := maxminddb.Open(path)
	if err != nil {
		return err
	}

	if path == CityDBPath {
		g.cityDB = db
	} else {
		g.asnDB = db
	}
	return nil
}

// buildASNPrefixMap builds a map of ASN prefixes for fast lookups.
func (g *GeoIPManager) buildASNPrefixMap() {
	slog.Info("building asn prefix map for fast lookups")
	startTime := time.Now()
	g.asnPrefixMap = make(map[uint][]*net.IPNet)
	if g.asnDB == nil {
		slog.Warn("asn database is not available, skipping prefix map build")
		return
	}
	networks := g.asnDB.Networks()
	for networks.Next() {
		var record ASNRecord
		subnet, err := networks.Network(&record)
		if err != nil {
			slog.Debug("skipping asn network due to error", "err", err)
			continue
		}
		if record.AutonomousSystemNumber > 0 {
			g.asnPrefixMap[record.AutonomousSystemNumber] = append(g.asnPrefixMap[record.AutonomousSystemNumber], subnet)
		}
	}
	slog.Info("finished building asn prefix map", "duration", time.Since(startTime))
}