ipinfo/internal/db/db.go

package internal

import (
	"context"
	"errors"
	"fmt"
	"log"
	"net"
	"os"
	"os/exec"
	"sync"
	"time"

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

// Database paths
const (
	CityDBPath = "./GeoLite2-City.mmdb"
	ASNDBPath  = "./GeoLite2-ASN.mmdb"
)

// Common errors
var (
	ErrDatabaseNotFound = errors.New("database file not found")
	ErrDatabaseOpen     = errors.New("failed to open database")
	ErrDownloadFailed   = errors.New("failed to download database")
)

type ASNRecord struct {
	AutonomousSystemNumber       uint   `maxminddb:"autonomous_system_number"`
	AutonomousSystemOrganization string `maxminddb:"autonomous_system_organization"`
}

// Handles MaxMind GeoIP database operations
type GeoIPManager struct {
	cityDB       *maxminddb.Reader
	asnDB        *maxminddb.Reader
	asnPrefixMap map[uint][]*net.IPNet
	mu           sync.RWMutex
}

// Creates and initializes a new GeoIP database
func NewGeoIPManager() (*GeoIPManager, error) {
	manager := &GeoIPManager{}
	if err := manager.Initialize(); err != nil {
		return nil, fmt.Errorf("initializing GeoIP manager: %w", err)
	}
	return manager, nil
}

// Initialize opens the GeoIP databases, downloading them if necessary
func (g *GeoIPManager) Initialize() error {
	g.mu.Lock()
	defer g.mu.Unlock()

	if err := g.openCityDB(); err != nil {
		return err
	}

	if err := g.openASNDB(); err != nil {
		return err
	}

	g.buildASNPrefixMap()

	return nil
}

// buildASNPrefixMap iterates the ASN database once and builds the lookup map.
func (g *GeoIPManager) buildASNPrefixMap() {
	log.Println("Building ASN prefix map for fast lookups...")
	startTime := time.Now()

	g.asnPrefixMap = make(map[uint][]*net.IPNet)
	networks := g.asnDB.Networks()

	for networks.Next() {
		var record ASNRecord
		subnet, err := networks.Network(&record)
		if err != nil {
			continue // Skip records that fail to decode
		}
		if record.AutonomousSystemNumber > 0 {
			g.asnPrefixMap[record.AutonomousSystemNumber] = append(g.asnPrefixMap[record.AutonomousSystemNumber], subnet)
		}
	}
	log.Printf("Finished building ASN prefix map in %v", time.Since(startTime))
}

// Opens the city database, downloading it if necessary
func (g *GeoIPManager) openCityDB() error {
	var err error
	g.cityDB, err = maxminddb.Open(CityDBPath)
	if err != nil {
		if os.IsNotExist(err) {
			log.Println("City database not found, attempting to download...")
			if err := g.downloadDatabases(); err != nil {
				return fmt.Errorf("%w: %v", ErrDownloadFailed, err)
			}
			g.cityDB, err = maxminddb.Open(CityDBPath)
			if err != nil {
				return fmt.Errorf("%w (city): %v", ErrDatabaseOpen, err)
			}
			return nil
		}
		return fmt.Errorf("%w (city): %v", ErrDatabaseOpen, err)
	}
	return nil
}

// Opens the ASN database, downloading it if necessary
func (g *GeoIPManager) openASNDB() error {
	var err error
	g.asnDB, err = maxminddb.Open(ASNDBPath)
	if err != nil {
		if os.IsNotExist(err) {
			log.Println("ASN database not found, attempting to download...")
			if err := g.downloadDatabases(); err != nil {
				return fmt.Errorf("%w: %v", ErrDownloadFailed, err)
			}
			g.asnDB, err = maxminddb.Open(ASNDBPath)
			if err != nil {
				return fmt.Errorf("%w (ASN): %v", ErrDatabaseOpen, err)
			}
			return nil
		}
		return fmt.Errorf("%w (ASN): %v", ErrDatabaseOpen, err)
	}
	return nil
}

// Downloads both GeoIP databases using geoipupdate
func (g *GeoIPManager) downloadDatabases() error {
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
	defer cancel()

	cmd := exec.CommandContext(ctx, "geoipupdate", "-d", "./")
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("failed to download databases: %v. Output: %s. Ensure geoipupdate is installed and configured", err, output)
	}
	return nil
}

// Close properly closes the database readers
func (g *GeoIPManager) Close() error {
	g.mu.Lock()
	defer g.mu.Unlock()

	var errs []error

	if g.cityDB != nil {
		if err := g.cityDB.Close(); err != nil {
			errs = append(errs, fmt.Errorf("closing city database: %w", err))
		}
	}

	if g.asnDB != nil {
		if err := g.asnDB.Close(); err != nil {
			errs = append(errs, fmt.Errorf("closing ASN database: %w", err))
		}
	}

	if len(errs) > 0 {
		return fmt.Errorf("errors closing databases: %v", errs)
	}
	return nil
}

// Safely provides read access to the City database
func (g *GeoIPManager) GetCityDB() *maxminddb.Reader {
	g.mu.RLock()
	defer g.mu.RUnlock()
	return g.cityDB
}

// Safely provides read access to the ASN database
func (g *GeoIPManager) GetASNDB() *maxminddb.Reader {
	g.mu.RLock()
	defer g.mu.RUnlock()
	return g.asnDB
}

// GetASNPrefixes returns the pre-computed list of prefixes for an ASN.
func (g *GeoIPManager) GetASNPrefixes(asn uint) []*net.IPNet {
	g.mu.RLock()
	defer g.mu.RUnlock()
	return g.asnPrefixMap[asn]
}

// Sets up automatic database updates
func (g *GeoIPManager) StartUpdater(ctx context.Context, updateInterval time.Duration) {
	log.Printf("Starting MaxMind GeoIP database updater with interval: %s", updateInterval)

	ticker := time.NewTicker(updateInterval)
	go func() {
		for {
			select {
			case <-ticker.C:
				log.Println("Performing scheduled GeoIP database update")
				if err := g.UpdateDatabases(); err != nil {
					log.Printf("Failed to update databases: %v", err)
				}
			case <-ctx.Done():
				ticker.Stop()
				log.Println("GeoIP database updater stopped")
				return
			}
		}
	}()
}

// Downloads fresh copies of the databases and reloads them
func (g *GeoIPManager) UpdateDatabases() error {
	// Download new databases
	if err := g.downloadDatabases(); err != nil {
		return err
	}

	g.mu.Lock()
	defer g.mu.Unlock()

	// Close existing databases
	if g.cityDB != nil {
		if err := g.cityDB.Close(); err != nil {
			log.Printf("Warning: error closing city database: %v", err)
		}
	}

	if g.asnDB != nil {
		if err := g.asnDB.Close(); err != nil {
			log.Printf("Warning: error closing ASN database: %v", err)
		}
	}

	// Reopen databases
	var err error
	g.cityDB, err = maxminddb.Open(CityDBPath)
	if err != nil {
		return fmt.Errorf("reopening city database: %w", err)
	}

	g.asnDB, err = maxminddb.Open(ASNDBPath)
	if err != nil {
		return fmt.Errorf("reopening ASN database: %w", err)
	}

	// Rebuild the prefix map with the new data.
	g.buildASNPrefixMap()

	log.Println("Successfully updated and reloaded GeoIP databases")
	return nil
}