From e052f767db8fee73abfcc6ccd24f1f0595e17bcb Mon Sep 17 00:00:00 2001
From: Toby <tobyxdd@gmail.com>
Date: Mon, 13 Nov 2023 20:27:08 -0800
Subject: [PATCH] feat: geoUpdateInterval

---
 app/cmd/server.go               | 10 ++++++----
 app/cmd/server_test.go          |  5 +++--
 app/cmd/server_test.yaml        |  1 +
 app/internal/utils/geoloader.go | 33 +++++++++++++++++++++++----------
 4 files changed, 33 insertions(+), 16 deletions(-)

diff --git a/app/cmd/server.go b/app/cmd/server.go
index 385849f..9ec95e9 100644
--- a/app/cmd/server.go
+++ b/app/cmd/server.go
@@ -147,10 +147,11 @@ type serverConfigResolver struct {
 }
 
 type serverConfigACL struct {
-	File    string   `mapstructure:"file"`
-	Inline  []string `mapstructure:"inline"`
-	GeoIP   string   `mapstructure:"geoip"`
-	GeoSite string   `mapstructure:"geosite"`
+	File              string        `mapstructure:"file"`
+	Inline            []string      `mapstructure:"inline"`
+	GeoIP             string        `mapstructure:"geoip"`
+	GeoSite           string        `mapstructure:"geosite"`
+	GeoUpdateInterval time.Duration `mapstructure:"geoUpdateInterval"`
 }
 
 type serverConfigOutboundDirect struct {
@@ -469,6 +470,7 @@ func (c *serverConfig) fillOutboundConfig(hyConfig *server.Config) error {
 	gLoader := &utils.GeoLoader{
 		GeoIPFilename:   c.ACL.GeoIP,
 		GeoSiteFilename: c.ACL.GeoSite,
+		UpdateInterval:  c.ACL.GeoUpdateInterval,
 		DownloadFunc:    geoDownloadFunc,
 		DownloadErrFunc: geoDownloadErrFunc,
 	}
diff --git a/app/cmd/server_test.go b/app/cmd/server_test.go
index 8eaafff..1c4d2f6 100644
--- a/app/cmd/server_test.go
+++ b/app/cmd/server_test.go
@@ -101,8 +101,9 @@ func TestServerConfig(t *testing.T) {
 				"lmao(ok)",
 				"kek(cringe,boba,tea)",
 			},
-			GeoIP:   "some.dat",
-			GeoSite: "some_site.dat",
+			GeoIP:             "some.dat",
+			GeoSite:           "some_site.dat",
+			GeoUpdateInterval: 168 * time.Hour,
 		},
 		Outbounds: []serverConfigOutboundEntry{
 			{
diff --git a/app/cmd/server_test.yaml b/app/cmd/server_test.yaml
index 77a791c..47d3b19 100644
--- a/app/cmd/server_test.yaml
+++ b/app/cmd/server_test.yaml
@@ -77,6 +77,7 @@ acl:
     - kek(cringe,boba,tea)
   geoip: some.dat
   geosite: some_site.dat
+  geoUpdateInterval: 168h
 
 outbounds:
   - name: goodstuff
diff --git a/app/internal/utils/geoloader.go b/app/internal/utils/geoloader.go
index ed8b993..56cb205 100644
--- a/app/internal/utils/geoloader.go
+++ b/app/internal/utils/geoloader.go
@@ -4,6 +4,7 @@ import (
 	"io"
 	"net/http"
 	"os"
+	"time"
 
 	"github.com/apernet/hysteria/extras/outbounds/acl"
 	"github.com/apernet/hysteria/extras/outbounds/acl/v2geo"
@@ -14,6 +15,8 @@ const (
 	geoipURL        = "https://cdn.jsdelivr.net/gh/Loyalsoldier/v2ray-rules-dat@release/geoip.dat"
 	geositeFilename = "geosite.dat"
 	geositeURL      = "https://cdn.jsdelivr.net/gh/Loyalsoldier/v2ray-rules-dat@release/geosite.dat"
+
+	geoDefaultUpdateInterval = 7 * 24 * time.Hour // 7 days
 )
 
 var _ acl.GeoLoader = (*GeoLoader)(nil)
@@ -24,6 +27,7 @@ var _ acl.GeoLoader = (*GeoLoader)(nil)
 type GeoLoader struct {
 	GeoIPFilename   string
 	GeoSiteFilename string
+	UpdateInterval  time.Duration
 
 	DownloadFunc    func(filename, url string)
 	DownloadErrFunc func(err error)
@@ -32,6 +36,19 @@ type GeoLoader struct {
 	geositeMap map[string]*v2geo.GeoSite
 }
 
+func (l *GeoLoader) shouldDownload(filename string) bool {
+	info, err := os.Stat(filename)
+	if os.IsNotExist(err) {
+		return true
+	}
+	dt := time.Now().Sub(info.ModTime())
+	if l.UpdateInterval == 0 {
+		return dt > geoDefaultUpdateInterval
+	} else {
+		return dt > l.UpdateInterval
+	}
+}
+
 func (l *GeoLoader) download(filename, url string) error {
 	l.DownloadFunc(filename, url)
 
@@ -64,15 +81,13 @@ func (l *GeoLoader) LoadGeoIP() (map[string]*v2geo.GeoIP, error) {
 		autoDL = true
 		filename = geoipFilename
 	}
-	m, err := v2geo.LoadGeoIP(filename)
-	if os.IsNotExist(err) && autoDL {
-		// It's ok, we will download it.
-		err = l.download(filename, geoipURL)
+	if autoDL && l.shouldDownload(filename) {
+		err := l.download(filename, geoipURL)
 		if err != nil {
 			return nil, err
 		}
-		m, err = v2geo.LoadGeoIP(filename)
 	}
+	m, err := v2geo.LoadGeoIP(filename)
 	if err != nil {
 		return nil, err
 	}
@@ -90,15 +105,13 @@ func (l *GeoLoader) LoadGeoSite() (map[string]*v2geo.GeoSite, error) {
 		autoDL = true
 		filename = geositeFilename
 	}
-	m, err := v2geo.LoadGeoSite(filename)
-	if os.IsNotExist(err) && autoDL {
-		// It's ok, we will download it.
-		err = l.download(filename, geositeURL)
+	if autoDL && l.shouldDownload(filename) {
+		err := l.download(filename, geositeURL)
 		if err != nil {
 			return nil, err
 		}
-		m, err = v2geo.LoadGeoSite(filename)
 	}
+	m, err := v2geo.LoadGeoSite(filename)
 	if err != nil {
 		return nil, err
 	}