chore(maxm): improved database update mechanism, fixed db being downloaded twice on first run

This commit is contained in:
yusing 2025-04-26 09:08:03 +08:00
parent 6a8f6fb4b5
commit e11579df10
2 changed files with 98 additions and 66 deletions

View file

@ -54,7 +54,7 @@ func (cfg *MaxMindConfig) LoadMaxMindDB(parent task.Parent) gperr.Error {
path := dbPath(cfg.Database) path := dbPath(cfg.Database)
reader, err := maxmindDBOpen(path) reader, err := maxmindDBOpen(path)
exists := true valid := true
if err != nil { if err != nil {
switch { switch {
case errors.Is(err, os.ErrNotExist): case errors.Is(err, os.ErrNotExist):
@ -65,20 +65,19 @@ func (cfg *MaxMindConfig) LoadMaxMindDB(parent task.Parent) gperr.Error {
return gperr.Wrap(err) return gperr.Wrap(err)
} }
} }
exists = false valid = false
} }
if !exists { if !valid {
cfg.logger.Info().Msg("MaxMind DB not found/invalid, downloading...") cfg.logger.Info().Msg("MaxMind DB not found/invalid, downloading...")
reader, err = cfg.download() if err = cfg.download(); err != nil {
if err != nil {
return ErrDownloadFailure.With(err) return ErrDownloadFailure.With(err)
} }
} else {
cfg.logger.Info().Msg("MaxMind DB loaded")
cfg.db.Reader = reader
go cfg.scheduleUpdate(parent)
} }
cfg.logger.Info().Msg("MaxMind DB loaded")
cfg.db.Reader = reader
go cfg.scheduleUpdate(parent)
return nil return nil
} }
@ -137,17 +136,10 @@ func (cfg *MaxMindConfig) update() {
Time("latest", remoteLastModified.Local()). Time("latest", remoteLastModified.Local()).
Time("current", cfg.lastUpdate). Time("current", cfg.lastUpdate).
Msg("MaxMind DB update available") Msg("MaxMind DB update available")
reader, err := cfg.download() if err = cfg.download(); err != nil {
if err != nil {
cfg.logger.Err(err).Msg("failed to update MaxMind DB") cfg.logger.Err(err).Msg("failed to update MaxMind DB")
return return
} }
cfg.db.Lock()
cfg.db.Close()
cfg.db.Reader = reader
cfg.setLastUpdate(*remoteLastModified)
cfg.db.Unlock()
cfg.logger.Info().Msg("MaxMind DB updated") cfg.logger.Info().Msg("MaxMind DB updated")
} }
@ -190,57 +182,87 @@ func (cfg *MaxMindConfig) checkLastest() (lastModifiedT *time.Time, err error) {
return &lastModifiedTime, nil return &lastModifiedTime, nil
} }
func (cfg *MaxMindConfig) download() (*maxminddb.Reader, error) { func (cfg *MaxMindConfig) download() error {
resp, err := newReq(cfg, http.MethodGet) resp, err := newReq(cfg, http.MethodGet)
if err != nil { if err != nil {
return nil, err return err
} }
defer resp.Body.Close() defer resp.Body.Close()
if resp.StatusCode != http.StatusOK { if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("%w: %d", ErrResponseNotOK, resp.StatusCode) return fmt.Errorf("%w: %d", ErrResponseNotOK, resp.StatusCode)
} }
path := dbPath(cfg.Database) dbFile := dbPath(cfg.Database)
tmpPath := path + "-tmp.tar.gz" tmpGZPath := dbFile + "-tmp.tar.gz"
file, err := os.OpenFile(tmpPath, os.O_CREATE|os.O_WRONLY, 0o644) tmpDBPath := dbFile + "-tmp"
if err != nil {
return nil, err
}
cfg.logger.Info().Msg("MaxMind DB downloading...") tmpGZFile, err := os.OpenFile(tmpGZPath, os.O_CREATE|os.O_RDWR, 0o644)
_, err = io.Copy(file, resp.Body)
if err != nil {
file.Close()
return nil, err
}
file.Close()
// extract .tar.gz and move only the dbFilename to path
err = extractFileFromTarGz(tmpPath, dbFilename(cfg.Database), path)
if err != nil {
return nil, gperr.New("failed to extract database from archive").With(err)
}
// cleanup the tar.gz file
_ = os.Remove(tmpPath)
db, err := maxmindDBOpen(path)
if err != nil {
return nil, err
}
return db, nil
}
func extractFileFromTarGz(tarGzPath, targetFilename, destPath string) error {
f, err := os.Open(tarGzPath)
if err != nil { if err != nil {
return err return err
} }
defer f.Close()
gzr, err := gzip.NewReader(f) // 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 { if err != nil {
return err return err
} }

View file

@ -1,6 +1,8 @@
package acl package acl
import ( import (
"archive/tar"
"compress/gzip"
"io" "io"
"net/http" "net/http"
"net/http/httptest" "net/http/httptest"
@ -144,9 +146,17 @@ func Test_MaxMindConfig_download(t *testing.T) {
logger: zerolog.Nop(), logger: zerolog.Nop(),
} }
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
io.Copy(w, strings.NewReader("FAKEMMDB")) gz := gzip.NewWriter(w)
t := tar.NewWriter(gz)
t.WriteHeader(&tar.Header{
Name: dbFilename(MaxMindGeoLite),
})
t.Write([]byte("1234"))
t.Close()
gz.Close()
})) }))
defer server.Close() defer server.Close()
oldURL := dbURL oldURL := dbURL
dbURL = func(MaxMindDatabaseType) string { return server.URL } dbURL = func(MaxMindDatabaseType) string { return server.URL }
defer func() { dbURL = oldURL }() defer func() { dbURL = oldURL }()
@ -163,26 +173,26 @@ func Test_MaxMindConfig_download(t *testing.T) {
} }
defer func() { maxmindDBOpen = origOpen }() defer func() { maxmindDBOpen = origOpen }()
rw := &fakeReadCloser{} req, err := http.NewRequest(http.MethodGet, server.URL, nil)
if err != nil {
t.Fatalf("newReq() error = %v", err)
}
rw := httptest.NewRecorder()
oldNewReq := newReq oldNewReq := newReq
newReq = func(cfg *MaxMindConfig, method string) (*http.Response, error) { newReq = func(cfg *MaxMindConfig, method string) (*http.Response, error) {
return &http.Response{ server.Config.Handler.ServeHTTP(rw, req)
StatusCode: http.StatusOK, return rw.Result(), nil
Body: rw,
}, nil
} }
defer func() { newReq = oldNewReq }() defer func() { newReq = oldNewReq }()
db, err := cfg.download() err = cfg.download()
if err != nil { if err != nil {
t.Fatalf("download() error = %v", err) t.Fatalf("download() error = %v", err)
} }
if db == nil { if cfg.db.Reader == nil {
t.Error("expected db instance") t.Error("expected db instance")
} }
if !rw.closed {
t.Error("expected rw to be closed")
}
} }
func Test_MaxMindConfig_loadMaxMindDB(t *testing.T) { func Test_MaxMindConfig_loadMaxMindDB(t *testing.T) {