Added Client V1 API to allow speed tests

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)
 }

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 ""
+}

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()
+}

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