Release v0.1.0: now with starlark script support.

You may send SIGHUP signal to the process to make it reread the script.
It will use older loaded version if parsing of new file failed.

go.mod

@@ -2,4 +2,7 @@ module git.eve.moe/jackyyf/navigator
 
 go 1.13
 
-require github.com/ipipdotnet/ipdb-go v1.2.0
+require (
+	github.com/ipipdotnet/ipdb-go v1.2.0
+	go.starlark.net v0.0.0-20191113183327-aaf7be003892
+)

go.sum

@@ -1,2 +1,4 @@
 github.com/ipipdotnet/ipdb-go v1.2.0 h1:Afa0qx/SgRevzIK8Qg1TevuD5M28kFLWbzPvU+GQJ08=
 github.com/ipipdotnet/ipdb-go v1.2.0/go.mod h1:6SFLNyXDBF6q99FQvbOZJQCc2rdPrB1V5DSy4S83RSw=
+go.starlark.net v0.0.0-20191113183327-aaf7be003892 h1:ZP11CRSzO9uOTTOVkH6yodtI3kSY69vUID8lx8B0M3s=
+go.starlark.net v0.0.0-20191113183327-aaf7be003892/go.mod h1:c1/X6cHgvdXj6pUlmWKMkuqRnW4K8x2vwt6JAaaircg=

ipgeo/ipdb.go

@@ -2,8 +2,12 @@ package ipgeo
 
 import (
 	"flag"
-	"github.com/ipipdotnet/ipdb-go"
 	"log"
+
+	"git.eve.moe/jackyyf/navigator/mapping/elf"
+	"github.com/ipipdotnet/ipdb-go"
+	"go.starlark.net/starlark"
+	"go.starlark.net/starlarkstruct"
 )
 
 var (
@@ -19,7 +23,7 @@ var (
 	db     *ipdb.City
 )
 
-func Init() {
+func Initialize() {
 	var err error
 	db, err = ipdb.NewCity(*ipipdb)
 	if err != nil {
@@ -39,7 +43,8 @@ func Init() {
 				}
 			}
 			if !ok {
-				log.Fatalln("This IPIP.net database has no required field", requiredField)
+				log.Fatalln("This IPIP.net database has no required field",
+					requiredField)
 			}
 		}
 	}
@@ -48,3 +53,29 @@ func Init() {
 func Get() *ipdb.City {
 	return db
 }
+
+func init() {
+	elf.RegisterFunc("geoLookup",
+		func(t *starlark.Thread, b *starlark.Builtin, args starlark.Tuple,
+			kwargs []starlark.Tuple) (starlark.Value, error) {
+			var ip string
+			if err := starlark.UnpackArgs(b.Name(), args, kwargs,
+				"ip", &ip); err != nil {
+				return nil, err
+			}
+			result, err := db.FindInfo(ip, "EN")
+			if err != nil {
+				return nil, err
+			}
+			return starlarkstruct.FromStringDict(
+				starlark.String("IPDBCityInfo"),
+				starlark.StringDict{
+					"CountryName":   starlark.String(result.CountryName),
+					"RegionName":    starlark.String(result.RegionName),
+					"CityName":      starlark.String(result.CityName),
+					"IspDomain":     starlark.String(result.IspDomain),
+					"CountryCode":   starlark.String(result.CountryCode),
+					"ContinentCode": starlark.String(result.ContinentCode),
+				}), nil
+		})
+}

main.go

@@ -64,6 +64,8 @@ func buildLocation(info *ipdb.CityInfo) string {
 
 func main() {
 	flag.Parse()
+	ipgeo.Initialize()
+	mapping.Initialize()
 
 	http.HandleFunc("/healthz", func(resp http.ResponseWriter, req *http.Request) {
 		resp.WriteHeader(200)
@@ -146,7 +148,6 @@ func main() {
 		log.Printf("%s => %s\n", ip, server)
 		fmt.Fprint(resp, server)
 	})
-	ipgeo.Init()
 	log.Println("HTTP server is running on", *listen_spec)
 	http.ListenAndServe(*listen_spec, nil)
 }

mapping/elf/core.go

@@ -0,0 +1,154 @@
+package elf
+
+import (
+	"flag"
+	"fmt"
+	"log"
+	"os"
+	"os/signal"
+	"sync"
+	"sync/atomic"
+	"syscall"
+
+	"go.starlark.net/starlark"
+)
+
+type threadPool struct {
+	pool  []*starlark.Thread
+	mutex *sync.Mutex
+}
+
+var (
+	scriptFile = flag.String("mapping-script", "map.starlark",
+		"Starlark script used to do mapping logic")
+	threadPoolSize = flag.Int("script-thread-pool", 128,
+		"Thread pool size for starlark execution engine")
+	builtinFunc  = make(starlark.StringDict)
+	requiredFunc = map[string]struct{}{
+		"getMapping": struct{}{},
+		"getNodes":   struct{}{},
+	}
+	parsedFunc  *starlark.StringDict
+	parseThread = &starlark.Thread{
+		Name: "parseThread",
+	}
+	reloadSignal = make(chan os.Signal)
+	pool         *threadPool
+	counter      uint64
+)
+
+func newThreadPool(size int) *threadPool {
+	ret := &threadPool{
+		pool:  make([]*starlark.Thread, size),
+		mutex: &sync.Mutex{},
+	}
+	for i := range ret.pool {
+		counter += 1
+		ret.pool[i] = &starlark.Thread{
+			Name: fmt.Sprintf("ElfThread-%d", counter),
+		}
+	}
+	return ret
+}
+
+func (pool *threadPool) Get() (ret *starlark.Thread) {
+	pool.mutex.Lock()
+	defer pool.mutex.Unlock()
+	if len(pool.pool) == 0 {
+		return &starlark.Thread{
+			Name: fmt.Sprintf("ElfThread-%d", atomic.AddUint64(&counter, 1)),
+		}
+	}
+	ret = pool.pool[len(pool.pool)-1]
+	pool.pool = pool.pool[:len(pool.pool)-1]
+	return
+}
+
+func (pool *threadPool) Put(thread *starlark.Thread) {
+	pool.mutex.Lock()
+	defer pool.mutex.Unlock()
+	if len(pool.pool) == *threadPoolSize {
+		return
+	}
+	pool.pool = append(pool.pool, thread)
+}
+
+func RegisterFunc(name string, impl func(*starlark.Thread, *starlark.Builtin,
+	starlark.Tuple, []starlark.Tuple) (starlark.Value, error)) {
+	if builtinFunc.Has(name) {
+		panic(fmt.Errorf("Function %s has already been declared as: %s",
+			name, builtinFunc[name].String()))
+	}
+	builtinFunc[name] = starlark.NewBuiltin(name, impl)
+	log.Println("Registered function", name, "to starlark")
+}
+
+func RequireFunc(name string) {
+	requiredFunc[name] = struct{}{}
+}
+
+func Initialize() {
+	globals, err := starlark.ExecFile(
+		parseThread, *scriptFile, nil, builtinFunc)
+	if err != nil {
+		panic(fmt.Errorf("Unable to parse starlark file: %s", err.Error()))
+	}
+	parsedFunc = &globals
+	for name := range requiredFunc {
+		if !parsedFunc.Has(name) {
+			panic(fmt.Errorf("Required function %s not provided", name))
+		}
+		if _, ok := globals[name].(starlark.Callable); !ok {
+			panic(fmt.Errorf(
+				"Variable %s is not callable as required function", name))
+		}
+
+	}
+	signal.Notify(reloadSignal, syscall.SIGHUP)
+	go func() {
+		for {
+			<-reloadSignal
+			reload()
+		}
+	}()
+	pool = newThreadPool(*threadPoolSize)
+}
+
+func reload() {
+	globals, err := starlark.ExecFile(
+		parseThread, *scriptFile, nil, builtinFunc)
+	if err != nil {
+		log.Println("Reload failed: Unable to parse starlark file:",
+			err.Error())
+		return
+	}
+	for name := range requiredFunc {
+		if !globals.Has(name) {
+			log.Println(
+				"Reload failed: Required function", name, "not provided")
+			return
+		}
+		if _, ok := globals[name].(starlark.Callable); !ok {
+			log.Println(
+				"Variable", name, "is not callable as required function")
+			return
+		}
+	}
+	log.Println("Reload success, new rules applied :)")
+	parsedFunc = &globals
+}
+
+func GetMapping(ip string) string {
+	thread := pool.Get()
+	ret, err := starlark.Call(thread, (*parsedFunc)["getMapping"],
+		starlark.Tuple{starlark.String(ip)}, nil)
+	if err != nil {
+		log.Println("Starlark execute error:", err.Error())
+		return ""
+	}
+	if r, ok := ret.(starlark.String); ok {
+		return string(r)
+	}
+	log.Println("Script returned unexpected result:", ret.String())
+	return ""
+}

mapping/map.go

@@ -2,39 +2,23 @@ package mapping
 
 import (
 	"flag"
-	"git.eve.moe/jackyyf/navigator/ipgeo"
-)
-
-const (
-	// Server IDs
-	WHOLESALE_INTERNET_10GE = "xe-mci1-us"
-	HETZNER_FSN_1GE         = "ge-fsn1-de"
-	HETZNER_HEL_1GE         = "ge-hel1-fi"
-	default_server          = WHOLESALE_INTERNET_10GE
-
-	// Served domain suffix
-	CHINA_MAINLAND_SUFFIX = ".eveedge.link"
-	GLOBAL_SUFFIX         = ".edge.eve.network"
-	default_suffix        = GLOBAL_SUFFIX
+	"git.eve.moe/jackyyf/navigator/mapping/elf"
 )
 
 var (
-	enableCNDomain     = flag.Bool("enable-cn-domain", true, "Enable china mainland specific domain")
-	enableCMExperiment = flag.Bool("cm-to-fsn", false, "Redirect all CM users to Hetzner FSN")
+	defaultServer = flag.String("fallback-node",
+		"xe-mci1-us.edge.eve.network",
+		"Default CDN node in case of any error when executing script")
 )
 
-// Get returns the edge node that should be used for client.
-func Get(ip string) string {
-	db := ipgeo.Get()
-	info_en, err := db.FindInfo(ip, "EN")
-	if err != nil {
-		return default_server + default_suffix
-	}
-	if *enableCMExperiment && info_en.IspDomain == "ChinaMobile" {
-		return HETZNER_FSN_1GE + CHINA_MAINLAND_SUFFIX
-	}
-	if *enableCNDomain && info_en.CountryCode == "CN" {
-		return default_server + CHINA_MAINLAND_SUFFIX
-	}
-	return default_server + GLOBAL_SUFFIX
+func Initialize() {
+	elf.Initialize()
+}
+
+func Get(ip string) string {
+	ret := elf.GetMapping(ip)
+	if ret == "" {
+		return *defaultServer
+	}
+	return ret
 }

rules/map.starlark

@@ -0,0 +1,22 @@
+WHOLESALE_INTERNET_10GE = "xe-mci1-us"
+HETZNER_FSN_1GE         = "ge-fsn1-de"
+HETZNER_HEL_1GE         = "ge-hel1-fi"
+default_server          = WHOLESALE_INTERNET_10GE
+
+CHINA_MAINLAND_SUFFIX = ".eveedge.link"
+GLOBAL_SUFFIX         = ".edge.eve.network"
+default_suffix        = GLOBAL_SUFFIX
+
+
+def getMapping(ip):
+	info = geoLookup(ip)
+	if not info:
+		return default_server + default_suffix
+	if info.IspDomain == "ChinaMobile":
+		return HETZNER_FSN_1GE + CHINA_MAINLAND_SUFFIX
+	if info.CountryCode == "CN":
+		return default_server + CHINA_MAINLAND_SUFFIX
+	return default_server + GLOBAL_SUFFIX
+
+def getNodes():
+	return ["xe-mci1-us", "ge-fsn1-de", "ge-lax1-us"]