feat(acl): add FORCE_RESOLVE_COUNTRY option to resolve country

internal/acl/config.go

@@ -9,6 +9,7 @@ import (
 	"github.com/puzpuzpuz/xsync/v3"
 	"github.com/rs/zerolog"
 	acl "github.com/yusing/go-proxy/internal/acl/types"
+	"github.com/yusing/go-proxy/internal/common"
 	"github.com/yusing/go-proxy/internal/gperr"
 	"github.com/yusing/go-proxy/internal/logging"
 	"github.com/yusing/go-proxy/internal/logging/accesslog"
@@ -153,7 +154,10 @@ func (c *Config) Start(parent *task.Task) gperr.Error {
 	return nil
 }
 
-func (c *config) cacheRecord(info *acl.IPInfo, allow bool) {
+func (c *Config) cacheRecord(info *acl.IPInfo, allow bool) {
+	if common.ForceResolveCountry && info.City == nil {
+		c.MaxMind.lookupCity(info)
+	}
 	c.ipCache.Store(info.Str, &checkCache{
 		IPInfo:  info,
 		allow:   allow,
@@ -175,7 +179,7 @@ func (c *Config) IPAllowed(ip net.IP) bool {
 		return false
 	}
 
-	// always allow private and loopback
+	// always allow loopback
 	// loopback is not logged
 	if ip.IsLoopback() {
 		return true

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
 	}

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) {

internal/common/env.go

@@ -58,6 +58,8 @@ var (
 	MetricsDisableDisk    = GetEnvBool("METRICS_DISABLE_DISK", false)
 	MetricsDisableNetwork = GetEnvBool("METRICS_DISABLE_NETWORK", false)
 	MetricsDisableSensors = GetEnvBool("METRICS_DISABLE_SENSORS", false)
+
+	ForceResolveCountry = GetEnvBool("FORCE_RESOLVE_COUNTRY", false)
 )
 
 func GetEnv[T any](key string, defaultValue T, parser func(string) (T, error)) T {

internal/notif/dispatcher.go

@@ -83,9 +83,6 @@ func (disp *Dispatcher) start() {
 }
 
 func (disp *Dispatcher) dispatch(msg *LogMessage) {
-	if true {
-		return
-	}
 	task := disp.task.Subtask("dispatcher")
 	defer task.Finish("notif dispatched")