Feature: add geoip-code option

itsevin · Aug 25, 2021 · e0d3f92 · e0d3f92
1 parent 121bc91
commit e0d3f92
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 17 deletions.
diff --git a/config/config.go b/config/config.go
@@ -71,9 +71,10 @@ type DNS struct {
 
 // FallbackFilter config
 type FallbackFilter struct {
-	GeoIP  bool         `yaml:"geoip"`
-	IPCIDR []*net.IPNet `yaml:"ipcidr"`
-	Domain []string     `yaml:"domain"`
+	GeoIP     bool         `yaml:"geoip"`
+	GeoIPCode string       `yaml:"geoip-code"`
+	IPCIDR    []*net.IPNet `yaml:"ipcidr"`
+	Domain    []string     `yaml:"domain"`
 }
 
 // Profile config
@@ -113,9 +114,10 @@ type RawDNS struct {
 }
 
 type RawFallbackFilter struct {
-	GeoIP  bool     `yaml:"geoip"`
-	IPCIDR []string `yaml:"ipcidr"`
-	Domain []string `yaml:"domain"`
+	GeoIP     bool     `yaml:"geoip"`
+	GeoIPCode string   `yaml:"geoip-code"`
+	IPCIDR    []string `yaml:"ipcidr"`
+	Domain    []string `yaml:"domain"`
 }
 
 type RawConfig struct {
@@ -172,8 +174,9 @@ func UnmarshalRawConfig(buf []byte) (*RawConfig, error) {
 			UseHosts:    true,
 			FakeIPRange: "198.18.0.1/16",
 			FallbackFilter: RawFallbackFilter{
-				GeoIP:  true,
-				IPCIDR: []string{},
+				GeoIP:     true,
+				GeoIPCode: "CN",
+				IPCIDR:    []string{},
 			},
 			DefaultNameserver: []string{
 				"114.114.114.114",
@@ -600,6 +603,7 @@ func parseDNS(cfg RawDNS, hosts *trie.DomainTrie) (*DNS, error) {
 	}
 
 	dnsCfg.FallbackFilter.GeoIP = cfg.FallbackFilter.GeoIP
+	dnsCfg.FallbackFilter.GeoIPCode = cfg.FallbackFilter.GeoIPCode
 	if fallbackip, err := parseFallbackIPCIDR(cfg.FallbackFilter.IPCIDR); err == nil {
 		dnsCfg.FallbackFilter.IPCIDR = fallbackip
 	}

diff --git a/dns/filters.go b/dns/filters.go
@@ -11,11 +11,13 @@ type fallbackIPFilter interface {
 	Match(net.IP) bool
 }
 
-type geoipFilter struct{}
+type geoipFilter struct {
+	code string
+}
 
 func (gf *geoipFilter) Match(ip net.IP) bool {
 	record, _ := mmdb.Instance().Country(ip)
-	return record.Country.IsoCode != "CN" && !ip.IsPrivate()
+	return record.Country.IsoCode != gf.code && !ip.IsPrivate()
 }
 
 type ipnetFilter struct {

diff --git a/dns/resolver.go b/dns/resolver.go
@@ -302,9 +302,10 @@ type NameServer struct {
 }
 
 type FallbackFilter struct {
-	GeoIP  bool
-	IPCIDR []*net.IPNet
-	Domain []string
+	GeoIP     bool
+	GeoIPCode string
+	IPCIDR    []*net.IPNet
+	Domain    []string
 }
 
 type Config struct {
@@ -344,7 +345,9 @@ func NewResolver(config Config) *Resolver {
 
 	fallbackIPFilters := []fallbackIPFilter{}
 	if config.FallbackFilter.GeoIP {
-		fallbackIPFilters = append(fallbackIPFilters, &geoipFilter{})
+		fallbackIPFilters = append(fallbackIPFilters, &geoipFilter{
+			code: config.FallbackFilter.GeoIPCode,
+		})
 	}
 	for _, ipnet := range config.FallbackFilter.IPCIDR {
 		fallbackIPFilters = append(fallbackIPFilters, &ipnetFilter{ipnet: ipnet})

diff --git a/hub/executor/executor.go b/hub/executor/executor.go
@@ -123,9 +123,10 @@ func updateDNS(c *config.DNS) {
 		Pool:         c.FakeIPRange,
 		Hosts:        c.Hosts,
 		FallbackFilter: dns.FallbackFilter{
-			GeoIP:  c.FallbackFilter.GeoIP,
-			IPCIDR: c.FallbackFilter.IPCIDR,
-			Domain: c.FallbackFilter.Domain,
+			GeoIP:     c.FallbackFilter.GeoIP,
+			GeoIPCode: c.FallbackFilter.GeoIPCode,
+			IPCIDR:    c.FallbackFilter.IPCIDR,
+			Domain:    c.FallbackFilter.Domain,
 		},
 		Default: c.DefaultNameserver,
 		Policy:  c.NameServerPolicy,