package acl import ( "archive/tar" "compress/gzip" "errors" "fmt" "io" "net/http" "os" "path/filepath" "time" "github.com/oschwald/maxminddb-golang" "github.com/yusing/go-proxy/internal/common" "github.com/yusing/go-proxy/internal/gperr" "github.com/yusing/go-proxy/internal/task" ) var ( updateInterval = 24 * time.Hour httpClient = &http.Client{ Timeout: 10 * time.Second, } ErrResponseNotOK = gperr.New("response not OK") ErrDownloadFailure = gperr.New("download failure") ) func dbPathImpl(dbType MaxMindDatabaseType) string { if dbType == MaxMindGeoLite { return filepath.Join(dataDir, "GeoLite2-City.mmdb") } return filepath.Join(dataDir, "GeoIP2-City.mmdb") } func dbURLimpl(dbType MaxMindDatabaseType) string { if dbType == MaxMindGeoLite { return "https://download.maxmind.com/geoip/databases/GeoLite2-City/download?suffix=tar.gz" } return "https://download.maxmind.com/geoip/databases/GeoIP2-City/download?suffix=tar.gz" } func dbFilename(dbType MaxMindDatabaseType) string { if dbType == MaxMindGeoLite { return "GeoLite2-City.mmdb" } return "GeoIP2-City.mmdb" } func (cfg *MaxMindConfig) LoadMaxMindDB(parent task.Parent) gperr.Error { if cfg.Database == "" { return nil } path := dbPath(cfg.Database) reader, err := maxmindDBOpen(path) valid := true if err != nil { switch { case errors.Is(err, os.ErrNotExist): default: // ignore invalid error, just download it again var invalidErr maxminddb.InvalidDatabaseError if !errors.As(err, &invalidErr) { return gperr.Wrap(err) } } valid = false } if !valid { cfg.logger.Info().Msg("MaxMind DB not found/invalid, downloading...") if err = cfg.download(); err != nil { return ErrDownloadFailure.With(err) } } else { cfg.logger.Info().Msg("MaxMind DB loaded") cfg.db.Reader = reader go cfg.scheduleUpdate(parent) } return nil } func (cfg *MaxMindConfig) loadLastUpdate() { f, err := os.Stat(dbPath(cfg.Database)) if err != nil { return } cfg.lastUpdate = f.ModTime() } func (cfg *MaxMindConfig) setLastUpdate(t time.Time) { cfg.lastUpdate = t _ = os.Chtimes(dbPath(cfg.Database), t, t) } func (cfg *MaxMindConfig) scheduleUpdate(parent task.Parent) { task := parent.Subtask("schedule_update", true) ticker := time.NewTicker(updateInterval) cfg.loadLastUpdate() cfg.update() defer func() { ticker.Stop() if cfg.db.Reader != nil { cfg.db.Reader.Close() } task.Finish(nil) }() for { select { case <-task.Context().Done(): return case <-ticker.C: cfg.update() } } } func (cfg *MaxMindConfig) update() { // check for update cfg.logger.Info().Msg("checking for MaxMind DB update...") remoteLastModified, err := cfg.checkLastest() if err != nil { cfg.logger.Err(err).Msg("failed to check MaxMind DB update") return } if remoteLastModified.Equal(cfg.lastUpdate) { cfg.logger.Info().Msg("MaxMind DB is up to date") return } cfg.logger.Info(). Time("latest", remoteLastModified.Local()). Time("current", cfg.lastUpdate). Msg("MaxMind DB update available") if err = cfg.download(); err != nil { cfg.logger.Err(err).Msg("failed to update MaxMind DB") return } cfg.logger.Info().Msg("MaxMind DB updated") } func (cfg *MaxMindConfig) newReq(method string) (*http.Response, error) { req, err := http.NewRequest(method, dbURL(cfg.Database), nil) if err != nil { return nil, err } req.SetBasicAuth(cfg.AccountID, cfg.LicenseKey) resp, err := httpClient.Do(req) if err != nil { return nil, err } return resp, nil } func (cfg *MaxMindConfig) checkLastest() (lastModifiedT *time.Time, err error) { resp, err := newReq(cfg, http.MethodHead) if err != nil { return nil, err } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { return nil, fmt.Errorf("%w: %d", ErrResponseNotOK, resp.StatusCode) } lastModified := resp.Header.Get("Last-Modified") if lastModified == "" { cfg.logger.Warn().Msg("MaxMind responded no last modified time, update skipped") return nil, nil } lastModifiedTime, err := time.Parse(http.TimeFormat, lastModified) if err != nil { cfg.logger.Warn().Err(err).Msg("MaxMind responded invalid last modified time, update skipped") return nil, err } return &lastModifiedTime, nil } func (cfg *MaxMindConfig) download() error { resp, err := newReq(cfg, http.MethodGet) if err != nil { return err } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { return fmt.Errorf("%w: %d", ErrResponseNotOK, resp.StatusCode) } dbFile := dbPath(cfg.Database) tmpGZPath := dbFile + "-tmp.tar.gz" tmpDBPath := dbFile + "-tmp" tmpGZFile, err := os.OpenFile(tmpGZPath, os.O_CREATE|os.O_RDWR, 0o644) if err != nil { return err } // cleanup the tar.gz file defer func() { _ = tmpGZFile.Close() _ = os.Remove(tmpGZPath) }() cfg.logger.Info().Msg("MaxMind DB downloading...") _, err = io.Copy(tmpGZFile, resp.Body) if err != nil { return err } if _, err := tmpGZFile.Seek(0, io.SeekStart); err != nil { return err } // extract .tar.gz and to database err = extractFileFromTarGz(tmpGZFile, dbFilename(cfg.Database), tmpDBPath) if err != nil { return gperr.New("failed to extract database from archive").With(err) } // test if the downloaded database is valid db, err := maxmindDBOpen(tmpDBPath) if err != nil { _ = os.Remove(tmpDBPath) return err } db.Close() err = os.Rename(tmpDBPath, dbFile) if err != nil { return err } cfg.db.Lock() defer cfg.db.Unlock() if cfg.db.Reader != nil { cfg.db.Reader.Close() } cfg.db.Reader, err = maxmindDBOpen(dbFile) if err != nil { return err } lastModifiedStr := resp.Header.Get("Last-Modified") lastModifiedTime, err := time.Parse(http.TimeFormat, lastModifiedStr) if err == nil { cfg.setLastUpdate(lastModifiedTime) } cfg.logger.Info().Msg("MaxMind DB downloaded") return nil } func extractFileFromTarGz(tarGzFile *os.File, targetFilename, destPath string) error { defer tarGzFile.Close() gzr, err := gzip.NewReader(tarGzFile) if err != nil { return err } defer gzr.Close() tr := tar.NewReader(gzr) for { hdr, err := tr.Next() if err == io.EOF { break // End of archive } if err != nil { return err } // Only extract the file that matches targetFilename (basename match) if filepath.Base(hdr.Name) == targetFilename { outFile, err := os.OpenFile(destPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, hdr.FileInfo().Mode()) if err != nil { return err } defer outFile.Close() _, err = io.Copy(outFile, tr) if err != nil { return err } return nil // Done } } return fmt.Errorf("file %s not found in archive", targetFilename) } var ( dataDir = common.DataDir dbURL = dbURLimpl dbPath = dbPathImpl maxmindDBOpen = maxminddb.Open newReq = (*MaxMindConfig).newReq )