ip-info/internal/data/ipdb.go

package data

import (
	"fmt"
	"github.com/rs/zerolog/log"
	"io"
	"os"
	"path/filepath"
	"sync"
	"thuanle.me/ip-info/configs"
)
import reader "github.com/oschwald/maxminddb-golang/v2"

type IpDb struct {
	r  *reader.Reader
	mu sync.RWMutex

	dbFile string
}

var (
	ins  *IpDb
	once sync.Once
)

func Ins() *IpDb {
	once.Do(func() {
		ins = &IpDb{}
		_ = ins.Reload()
	})
	return ins
}

func CleanupDataDir() error {
	if _, err := os.Stat(configs.GeoDbFolder); os.IsNotExist(err) {
		log.Info().Str("dir", configs.GeoDbFolder).Msg("Creating data folder")
		err := os.MkdirAll(configs.GeoDbFolder, os.ModePerm)
		if err != nil {
			return err
		}
	}

	log.Info().Str("dir", configs.GeoDbFolder).Msg("Cleaning temp files")
	dir, err := os.Open(configs.GeoDbFolder)
	if err != nil {
		return err
	}
	defer dir.Close()

	// List all files in the directory
	files, err := dir.Readdir(-1) // -1 means read all files
	if err != nil {
		return err
	}

	// Delete each file
	for _, file := range files {
		if !file.IsDir() && file.Name() != configs.MmdbDbFileName {
			filePath := filepath.Join(configs.GeoDbFolder, file.Name())
			err := os.Remove(filePath)
			if err != nil {
				return fmt.Errorf("failed to delete file %s: %w", filePath, err)
			}
			log.Info().Str("file", filePath).Msg("Deleted temp file")
		}
	}

	return nil
}

func (d *IpDb) Reload() error {
	d.mu.Lock()
	defer d.mu.Unlock()

	wilDeleteFile := d.dbFile

	tmpFile, err := cloneDBFile()
	if err != nil {
		log.Err(err).Msg("Failed to clone db file")
		return err
	}

	d.dbFile = tmpFile
	r, err := reader.Open(tmpFile)
	if err != nil {
		log.Err(err).Msg("Failed to open mmdb")
		return err
	}
	log.Info().Str("file", tmpFile).Msg("MMDB reloaded")

	tmpR := d.r
	d.r = r

	if tmpR != nil {
		_ = tmpR.Close()

		if wilDeleteFile != "" {
			log.Info().Str("file", wilDeleteFile).Msg("Deleting old mmdb")
			_ = os.Remove(wilDeleteFile)
		}
	}

	return nil
}

func (d *IpDb) IsLoaded() bool {
	d.mu.RLock()
	defer d.mu.RUnlock()
	return d.r != nil
}

func cloneDBFile() (string, error) {
	srcFile, err := os.Open(configs.MmdbDbFile)
	if err != nil {
		return "", err
	}
	defer srcFile.Close()

	tmpFile, err := os.CreateTemp(configs.GeoDbFolder, "geoip-v4-*.mmdb")
	if err != nil {
		return "", fmt.Errorf("failed to create temporary file: %v", err)
	}
	defer tmpFile.Close()

	buf := make([]byte, 1024*1024) // 1 MB buffer

	// Copy the data from srcFile to dstFile
	_, err = io.CopyBuffer(tmpFile, srcFile, buf)
	if err != nil {
		return "", err
	}

	return tmpFile.Name(), nil
}