diff --git a/main.go b/main.go index d3c9796..92d5ef5 100644 --- a/main.go +++ b/main.go @@ -8,7 +8,6 @@ import ( "net" "net/http" _ "net/http/pprof" - "strconv" "strings" "git.eve.moe/jackyyf/navigator/ipgeo" @@ -36,16 +35,15 @@ func responseWithError(resp http.ResponseWriter, statusCode int, message string) func responseWithJsonError(resp http.ResponseWriter, statusCode int, message string) { resp.Header().Set("Content-Type", "application/json") - body, err := json.Marshal(&errorMessage{ + resp.WriteHeader(statusCode) + encoder := json.NewEncoder(resp) + err := encoder.Encode(&errorMessage{ Error: message, }) if err != nil { // This should never happen panic("json marshal failed, check code") } - resp.Header().Set("Content-Length", strconv.Itoa(len(body))) - resp.WriteHeader(statusCode) - resp.Write(body) } func buildLocation(info *ipdb.CityInfo) string { @@ -148,6 +146,37 @@ func main() { log.Printf("%s => %s\n", ip, server) fmt.Fprint(resp, server) }) + clientApi := http.NewServeMux() + http.Handle("/client/", http.StripPrefix("/client", clientApi)) + clientV1Api := http.NewServeMux() + clientApi.Handle("/v1/", http.StripPrefix("/v1", clientV1Api)) + clientV1Api.HandleFunc("/getNodes", func(resp http.ResponseWriter, req *http.Request) { + host, _, err := net.SplitHostPort(req.RemoteAddr) + if err != nil { + responseWithJsonError(resp, http.StatusPreconditionFailed, errIPv4Only) + return + } + if net.ParseIP(host).To4() == nil { + responseWithJsonError(resp, http.StatusPreconditionFailed, errIPv4Only) + return + } + nodes := mapping.GetNodes() + if nodes == nil { + responseWithJsonError(resp, http.StatusInternalServerError, "Unable to get nodes") + return + } + suffix := mapping.GetSuffix(host) + resp.Header().Set("Content-Type", "application/json") + resp.WriteHeader(http.StatusOK) + jsonEncoder := json.NewEncoder(resp) + jsonEncoder.Encode(map[string]interface{}{ + "nodes": nodes, + "suffix": suffix, + }) + }) + clientV1Api.HandleFunc("/getSuffix", func(resp http.ResponseWriter, req *http.Request) { + }) + 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 index 6d99d14..09d3b42 100644 --- a/mapping/elf/core.go +++ b/mapping/elf/core.go @@ -27,6 +27,7 @@ var ( requiredFunc = map[string]struct{}{ "getMapping": struct{}{}, "getNodes": struct{}{}, + "getSuffix": struct{}{}, } parsedFunc *starlark.StringDict parseThread = &starlark.Thread{ @@ -152,3 +153,42 @@ func GetMapping(ip string) string { log.Println("Script returned unexpected result:", ret.String()) return "" } + +func GetNodes() (res []string) { + thread := pool.Get() + ret, err := starlark.Call(thread, (*parsedFunc)["getNodes"], nil, nil) + if err != nil { + log.Println("Starlark execute error:", err.Error()) + return nil + } + if r, ok := ret.(*starlark.List); ok { + res = make([]string, r.Len()) + for i := 0; i < r.Len(); i++ { + v := r.Index(i) + if s, ok := v.(starlark.String); ok { + res[i] = string(s) + } else { + log.Println("Script returned unexpected result:", + ret.String()) + } + } + } else { + log.Println("Script returned unexpected result:", ret.String()) + } + return +} + +func GetSuffix(ip string) string { + thread := pool.Get() + ret, err := starlark.Call(thread, (*parsedFunc)["getSuffix"], + 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 7600c3e..646cadc 100644 --- a/mapping/map.go +++ b/mapping/map.go @@ -8,7 +8,10 @@ import ( var ( defaultServer = flag.String("fallback-node", "xe-mci1-us.edge.eve.network", - "Default CDN node in case of any error when executing script") + "Default CDN node in case of any script error") + defaultSuffix = flag.String("fallback-suffix", + ".edge.eve.network", + "Default CDN suffix in case of any script error") ) func Initialize() { @@ -22,3 +25,15 @@ func Get(ip string) string { } return ret } + +func GetSuffix(ip string) string { + ret := elf.GetSuffix(ip) + if ret == "" { + return *defaultSuffix + } + return ret +} + +func GetNodes() []string { + return elf.GetNodes() +} diff --git a/rules/map.starlark b/rules/map.starlark index ab3d603..2ed82eb 100644 --- a/rules/map.starlark +++ b/rules/map.starlark @@ -20,3 +20,11 @@ def getMapping(ip): def getNodes(): return ["xe-mci1-us", "ge-fsn1-de", "ge-lax1-us"] + +def getSuffix(ip): + info = geoLookup(ip) + if not info: + return default_suffix + if info.CountryCode == "CN": + return CHINA_MAINLAND_SUFFIX + return GLOBAL_SUFFIX