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"]