From fbc2301287f40130ff640862916777674c80bb33 Mon Sep 17 00:00:00 2001 From: Yifu Yu Date: Sun, 17 Nov 2019 03:51:15 +0800 Subject: [PATCH] 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 | 5 +- go.sum | 2 + ipgeo/ipdb.go | 37 ++++++++++- main.go | 3 +- mapping/elf/core.go | 154 ++++++++++++++++++++++++++++++++++++++++++++ mapping/map.go | 44 ++++--------- rules/map.starlark | 22 +++++++ 7 files changed, 232 insertions(+), 35 deletions(-) create mode 100644 mapping/elf/core.go create mode 100644 rules/map.starlark diff --git a/go.mod b/go.mod index 7dd6c1e..017f42c 100644 --- a/go.mod +++ b/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 +) diff --git a/go.sum b/go.sum index 5a5932d..5923efc 100644 --- a/go.sum +++ b/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= diff --git a/ipgeo/ipdb.go b/ipgeo/ipdb.go index cd57b6a..0ae0e82 100644 --- a/ipgeo/ipdb.go +++ b/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 + }) +} diff --git a/main.go b/main.go index afdb0e7..d3c9796 100644 --- a/main.go +++ b/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) } diff --git a/mapping/elf/core.go b/mapping/elf/core.go new file mode 100644 index 0000000..6d99d14 --- /dev/null +++ b/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 "" +} diff --git a/mapping/map.go b/mapping/map.go index 17e9e83..7600c3e 100644 --- a/mapping/map.go +++ b/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 } diff --git a/rules/map.starlark b/rules/map.starlark new file mode 100644 index 0000000..ab3d603 --- /dev/null +++ b/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"]