Merge pull request #4 from tystuyfzand/feature/version-check

Improve checks and rule speed, add JSON support

.drone.yml

@@ -16,7 +16,7 @@ steps:
     environment:
       CGO_ENABLED: '0'
   - name: build
-    image: tystuyfzand/goc:1.20
+    image: tystuyfzand/goc:1.21
     volumes:
       - name: build
         path: /build

cmd/db/genaccessors.go

@@ -0,0 +1,115 @@
+package main
+
+import (
+	"fmt"
+	"github.com/armbian/redirector/db"
+	"github.com/samber/lo"
+	"reflect"
+	"strings"
+)
+
+var (
+	structTags = []string{"json", "yaml", "maxminddb"}
+)
+
+// This is a VERY messy way to generate static recursive field getters, which transform rule strings
+// into the field value
+func main() {
+	var asn db.ASN
+
+	accessors(asn)
+
+	var city db.City
+
+	accessors(city)
+	accessors(city.Continent)
+	accessors(city.Country)
+	accessors(city.Location)
+	accessors(city.RegisteredCountry)
+}
+
+func accessors(val any) {
+	v := reflect.ValueOf(val)
+
+	funcName := fmt.Sprintf("get%s", v.Type().Name())
+
+	valueType := v.Type()
+
+	// Check if we need to include an index check
+
+	var s strings.Builder
+
+	s.WriteString(fmt.Sprintf("func %s(v %s, keys []string) (any, bool) {\n", funcName, valueType.Name()))
+	s.WriteString("\tkey := keys[0]\n\n")
+
+	s.WriteString("\tswitch key {\n")
+
+	for i := 0; i < valueType.NumField(); i++ {
+		fieldType := valueType.Field(i)
+
+		var fieldNames = []string{fieldType.Name}
+
+		for _, tag := range structTags {
+			// Append tags to possible names
+			tagVal := fieldType.Tag.Get(tag)
+
+			if tagVal == "" {
+				continue
+			}
+
+			fieldNames = append(fieldNames, tagVal)
+		}
+
+		s.WriteString("\t\tcase ")
+
+		formattedNames := lo.Map(lo.Uniq(fieldNames), func(name string, _ int) string {
+			return fmt.Sprintf(`"%s"`, name)
+		})
+
+		s.WriteString(strings.Join(formattedNames, ", "))
+		s.WriteString(":\n")
+
+		// If struct/map/etc, access. If string/other value, return.
+		s.WriteString("\t\t\t")
+
+		// Check kind as it's an int, starting at bool and ending at Complex128. Then include String.
+		if fieldType.Type.Kind() >= reflect.Bool && fieldType.Type.Kind() <= reflect.Complex128 ||
+			fieldType.Type.Kind() == reflect.String {
+			s.WriteString("return ")
+			s.WriteString("v.")
+			s.WriteString(fieldType.Name)
+			s.WriteString(", true")
+		} else if fieldType.Type.Kind() == reflect.Map {
+			// Handle slice logic (index, so [0] off the key)
+			s.WriteString("index := getMapIndex(key)\n\n")
+
+			s.WriteString("\t\t\tif index == \"\" {\n")
+			s.WriteString("\t\t\t\treturn nil, false\n")
+			s.WriteString("\t\t\t}\n\n")
+
+			s.WriteString("\t\t\tm, found := v." + fieldType.Name + "[index]\n")
+
+			s.WriteString("\t\t\treturn m, found")
+		} else if fieldType.Type.Kind() == reflect.Slice {
+			// Handle slice logic (index, so [0] off the key)
+			s.WriteString("index := getSliceIndex(key)\n\n")
+
+			s.WriteString("\t\t\tif index == nil {\n")
+			s.WriteString("\t\t\t\treturn nil, false\n")
+			s.WriteString("\t\t\t}\n\n")
+
+			s.WriteString(fmt.Sprintf("\t\t\treturn v.%s[index]", fieldType.Name))
+		} else if fieldType.Type.Kind() == reflect.Struct {
+			s.WriteString("return get" + fieldType.Type.Name() + "(v." + fieldType.Name + ", keys[1:])")
+		} else {
+			s.WriteString("return nil, true")
+		}
+		s.WriteString("\n")
+	}
+
+	s.WriteString("\t}\n\n")
+	s.WriteString("\treturn nil, false\n")
+	s.WriteString("}\n")
+
+	fmt.Println(s.String())
+}

cmd/main.go

@@ -26,7 +26,7 @@ func main() {
 	viper.SetDefault("bind", ":8080")
 	viper.SetDefault("cacheSize", 1024)
 	viper.SetDefault("topChoices", 3)
-	viper.SetDefault("reloadKey", redirector.RandomSequence(32))
+	viper.SetDefault("reloadKey", util.RandomSequence(32))
 
 	viper.SetConfigName("dlrouter")        // name of config file (without extension)
 	viper.SetConfigType("yaml")            // REQUIRED if the config file does not have the extension in the name

config.go

@@ -3,6 +3,7 @@ package redirector
 import (
 	"crypto/tls"
 	"crypto/x509"
+	"github.com/armbian/redirector/db"
 	lru "github.com/hashicorp/golang-lru"
 	"github.com/oschwald/maxminddb-golang"
 	"github.com/pkg/errors"
@@ -168,7 +169,7 @@ func (r *Redirector) ReloadConfig() error {
 	}
 
 	// Force check
-	go r.servers.Check(r.checks)
+	go r.servers.Check(r, r.checks)
 
 	return nil
 }
@@ -305,7 +306,7 @@ func (r *Redirector) addServer(server ServerConfig, u *url.URL) (*Server, error)
 		return nil, err
 	}
 
-	var city City
+	var city db.City
 	err = r.db.Lookup(ips[0], &city)
 
 	if err != nil {

db/accessors.go

@@ -0,0 +1,164 @@
+package db
+
+import (
+	"regexp"
+	"strconv"
+	"strings"
+)
+
+var (
+	mapIndexRe   = regexp.MustCompile("\\[([a-zA-Z0-9]+)]")
+	sliceIndexRe = regexp.MustCompile("\\[(-?\\d+)]")
+)
+
+func getMapIndex(key string) string {
+	m := mapIndexRe.FindStringSubmatch(key)
+
+	if m == nil {
+		return ""
+	}
+
+	return m[1]
+}
+
+func getSliceIndex(key string) *int {
+	m := sliceIndexRe.FindStringSubmatch(key)
+
+	if m == nil {
+		return nil
+	}
+
+	v, err := strconv.Atoi(m[1])
+
+	if err != nil {
+		return nil
+	}
+
+	return &v
+}
+
+// GetValue is a generated, optimized value getter based on string keys
+// This is the only function in this file that should be edited when adding new types.
+func GetValue(val any, key string) (any, bool) {
+	keysSplit := strings.Split(key, ".")
+
+	switch keysSplit[0] {
+	case "asn":
+		return getASN(val.(ASN), keysSplit[1:])
+	case "city":
+		return getCity(val.(City), keysSplit[1:])
+	}
+
+	return nil, false
+}
+
+func getASN(v ASN, keys []string) (any, bool) {
+	key := keys[0]
+
+	switch key {
+	case "AutonomousSystemNumber", "autonomous_system_number":
+		return v.AutonomousSystemNumber, true
+	case "AutonomousSystemOrganization", "autonomous_system_organization":
+		return v.AutonomousSystemOrganization, true
+	}
+
+	return nil, false
+}
+
+func getCity(v City, keys []string) (any, bool) {
+	key := keys[0]
+
+	switch key {
+	case "Continent", "continent":
+		return getContinent(v.Continent, keys[1:])
+	case "Country", "country":
+		return getCountry(v.Country, keys[1:])
+	case "Location", "location":
+		return getLocation(v.Location, keys[1:])
+	case "RegisteredCountry", "registered_country":
+		return getRegisteredCountry(v.RegisteredCountry, keys[1:])
+	}
+
+	return nil, false
+}
+
+func getContinent(v Continent, keys []string) (any, bool) {
+	key := keys[0]
+
+	switch key {
+	case "Code", "code":
+		return v.Code, true
+	case "GeoNameID", "geoname_id":
+		return v.GeoNameID, true
+	case "Names", "names":
+		index := getMapIndex(key)
+
+		if index == "" {
+			return nil, false
+		}
+
+		m, found := v.Names[index]
+		return m, found
+	}
+
+	return nil, false
+}
+
+func getCountry(v Country, keys []string) (any, bool) {
+	key := keys[0]
+
+	switch key {
+	case "GeoNameID", "geoname_id":
+		return v.GeoNameID, true
+	case "IsoCode", "iso_code":
+		return v.IsoCode, true
+	case "Names", "names":
+		index := getMapIndex(key)
+
+		if index == "" {
+			return nil, false
+		}
+
+		m, found := v.Names[index]
+		return m, found
+	}
+
+	return nil, false
+}
+
+func getLocation(v Location, keys []string) (any, bool) {
+	key := keys[0]
+
+	switch key {
+	case "AccuracyRadius", "accuracy_radius":
+		return v.AccuracyRadius, true
+	case "Latitude", "latitude":
+		return v.Latitude, true
+	case "Longitude", "longitude":
+		return v.Longitude, true
+	}
+
+	return nil, false
+}
+
+func getRegisteredCountry(v RegisteredCountry, keys []string) (any, bool) {
+	key := keys[0]
+
+	switch key {
+	case "GeoNameID", "geoname_id":
+		return v.GeoNameID, true
+	case "IsoCode", "iso_code":
+		return v.IsoCode, true
+	case "Names", "names":
+		index := getMapIndex(key)
+
+		if index == "" {
+			return nil, false
+		}
+
+		m, found := v.Names[index]
+		return m, found
+	}
+
+	return nil, false
+}

db/structs.go

@@ -0,0 +1,40 @@
+package db
+
+// City represents a MaxmindDB city.
+// This used to only be used on load, but is now used with rules as well.
+type City struct {
+	Continent         Continent         `maxminddb:"continent" json:"continent"`
+	Country           Country           `maxminddb:"country" json:"country"`
+	Location          Location          `maxminddb:"location"`
+	RegisteredCountry RegisteredCountry `maxminddb:"registered_country" json:"registered_country"`
+}
+
+type Continent struct {
+	Code      string            `maxminddb:"code" json:"code"`
+	GeoNameID uint              `maxminddb:"geoname_id" json:"geoname_id"`
+	Names     map[string]string `maxminddb:"names" json:"names"`
+}
+
+type Country struct {
+	GeoNameID uint              `maxminddb:"geoname_id" json:"geoname_id"`
+	IsoCode   string            `maxminddb:"iso_code" json:"iso_code"`
+	Names     map[string]string `maxminddb:"names" json:"names"`
+}
+
+type Location struct {
+	AccuracyRadius uint16  `maxminddb:"accuracy_radius" json:"accuracy_radius"`
+	Latitude       float64 `maxminddb:"latitude" json:"latitude"`
+	Longitude      float64 `maxminddb:"longitude" json:"longitude"`
+}
+
+type RegisteredCountry struct {
+	GeoNameID uint              `maxminddb:"geoname_id" json:"geoname_id"`
+	IsoCode   string            `maxminddb:"iso_code" json:"iso_code"`
+	Names     map[string]string `maxminddb:"names" json:"names"`
+}
+
+// The ASN struct corresponds to the data in the GeoLite2 ASN database.
+type ASN struct {
+	AutonomousSystemNumber       uint   `maxminddb:"autonomous_system_number"`
+	AutonomousSystemOrganization string `maxminddb:"autonomous_system_organization"`
+}