From e11579df10457f4fae09e0caba4d615109bd77a0 Mon Sep 17 00:00:00 2001 From: yusing Date: Sat, 26 Apr 2025 09:08:03 +0800 Subject: [PATCH] chore(maxm): improved database update mechanism, fixed db being downloaded twice on first run --- internal/acl/maxmind.go | 132 ++++++++++++++++++++--------------- internal/acl/maxmind_test.go | 32 ++++++--- 2 files changed, 98 insertions(+), 66 deletions(-) diff --git a/internal/acl/maxmind.go b/internal/acl/maxmind.go index 20684470..eed630e5 100644 --- a/internal/acl/maxmind.go +++ b/internal/acl/maxmind.go @@ -54,7 +54,7 @@ func (cfg *MaxMindConfig) LoadMaxMindDB(parent task.Parent) gperr.Error { path := dbPath(cfg.Database) reader, err := maxmindDBOpen(path) - exists := true + valid := true if err != nil { switch { case errors.Is(err, os.ErrNotExist): @@ -65,20 +65,19 @@ func (cfg *MaxMindConfig) LoadMaxMindDB(parent task.Parent) gperr.Error { return gperr.Wrap(err) } } - exists = false + valid = false } - if !exists { + if !valid { cfg.logger.Info().Msg("MaxMind DB not found/invalid, downloading...") - reader, err = cfg.download() - if err != nil { + 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) } - cfg.logger.Info().Msg("MaxMind DB loaded") - - cfg.db.Reader = reader - go cfg.scheduleUpdate(parent) return nil } @@ -137,17 +136,10 @@ func (cfg *MaxMindConfig) update() { Time("latest", remoteLastModified.Local()). Time("current", cfg.lastUpdate). Msg("MaxMind DB update available") - reader, err := cfg.download() - if err != nil { + if err = cfg.download(); err != nil { cfg.logger.Err(err).Msg("failed to update MaxMind DB") return } - cfg.db.Lock() - cfg.db.Close() - cfg.db.Reader = reader - cfg.setLastUpdate(*remoteLastModified) - cfg.db.Unlock() - cfg.logger.Info().Msg("MaxMind DB updated") } @@ -190,57 +182,87 @@ func (cfg *MaxMindConfig) checkLastest() (lastModifiedT *time.Time, err error) { return &lastModifiedTime, nil } -func (cfg *MaxMindConfig) download() (*maxminddb.Reader, error) { +func (cfg *MaxMindConfig) download() error { resp, err := newReq(cfg, http.MethodGet) if err != nil { - return nil, err + return err } defer resp.Body.Close() 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) - tmpPath := path + "-tmp.tar.gz" - file, err := os.OpenFile(tmpPath, os.O_CREATE|os.O_WRONLY, 0o644) - if err != nil { - return nil, err - } + dbFile := dbPath(cfg.Database) + tmpGZPath := dbFile + "-tmp.tar.gz" + tmpDBPath := dbFile + "-tmp" - cfg.logger.Info().Msg("MaxMind DB downloading...") - - _, 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) + tmpGZFile, err := os.OpenFile(tmpGZPath, os.O_CREATE|os.O_RDWR, 0o644) if err != nil { 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 { return err } diff --git a/internal/acl/maxmind_test.go b/internal/acl/maxmind_test.go index 62f2cca8..6c0c0c95 100644 --- a/internal/acl/maxmind_test.go +++ b/internal/acl/maxmind_test.go @@ -1,6 +1,8 @@ package acl import ( + "archive/tar" + "compress/gzip" "io" "net/http" "net/http/httptest" @@ -144,9 +146,17 @@ func Test_MaxMindConfig_download(t *testing.T) { logger: zerolog.Nop(), } 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() + oldURL := dbURL dbURL = func(MaxMindDatabaseType) string { return server.URL } defer func() { dbURL = oldURL }() @@ -163,26 +173,26 @@ func Test_MaxMindConfig_download(t *testing.T) { } 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 newReq = func(cfg *MaxMindConfig, method string) (*http.Response, error) { - return &http.Response{ - StatusCode: http.StatusOK, - Body: rw, - }, nil + server.Config.Handler.ServeHTTP(rw, req) + return rw.Result(), nil } defer func() { newReq = oldNewReq }() - db, err := cfg.download() + err = cfg.download() if err != nil { t.Fatalf("download() error = %v", err) } - if db == nil { + if cfg.db.Reader == nil { t.Error("expected db instance") } - if !rw.closed { - t.Error("expected rw to be closed") - } } func Test_MaxMindConfig_loadMaxMindDB(t *testing.T) {